Got error when training GPT2 with FSDP and activation checkpoint

I'm trying to train GPT2 with FSDP.

My environment is below.

PyTorch: 1.10.0+cu113
Fairscale: 0.4.5
transformers: 4.16.2
Tesla A100 x8

When I set `CUDA_LAUNCH_BLOCKING=1`, I got:
```python
Traceback (most recent call last):
File "/home/lclhx/DeepSpeed-FairScale-Benchmark/run_benchmark.py", line 65, in
main()
File "/home/lclhx/DeepSpeed-FairScale-Benchmark/run_benchmark.py", line 51, in main
run_iter_func(model, optimizer, criterion)
File "/home/lclhx/DeepSpeed-FairScale-Benchmark/benchmark/fairscale/fairscale.py", line 23, in run_iter
loss.backward()
File "/home/lclhx/.conda/envs/colossal/lib/python3.9/site-packages/torch/_tensor.py", line 307, in backward
torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
File "/home/lclhx/.conda/envs/colossal/lib/python3.9/site-packages/torch/autograd/__init__.py", line 154, in backward
Variable._execution_engine.run_backward(
File "/home/lclhx/.conda/envs/colossal/lib/python3.9/site-packages/torch/autograd/function.py", line 199, in apply
return user_fn(self, *args)
File "/home/lclhx/.conda/envs/colossal/lib/python3.9/site-packages/fairscale/nn/checkpoint/checkpoint_activations.py", line 348, in backward
torch.autograd.backward(outputs_with_grad, args_with_grad)
File "/home/lclhx/.conda/envs/colossal/lib/python3.9/site-packages/torch/autograd/__init__.py", line 154, in backward
Variable._execution_engine.run_backward(
File "/home/lclhx/.conda/envs/colossal/lib/python3.9/site-packages/torch/autograd/function.py", line 199, in apply
return user_fn(self, *args)
File "/home/lclhx/.conda/envs/colossal/lib/python3.9/site-packages/torch/utils/checkpoint.py", line 122, in backward
outputs = ctx.run_function(*detached_inputs)
File "/home/lclhx/.conda/envs/colossal/lib/python3.9/contextlib.py", line 137, in __exit__
self.gen.throw(typ, value, traceback)
File "/home/lclhx/.conda/envs/colossal/lib/python3.9/site-packages/torch/random.py", line 129, in fork_rng
torch.cuda.set_rng_state(gpu_rng_state, device)
File "/home/lclhx/.conda/envs/colossal/lib/python3.9/site-packages/torch/cuda/random.py", line 64, in set_rng_state
_lazy_call(cb)
File "/home/lclhx/.conda/envs/colossal/lib/python3.9/site-packages/torch/cuda/__init__.py", line 153, in _lazy_call
callable()
File "/home/lclhx/.conda/envs/colossal/lib/python3.9/site-packages/torch/cuda/random.py", line 62, in cb
default_generator.set_state(new_state_copy)
RuntimeError: CUDA error: an illegal memory access was encountered
```

When `CUDA_LAUNCH_BLOCKING` was not set, I got:
```python
Traceback (most recent call last):
File "/home/lclhx/DeepSpeed-FairScale-Benchmark/run_benchmark.py", line 65, in
main()
File "/home/lclhx/DeepSpeed-FairScale-Benchmark/run_benchmark.py", line 51, in main
run_iter_func(model, optimizer, criterion)
File "/home/lclhx/DeepSpeed-FairScale-Benchmark/benchmark/fairscale/fairscale.py", line 23, in run_iter
loss.backward()
File "/home/lclhx/.conda/envs/colossal/lib/python3.9/site-packages/torch/_tensor.py", line 307, in backward
torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
File "/home/lclhx/.conda/envs/colossal/lib/python3.9/site-packages/torch/autograd/__init__.py", line 154, in backward
Variable._execution_engine.run_backward(
File "/home/lclhx/.conda/envs/colossal/lib/python3.9/site-packages/torch/autograd/function.py", line 199, in apply
return user_fn(self, *args)
File "/home/lclhx/.conda/envs/colossal/lib/python3.9/site-packages/fairscale/nn/checkpoint/checkpoint_activations.py", line 348, in backward
torch.autograd.backward(outputs_with_grad, args_with_grad)
File "/home/lclhx/.conda/envs/colossal/lib/python3.9/site-packages/torch/autograd/__init__.py", line 154, in backward
Variable._execution_engine.run_backward(
File "/home/lclhx/.conda/envs/colossal/lib/python3.9/site-packages/torch/autograd/function.py", line 199, in apply
return user_fn(self, *args)
File "/home/lclhx/.conda/envs/colossal/lib/python3.9/site-packages/torch/utils/checkpoint.py", line 122, in backward
outputs = ctx.run_function(*detached_inputs)
File "/home/lclhx/.conda/envs/colossal/lib/python3.9/site-packages/transformers/models/gpt2/modeling_gpt2.py", line 881, in custom_forward
return module(*inputs, use_cache, output_attentions)
File "/home/lclhx/.conda/envs/colossal/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
return forward_call(*input, **kwargs)
File "/home/lclhx/.conda/envs/colossal/lib/python3.9/site-packages/transformers/models/gpt2/modeling_gpt2.py", line 400, in forward
attn_outputs = self.attn(
File "/home/lclhx/.conda/envs/colossal/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
return forward_call(*input, **kwargs)
File "/home/lclhx/.conda/envs/colossal/lib/python3.9/site-packages/transformers/models/gpt2/modeling_gpt2.py", line 321, in forward
raise e
File "/home/lclhx/.conda/envs/colossal/lib/python3.9/site-packages/transformers/models/gpt2/modeling_gpt2.py", line 318, in forward
query, key, value = self.c_attn(hidden_states).split(self.split_size, dim=2)
File "/home/lclhx/.conda/envs/colossal/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
return forward_call(*input, **kwargs)
File "/home/lclhx/.conda/envs/colossal/lib/python3.9/site-packages/transformers/modeling_utils.py", line 1837, in forward
x = torch.addmm(self.bias, x.view(-1, x.size(-1)), self.weight)
RuntimeError: setStorage: sizes [16384, 3072], strides [0, 1], storage offset 332776448, and itemsize 2 requiring a storage size of 665559040 are out of bounds for storage of size 0
```

I train my model like:
```python
def run_iter(model, optimizer, criterion):
# img = torch.rand(BATCH_SIZE, 3, IMG_SIZE, IMG_SIZE).cuda()
# label = torch.randint(0, NUM_CLASS, (BATCH_SIZE, )).cuda()
input_ids = torch.randint(0, VOCAB_SIZE, (BATCH_SIZE, SEQ_LEN)).cuda()
mask = torch.ones((BATCH_SIZE, SEQ_LEN), dtype=torch.int64, device=torch.cuda.current_device())
optimizer.zero_grad()
model.zero_grad(set_to_none=True)

with torch.cuda.amp.autocast():
# out = model(img)
# loss = criterion(out, label)
out = model(input_ids, mask)
loss = criterion(out, input_ids)
loss.backward()
optimizer.step()

def init(model, criterion, stage):

if stage < 2:
model = DDP(model, device_ids=[dist.get_rank()])
elif stage == 2:
model = ShardedDataParallel(model, optimizer)
elif stage == 3:
model = FullyShardedDataParallel(model, mixed_precision=True, reshard_after_forward=False, disable_reshard_on_root=False)
if stage == 1 or stage == 2:
optimizer = OSS(params=model.parameters(), optim=OPTIM, lr=0.001)
else:
optimizer = OPTIM(model.parameters(), lr=0.001)
# criterion = torch.nn.CrossEntropyLoss()
return model, optimizer, criterion
```

The GPT2 provided by `transformers` use torch's checkpoint. I also tried to use fairscale's `checkpoint_wrapper` by modifying the source code of `transformers`. However, I still got the error. Could you help me figure out this problem?

Reproduce the GPT-2 training case with the listed PyTorch, FairScale, and Transformers versions, starting from benchmark/fairscale/fairscale.py and its run_iter function. Trace the activation-checkpoint backward path through fairscale/nn/checkpoint/checkpoint_activations.py and transformers/models/gpt2/modeling_gpt2.py. Done means the FSDP training path no longer produces the reported illegal-memory-access or out-of-bounds-storage errors.

Written by the indexing model from the issue text.