From 781f2e2cc6daca8ba9886d1a797f19082c247389 Mon Sep 17 00:00:00 2001 From: taking-lying-flat <1615405@qq.com> Date: Sun, 6 Sep 2026 18:01:47 +0800 Subject: [PATCH] Fix SequenceTiledCompute backward for empty trailing shards Signed-off-by: taking-lying-flat <1615405@qq.com> --- .../runtime/sequence_parallel/ulysses_sp.py | 13 ++-- tests/unit/ulysses_alst/test_tiled_compute.py | 61 +++++++++++++++++++ 2 files changed, 69 insertions(+), 5 deletions(-) diff --git a/deepspeed/runtime/sequence_parallel/ulysses_sp.py b/deepspeed/runtime/sequence_parallel/ulysses_sp.py index c341d443726b..bd8d23caede8 100644 --- a/deepspeed/runtime/sequence_parallel/ulysses_sp.py +++ b/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: diff --git a/tests/unit/ulysses_alst/test_tiled_compute.py b/tests/unit/ulysses_alst/test_tiled_compute.py index 2a4fc8f6a797..79526ea80c4a 100644 --- a/tests/unit/ulysses_alst/test_tiled_compute.py +++ b/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: