Repository: microsoft/DeepSpeed

Issue:
[BUG] `--partition-activations` in Meg-DS breaks in the Deepspeed land for TP>1 & PP>1
**Describe the bug**
`--partition-activations` in Meg-DS breaks in the Deepspeed land.

it only works with TP=2 or PP=2 - not both. Fails with TP=2 and PP=2.

Here is how to reproduce:

use this PR https://github.com/bigscience-workshop/Megatron-DeepSpeed/pull/184
**on a 4+ gpu instance**

set up:
```
git clone https://github.com/bigscience-workshop/Megatron-DeepSpeed
cd Megatron-DeepSpeed
gh pr checkout 184
pip install -r requirements.txt
pip install pytest-timeout
```
run test:

```
pytest tests/test_training.py::MegDSTestTraining::test_training_all_0_base 
```

It (and other tests) fail with:

```
E           Traceback (most recent call last):
E             File "/actions-runner/_work/Megatron-DeepSpeed/Megatron-DeepSpeed/pretrain_gpt.py", line 268, in <module>
E               pretrain(train_valid_test_datasets_provider, model_provider, forward_step,
E             File "/actions-runner/_work/Megatron-DeepSpeed/Megatron-DeepSpeed/megatron/training.py", line 181, in pretrain
E               iteration = train(forward_step_func,
E             File "/actions-runner/_work/Megatron-DeepSpeed/Megatron-DeepSpeed/megatron/training.py", line 766, in train
E               train_step(forward_step_func,
E             File "/actions-runner/_work/Megatron-DeepSpeed/Megatron-DeepSpeed/megatron/training.py", line 427, in train_step
E               loss = model[0].train_batch(data_iter=data_iterator)
E             File "/usr/local/lib/python3.8/dist-packages/deepspeed/runtime/pipe/engine.py", line 329, in train_batch
E               self._exec_schedule(sched)
E             File "/usr/local/lib/python3.8/dist-packages/deepspeed/runtime/pipe/engine.py", line 1335, in _exec_schedule
E               self._exec_instr(**cmd.kwargs)
E             File "/usr/local/lib/python3.8/dist-packages/deepspeed/runtime/pipe/engine.py", line 691, in _exec_backward_pass
E               super().backward(self.loss)
E             File "/usr/local/lib/python3.8/dist-packages/deepspeed/runtime/engine.py", line 1460, in backward
E               self.optimizer.backward(loss)
E             File "/usr/local/lib/python3.8/dist-packages/deepspeed/runtime/zero/stage2.py", line 1875, in backward
E               self.loss_scaler.backward(loss.float(), retain_graph=retain_graph)
E             File "/usr/local/lib/python3.8/dist-packages/deepspeed/runtime/fp16/loss_scaler.py", line 53, in backward
E               scaled_loss.backward(retain_graph=retain_graph)
E             File "/usr/local/lib/python3.8/dist-packages/torch/tensor.py", line 245, in backward
E               torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
E             File "/usr/local/lib/python3.8/dist-packages/torch/autograd/__init__.py", line 145, in backward
E               Variable._execution_engine.run_backward(
E           RuntimeError: The size of tensor a (4096) must match the size of tensor b (64) at non-singleton dimension 2
```

Thank you!

@ShadenSmith


Discussion:

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
@@ -506,7 +506,7 @@ def forward(ctx, run_function, all_outputs, *args):
 
         def save_args_for_backward(*all_args):
             tensor_args, non_tensor_args, tensor_flags = extract_tensors(all_objects=all_args)
-            ctx.save_for_backward(*tensor_args)
+            ctx.deepspeed_saved_tensors = tensor_args
             ctx.non_tensor_args = non_tensor_args
             ctx.tensor_flags = tensor_flags
 
@@ -656,16 +656,16 @@ def backward(ctx, *grads):
         if PARTITION_ACTIVATIONS:
             # with torch.cuda.stream(transport_stream):
             inputs = gather_partitioned_activations(
-                ctx.saved_tensors,
+                ctx.deepspeed_saved_tensors,
                 device=cuda_device if CPU_CHECKPOINT else None)
             detached_inputs = detach_variable(inputs)
         elif CPU_CHECKPOINT:
-            inputs = move_to_device(ctx.saved_tensors,
+            inputs = move_to_device(ctx.deepspeed_saved_tensors,
                                     cuda_device,
                                     is_activation_to_checkpoint)
             detached_inputs = detach_variable(inputs)
         else:
-            inputs = ctx.saved_tensors
+            inputs = ctx.deepspeed_saved_tensors
             detached_inputs = detach_variable(inputs)
 
         # Add non tensor input args
diff --git a/deepspeed/utils/debug.py b/deepspeed/utils/debug.py
--- a/deepspeed/utils/debug.py
+++ b/deepspeed/utils/debug.py
@@ -120,3 +120,21 @@ def print_rank_0(message, debug=False, force=False):
     for m in msgs:
         fh.write(f"{m}\n")
     fh.flush()
+
+
+def print_backward_tensors(tensor):
+    def _print_bwd_tensors(grad_fn):
+        print(f"Backward tensors in {grad_fn}")
+        for funcs in grad_fn.next_functions:
+            if funcs[0]:
+                try:
+                    tensor = getattr(funcs[0], 'variable')
+                    print(funcs[0])
+                    print(
+                        f"Tensor - id: {id(tensor)}, shape: {tensor.shape}, data: {tensor}, grad: {tensor.grad}"
+                    )
+                except AttributeError as e:
+                    _print_bwd_tensors(funcs[0])
+
+    if hasattr(tensor, 'grad_fn'):
+        _print_bwd_tensors(tensor.grad_fn)


Test patch: