Fix SequenceTiledCompute backward for empty trailing shards

deepspeed/runtime/sequence_parallel/ulysses_sp.py

@@ -821,6 +821,7 @@ def forward(
 
         with torch.no_grad():
             shard_step = math.ceil(seqlen / shards)
+            ctx.shard_step = shard_step
             output_shards = []
 
             for i in range(shards):
@@ -877,8 +878,6 @@ def backward(ctx, *grads) -> torch.Tensor:
         else:
             grad_requiring_tensor_grad = torch.empty_like(grad_requiring_tensor)
 
-        kwargs_to_shard_shards = {k: list(torch.chunk(v, chunks=shards, dim=1)) for k, v in kwargs_to_shard.items()}
-
         for i in range(shards):
             # when fn involves one or more model weights deepspeed will normally push a grad to
             # reduce per sub-module call, so since we only want it to add a grad for the last
@@ -894,14 +893,18 @@ def backward(ctx, *grads) -> torch.Tensor:
                     for param in compute_params:
                         param.ds_grad_is_ready = True
 
-            kwargs_to_shard_shard = {k: v[i] for k, v in kwargs_to_shard_shards.items()}
+            # Match forward's empty trailing slices, with offsets that remain valid for narrow().
+            shard_offset = min(i * ctx.shard_step, ctx.seqlen)
+            kwargs_to_shard_shard = {
+                k: v[:, shard_offset:shard_offset + ctx.shard_step]
+                for k, v in kwargs_to_shard.items()
+            }
             grad_requiring_tensor_shard = kwargs_to_shard_shard[grad_requiring_tensor_key]
 
             grad_requiring_tensor_shard.requires_grad_(grad_requiring_tensor_requires_grad)
 
             # if seqlen is not exactly divisible by shards the last step will be shorter than shard_step
-            shard_step = kwargs_to_shard_shards[grad_requiring_tensor_key][i].shape[1]
-            shard_offset = i * kwargs_to_shard_shards[grad_requiring_tensor_key][0].shape[1]
+            shard_step = grad_requiring_tensor_shard.shape[1]
 
             if grad_requiring_tensor.shape[0] == 1:
                 # on narrow the shard's stride is unaffected with dim0==1 (bs) so we use the most efficient `narrow` alias:

tests/unit/ulysses_alst/test_tiled_compute.py

@@ -7,6 +7,7 @@
 """
 
 from deepspeed.runtime.sequence_parallel.ulysses_sp import TiledMLP, sequence_tiled_compute, TiledFusedLogitsLoss
+from deepspeed.accelerator import get_accelerator
 from deepspeed.utils import safe_get_full_grad
 from torch.nn import Linear, Module
 from unit.common import DistributedTest, preferred_dtype
@@ -231,6 +232,66 @@ def test_tiled_mlp(self, zero_stage, batch_size):
         torch_assert_close(x_grad_a, x_grad_c)
 
 
+@pytest.mark.parametrize("batch_size", [1, 2])
+@pytest.mark.parametrize("seqlen,shards", [(5, 4), (2, 4), (7, 4), (8, 4)])
+@pytest.mark.parametrize("output_reduction", [None, "sum", "mean"])
+def test_sequence_tiled_compute_shard_boundaries(batch_size, seqlen, shards, output_reduction):
+    # Empty trailing tiles must preserve both gradients and updates relative to untiled computation.
+    device = get_accelerator().device_name()
+    dtype = torch.float32
+    hidden_dim = 8
+    torch.manual_seed(42)
+    tiled_model = SimpleMLP(hidden_dim).to(device=device, dtype=dtype)
+    reference_model = SimpleMLP(hidden_dim).to(device=device, dtype=dtype)
+    reference_model.load_state_dict(tiled_model.state_dict())
+    tiled_optimizer = torch.optim.SGD(tiled_model.parameters(), lr=0.01)
+    reference_optimizer = torch.optim.SGD(reference_model.parameters(), lr=0.01)
+
+    def compute(x, scale, model):
+        output = mlp_forward_orig(model, x) * scale
+        return output if output_reduction is None else output.sum()
+
+    for _ in range(2):
+        tiled_optimizer.zero_grad()
+        reference_optimizer.zero_grad()
+        x = torch.rand((batch_size, seqlen, hidden_dim), device=device, dtype=dtype, requires_grad=True)
+        x_reference = x.detach().clone().requires_grad_(True)
+        scale = torch.rand_like(x)
+
+        output = sequence_tiled_compute(
+            compute,
+            seqlen,
+            shards,
+            kwargs_to_shard={
+                "x": x,
+                "scale": scale
+            },
+            kwargs_to_pass={"model": tiled_model},
+            grad_requiring_tensor_key="x",
+            compute_params=list(tiled_model.parameters()),
+            output_unshard_dimension=1 if output_reduction is None else 0,
+            output_reduction=output_reduction,
+        )
+        expected = mlp_forward_orig(reference_model, x_reference) * scale
+        if output_reduction is not None:
+            expected = expected.sum()
+            if output_reduction == "mean":
+                # Averaging per-tile sums divides the full sum by the requested tile count.
+                expected = expected / shards
+
+        torch.testing.assert_close(output, expected)
+        output.sum().backward()
+        expected.sum().backward()
+        torch.testing.assert_close(x.grad, x_reference.grad)
+        for param, reference_param in zip(tiled_model.parameters(), reference_model.parameters()):
+            torch.testing.assert_close(param.grad, reference_param.grad)
+
+        tiled_optimizer.step()
+        reference_optimizer.step()
+        for param, reference_param in zip(tiled_model.parameters(), reference_model.parameters()):
+            torch.testing.assert_close(param, reference_param)
+
+
 @pytest.mark.parametrize("shards", [2, 4])
 @pytest.mark.parametrize("layout", ["transposed", "channel_slice"])
 class TestTiledMLPInputLayout: