From f0f4776096adf7cd588b54a85a3f220790ae4191 Mon Sep 17 00:00:00 2001 From: Vineeth Sai Date: Mon, 7 Sep 2026 10:51:32 -0700 Subject: [PATCH] Rebuild the full argument list when partitioning activations `get_partitioned_activations_for_backward` saves a (value, size) pair for every argument, so both the tensor flags and the non-tensor values hold two entries per argument. `merge_tensors` collapsed the flags by skipping the one after a True, which drops nothing for a non-tensor argument, whose value and its None size are both non-tensors, and it never collapsed the non-tensors at all. The recompute was handed a stray None after every non-tensor argument, shifting everything that followed. Every argument owns an even index in both lists, so drop the odd ones. Signed-off-by: Vineeth Sai --- .../activation_checkpointing/checkpointing.py | 15 ++--- .../test_activation_checkpointing.py | 58 +++++++++++++++++++ 2 files changed, 64 insertions(+), 9 deletions(-) diff --git a/deepspeed/runtime/activation_checkpointing/checkpointing.py b/deepspeed/runtime/activation_checkpointing/checkpointing.py index c8027c653a25..bfb8ea00fd2f 100644 --- a/deepspeed/runtime/activation_checkpointing/checkpointing.py +++ b/deepspeed/runtime/activation_checkpointing/checkpointing.py @@ -350,16 +350,13 @@ def merge_tensors(tensor_objects, non_tensor_objects, tensor_flags): real_tensor_flags = None - # remove the flags that are assigned to the size of the flattened tensors + # get_partitioned_activations_for_backward() saved a (value, size) pair for every argument, so + # the flags and the non-tensors both hold one extra entry per argument. A tensor argument pairs + # with a tensor size and a non-tensor argument pairs with None, so every argument owns an even + # index in each list and dropping the odd ones restores one entry per original argument. if PARTITION_ACTIVATIONS: - real_tensor_flags = [] - previous_flag = False - for flag in tensor_flags: - if previous_flag: - previous_flag = False - continue - previous_flag = flag - real_tensor_flags.append(flag) + real_tensor_flags = tensor_flags[::2] + non_tensor_objects = non_tensor_objects[::2] else: real_tensor_flags = tensor_flags diff --git a/tests/unit/runtime/activation_checkpointing/test_activation_checkpointing.py b/tests/unit/runtime/activation_checkpointing/test_activation_checkpointing.py index 4fe83cc499a6..2462e748797d 100644 --- a/tests/unit/runtime/activation_checkpointing/test_activation_checkpointing.py +++ b/tests/unit/runtime/activation_checkpointing/test_activation_checkpointing.py @@ -496,3 +496,61 @@ def test_configure_with_contiguous_checkpointing_requires_num_checkpoints(): cp.mpu, cp.deepspeed_checkpointing_enabled, ) = saved + + +def _partitioned_backward_args(args): + """Run an argument list through the save and restore path `partition_activations` uses. + + This mirrors `CheckpointFunction.forward`/`backward` without an accelerator, so it runs on CPU. + """ + cp = deepspeed.checkpointing + saved = _snapshot_ckpt_config() + saved_mp = (cp.mp_group, cp.mp_size) + try: + cp.PARTITION_ACTIVATIONS = True + cp.CONTIGUOUS_CHECKPOINTING = False + cp.mp_group, cp.mp_size = None, 1 + + inputs = tuple(a.clone() if torch.is_tensor(a) else a for a in args) + new_args = cp.get_partitioned_activations_for_backward(list(args), inputs, False) + tensor_args, non_tensor_args, tensor_flags = cp.extract_tensors(all_objects=tuple(new_args)) + + for tensor in tensor_args: + if tensor is not None and getattr(tensor, 'saved_data', None) is not None: + tensor.data = tensor.saved_data.to(tensor.device) + tensor.saved_data = None + + gathered = cp.gather_partitioned_activations(tensor_args) + return cp.merge_tensors(tensor_objects=gathered, non_tensor_objects=non_tensor_args, tensor_flags=tensor_flags) + finally: + _restore_ckpt_config(saved) + cp.mp_group, cp.mp_size = saved_mp + + +@pytest.mark.parametrize('non_tensor', [None, 2, True, 2.5, (None, 2.5)]) +def test_partitioned_non_tensor_args_survive_the_round_trip(non_tensor): + """The recompute must be handed the arguments the forward pass received, not extra `None`s.""" + tensor = torch.rand(HIDDEN_DIM, requires_grad=True) + args = (tensor, non_tensor) + + merged = _partitioned_backward_args(args) + + assert len(merged) == len(args) + assert torch.is_tensor(merged[0]) + if non_tensor is None: + assert merged[1] is None + else: + assert merged[1] == non_tensor + + +def test_partitioned_args_keep_their_order_around_a_non_tensor(): + """A non-tensor argument must not shift the arguments that follow it.""" + first = torch.rand(HIDDEN_DIM, requires_grad=True) + second = torch.rand(HIDDEN_DIM, requires_grad=True) + args = (first, None, True, second) + + merged = _partitioned_backward_args(args) + + assert [torch.is_tensor(item) for item in merged] == [True, False, False, True] + assert merged[1] is None + assert merged[2] is True