Repository: microsoft/DeepSpeed

Issue:
[deepspeed checkpointing] AttributeError: 'NoneType' object has no attribute 'numel'
So I took a public GPT-2 class implementation (not Megatron-LM) and I added deepspeed checkpointing to it for all 48 layers.
In my train script for this class, I added the following line:

```
deepspeed.checkpointing.configure(mpu_=None, deepspeed_config=args.deepspeed_config)
```

My deepspeed config JSON is as follows:

```
{
  "train_batch_size": 128,
  "gradient_accumulation_steps": 8,
  "gradient_clipping": 1.0,
  "optimizer": {
    "type": "adam",
    "params": {
      "lr": 6.25e-5
    }
  },
  "fp16": {
    "enabled": true
  },
  "zero_optimization": {
    "stage": 2,
    "cpu_offload": true,
    "contiguous_gradients": true,
    "overlap_comm": false,
    "allgather_bucket_size": 500000000
  },

  "activation_checkpointing": {
    "partition_activations": true,
    "contiguous_memory_optimization": true,
    "number_checkpoints": 48,
    "cpu_checkpointing": true
  }

}
```

When I try running my script, I get the following error:

```
  File "/path/to/my/modeling_gpt2.py", line 221, in forward
    encoder_attention_mask)
  File "/usr/local/lib/python3.6/dist-packages/deepspeed/runtime/activation_checkpointing/checkpointing.py", line 582, in checkpoint
    return CheckpointFunction.apply(function, *args)
  File "/usr/local/lib/python3.6/dist-packages/deepspeed/runtime/activation_checkpointing/checkpointing.py", line 376, in forward
    partition_size = get_partition_size(item)
  File "/usr/local/lib/python3.6/dist-packages/deepspeed/runtime/activation_checkpointing/checkpointing.py", line 275, in get_partition_size
    size = item.numel()
AttributeError: 'NoneType' object has no attribute 'numel'
```

Any ideas what's going on?

@tjruwase @ShadenSmith 


Discussion:
Hi @g-karthik , that is an unfriendly way of reporting that the `partition_activations`, `contiguous_memory_optimization`, and `cpu_checkpointing` features require model parallelism. The `partition_activations` feature assumes that activations are replicated among a group of model-parallel workers (as is the case in Megatron) and then partitions them to save additional memory with activation checkpointing. The `contiguous_memory_optimization` saves additional memory by avoiding fragementation and allocating all activations up front. The `cpu_checkpointing` feature then offloads the saved activations to CPU memory.

I will file an issue to catch and report the configuration error when model parallelism is not enabled.
Ah I see, thanks @ShadenSmith! I made those `false` and get a different `NoneType` error now :(

```
  File "/path/to/my/modeling_gpt2.py", line 221, in forward
    encoder_attention_mask)
  File "/usr/local/lib/python3.6/dist-packages/deepspeed/runtime/activation_checkpointing/checkpointing.py", line 582, in checkpoint
    return CheckpointFunction.apply(function, *args)
  File "/usr/local/lib/python3.6/dist-packages/deepspeed/runtime/activation_checkpointing/checkpointing.py", line 416, in forward
    inputs_cuda = [item.to(cuda_device) for item in args]
  File "/usr/local/lib/python3.6/dist-packages/deepspeed/runtime/activation_checkpointing/checkpointing.py", line 416, in <listcomp>
    inputs_cuda = [item.to(cuda_device) for item in args]
AttributeError: 'NoneType' object has no attribute 'to'
```
Btw as a more foundational question, why couldn't we have upfront contiguous memory allocation for activations in the case of no-model-parallelism? Same question for offloading activations to CPU with `cpu_checkpointing=True`.

I understand the behavior of `partition_activations` is strongly tied to model-parallelism, but the others shouldn't necessarily need to, right?
Hm, can you share a bit about the model code that uses checkpointing? My first response still applies, but now I'm wondering if there's a bigger issue that we're not catching. This looks like `checkpoint()` is not getting the `args`, which should be the input to the checkpointed method.
@ShadenSmith yep, this is kinda what it looks like inside the model's `forward()`:

```
            checkpoint = deepspeed.checkpointing.checkpoint

            def create_custom_forward(module):
                def custom_forward(*inputs):
                    # checkpointing only works with tuple returns, not with lists
                    return tuple(output for output in module(*inputs, use_cache, output_attentions))

                return custom_forward

            outputs = checkpoint(create_custom_forward(block),
                                 hidden_states,
                                 layer_past,
                                 attention_mask,
                                 head_mask[i],
                                 encoder_hidden_states,
                                 encoder_attention_mask)


```

I assumed `deepspeed.checkpointing.checkpoint` is like a drop-in replacement for `torch.utils.checkpoint.checkpoint` since the signatures are the same. Is that a fair assumption to make?
Ah @ShadenSmith I think I know what the issue is, but correct me if I'm wrong. The default `torch.utils.checkpoint.checkpoint` [implementation](https://pytorch.org/docs/stable/_modules/torch/utils/checkpoint.html#checkpoint) does not assume anything about availability of all `*args`, i.e., the `CheckpointFunction`'s `forward()` method implementation allows for some arguments to be `None`.

However, in the corresponding deepspeed `CheckpointFunction`'s `forward()` method, the implementation [seems](https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/runtime/activation_checkpointing/checkpointing.py#L375) to [assume](https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/runtime/activation_checkpointing/checkpointing.py#L416) at multiple [places](https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/runtime/activation_checkpointing/checkpointing.py#L441) that *all* passed `args` are not `None`.

Perhaps these enumerations need to be modified and made more generic to account for cases where a layer being checkpointed could have some args as `None`, depending on the consumer's choice?

EDIT: Yep, I ran a test where I simply replaced `deepspeed.checkpointing.checkpoint` with `torch.utils.checkpoint.checkpoint` (all else equal), I can confirm the job runs successfully with the replacement. So there definitely needs to be some updates to the deepspeed checkpointing implementation.

Also, do you happen to have performed any side-by-side benchmarking of `torch.utils.checkpoint.checkpoint` and `deepspeed.checkpointing.checkpoint`? I'd love to know more about it if y'all have done it! :)
@ShadenSmith happy new year! Just bumping this back up your notifications!
Hey @g-karthik , happy new year and thanks for the ping! I hope you had a nice holiday. I spent some time away to focus on [python scalability](https://twitter.com/shaden_smith/status/1344717219259768833) :-).

Thanks for the deep dive on debugging this issue. I think you're right. DeepSpeed's activation checkpointing should work as a drop-in replacement in the simple case without partitioning, etc.

I don't know if we've done any benchmarking beyond at large scale with model parallelism where it's critical for huge activations. Maybe @samyam knows more.


@ShadenSmith The `args` fix should be straightforward, unless I'm missing something? i.e., just need to ignore those args that are `None`?

I've been using `torch.utils.checkpoint.checkpoint` as a workaround so far and I'm unable to fit the GPT-2 10B configuration (with 1024 sequence length) in my Tesla V100 GPU (32 GB memory). The best I've fitted with this workaround is the GPT-2 8B configuration (with 512 sequence length). I'm not quite sure yet if using `deepspeed.checkpointing.checkpoint` will help me fit 10B with 1024 sequence length, but I'll at least be able to try after the above `args` fix.

Resolving patch:
diff --git a/deepspeed/runtime/activation_checkpointing/checkpointing.py b/deepspeed/runtime/activation_checkpointing/checkpointing.py
--- a/deepspeed/runtime/activation_checkpointing/checkpointing.py
+++ b/deepspeed/runtime/activation_checkpointing/checkpointing.py
@@ -373,6 +373,10 @@ def forward(ctx, run_function, *args):
 
             inputs = []
             for i, item in enumerate(args[:-1]):
+                if not torch.is_tensor(item):
+                    inputs.append(item)
+                    continue
+
                 partition_size = get_partition_size(item)
                 partition = item.detach().contiguous().view(-1).narrow(
                     0,
@@ -413,7 +417,12 @@ def forward(ctx, run_function, *args):
             inputs.append(args[-1])
 
         #just in case something funky is happening such as reuse of inputs
-        inputs_cuda = [item.to(cuda_device) for item in args]
+        inputs_cuda = []
+        for item in args:
+            if torch.is_tensor(item):
+                inputs_cuda.append(item.to(cuda_device))
+            else:
+                inputs_cuda.append(item)
 
         # Copy the rng states.
         ctx.fwd_cpu_rng_state = torch.get_rng_state()
@@ -439,6 +448,10 @@ def forward(ctx, run_function, *args):
         if PARTITION_ACTIVATIONS:
             new_args = []
             for i, (arg, inp) in enumerate(zip(args, inputs)):
+                if not torch.is_tensor(arg):
+                    new_args.append(arg)
+                    continue
+
                 size = torch.tensor(arg.size())
 
                 arg.data = inp.data
@@ -573,7 +586,14 @@ def backward(ctx, *grads):
             timers.log(['backward'])
         if SYNCHRONIZE:
             torch.cuda.synchronize()
-        return (None, ) + tuple(inp.grad for inp in detached_inputs)
+        ret_list = [None]  # first None for ctx
+        for inp in detached_inputs:
+            if torch.is_tensor(inp):
+                ret_list.append(inp.grad)
+            else:
+                ret_list.append(None)
+
+        return tuple(ret_list)
 
 
 def checkpoint(function, *args):


Test patch:
diff --git a/tests/unit/test_activation_checkpointing.py b/tests/unit/test_activation_checkpointing.py
--- a/tests/unit/test_activation_checkpointing.py
+++ b/tests/unit/test_activation_checkpointing.py
@@ -23,7 +23,7 @@ def _compute(module, *inputs, do_checkpoint=False):
 
     sum(o.sum() for o in outputs if o.requires_grad).backward()
     grads = [p.grad for p in module.parameters()]
-    input_grads = [inp.grad for inp in inputs]
+    input_grads = [inp.grad for inp in inputs if torch.is_tensor(inp)]
 
     return {
         'outputs': outputs,
@@ -32,6 +32,18 @@ def _compute(module, *inputs, do_checkpoint=False):
     }
 
 
+def _prep_inputs(*inputs):
+    _inputs = []
+
+    for inp in inputs:
+        inp = deepcopy(inp)
+        if torch.is_tensor(inp):
+            inp = inp.cuda()
+        _inputs.append(inp)
+
+    return tuple(_inputs)
+
+
 # This is distributed because checkpoint() assumes that torch.distributed is initialized.
 # torch.distributed is used with activation partitioning, but not for these simple cases.
 @distributed_test(world_size=1)
@@ -43,11 +55,11 @@ def _test_activation_checkpoint(module, *inputs):
     module.eval()
 
     module_ = deepcopy(module)
-    inputs_ = tuple(deepcopy(inp).cuda() for inp in inputs)
+    inputs_ = _prep_inputs(*inputs)
     base = _compute(module_, *inputs_, do_checkpoint=False)
 
     module_ = deepcopy(module)
-    inputs_ = tuple(deepcopy(inp).cuda() for inp in inputs)
+    inputs_ = _prep_inputs(*inputs)
     test = _compute(module_, *inputs_, do_checkpoint=True)
 
     for group in base.keys():
@@ -155,3 +167,15 @@ def test_ckpt_inputs2_outputs3(mask):
     inputs = torch.rand(HIDDEN_DIM)
     inputs.requires_grad = True
     _test_activation_checkpoint(module, inputs, mask)
+
+
+class DropMaskLinear(torch.nn.Linear):
+    def forward(self, x, mask):
+        return super().forward(x)
+
+
+def test_ckpt_arg_none():
+    module = DropMaskLinear(HIDDEN_DIM, HIDDEN_DIM)
+    inputs = (torch.rand(HIDDEN_DIM), None)
+    inputs[0].requires_grad = True
+    _test_activation_checkpoint(module, *inputs)