From 307ff1572f2d769396ed661458149680c8a5f28e Mon Sep 17 00:00:00 2001 From: alanhuangyoo Date: Sun, 30 Aug 2026 09:28:57 +0800 Subject: [PATCH 1/2] Reshape instead of view in TiledFusedLogitsLoss #8348 fixed this in TiledMLP.backward: the flatten of batch and sequence into one axis needs a copy for a caller that hands in a non-contiguous activation, so it cannot be a view. TiledFusedLogitsLoss.forward does the same flatten, under the same comment, and is still on view: # flatten bs+seqlen to avoid having stride issues when narrowing into seqlen w/ bs>1 x = x.view(-1, *x.shape[2:]) y = y.view(-1, *y.shape[2:]) A transposed activation -- the layout #8348's test covers -- fails there: RuntimeError: view size is not compatible with input tensor's size and stride (at least one dimension spans across two contiguous subspaces) A channel slice happens to survive, because its row stride still admits the flatten, so only the transposed case is reachable today. y and mask go through the same flatten and are changed with it. The unflatten at the end stays a view: x_grad comes from zeros_like() of the already-flattened x, so it is contiguous by construction. Signed-off-by: alanhuangyoo --- .../runtime/sequence_parallel/ulysses_sp.py | 8 ++-- tests/unit/ulysses_alst/test_tiled_compute.py | 41 +++++++++++++++++++ 2 files changed, 46 insertions(+), 3 deletions(-) diff --git a/deepspeed/runtime/sequence_parallel/ulysses_sp.py b/deepspeed/runtime/sequence_parallel/ulysses_sp.py index c341d443726b..60c9120def35 100644 --- a/deepspeed/runtime/sequence_parallel/ulysses_sp.py +++ b/deepspeed/runtime/sequence_parallel/ulysses_sp.py @@ -1145,10 +1145,12 @@ def forward( bs, seqlen = x.shape[:2] # flatten bs+seqlen to avoid having stride issues when narrowing into seqlen w/ bs>1 - x = x.view(-1, *x.shape[2:]) - y = y.view(-1, *y.shape[2:]) + # reshape rather than view: a caller may pass a non-contiguous x, y or mask (a transposed or + # channel-sliced activation), for which view cannot produce the flattened shape + x = x.reshape(-1, *x.shape[2:]) + y = y.reshape(-1, *y.shape[2:]) if mask is not None: - mask = mask.view(-1) + mask = mask.reshape(-1) incoming_grad = torch.tensor(1.0, dtype=x.dtype, device=x.device) # we are faking the incoming gradient, and since we perform a reduction outside of `autograd.backward` below we need to pre-adjust the incoming gradient. in the case of "sum" the gradient is 1.0, in the case of "mean" it's 1.0/num_elements, which in this case is 1/shards. diff --git a/tests/unit/ulysses_alst/test_tiled_compute.py b/tests/unit/ulysses_alst/test_tiled_compute.py index 2a4fc8f6a797..5e670c830bac 100644 --- a/tests/unit/ulysses_alst/test_tiled_compute.py +++ b/tests/unit/ulysses_alst/test_tiled_compute.py @@ -401,3 +401,44 @@ def loss_fn(self, x, y): # restore MyModel.forward = MyModel.forward_orig + + +@pytest.mark.parametrize("shards", [2, 4]) +class TestTiledFusedLogitsLossInputLayout: + """ + Same caller contract as TestTiledMLPInputLayout, on the loss. TiledFusedLogitsLoss flattens batch and + sequence into one axis before sharding, and a transposed activation cannot be flattened by a view. + """ + + def make_model(self, hidden_dim, vocab_size, dtype): + model = Linear(hidden_dim, vocab_size, bias=False, dtype=dtype) + torch.nn.init.normal_(model.weight, std=0.02) + return model + + @staticmethod + def loss_fn(model, x_shard, y_shard): + return torch.nn.functional.cross_entropy(model(x_shard), y_shard, reduction="sum") + + def test_transposed_input_matches_a_contiguous_copy(self, + shards, + batch_size=2, + seqlen=12, + hidden_dim=16, + vocab_size=32): + dtype = torch.float32 + torch.manual_seed(0) + # [bs, hidden, seqlen] transposed into [bs, seqlen, hidden] keeps the original strides + source = torch.rand((batch_size, hidden_dim, seqlen), dtype=dtype) + strided = source.transpose(1, 2).detach().requires_grad_(True) + assert not strided.is_contiguous(), "input is contiguous, so it does not exercise the flattening" + contiguous = strided.detach().clone().contiguous().requires_grad_(True) + y = torch.randint(0, vocab_size, (batch_size, seqlen)) + + model = self.make_model(hidden_dim, vocab_size, dtype) + losses = [] + for x in (strided, contiguous): + model.zero_grad() + loss = TiledFusedLogitsLoss.apply(self.loss_fn, model, x, y, None, shards, list(model.parameters()), "sum") + losses.append(loss) + + torch_assert_close(losses[0], losses[1]) From ac714ffef3290b6d0b8a0aece327171317a59d88 Mon Sep 17 00:00:00 2001 From: alanhuangyoo Date: Fri, 4 Sep 2026 12:21:20 +0800 Subject: [PATCH 2/2] Assert the gradients in the tiled-logits layout test The flatten under test also feeds the backward: x_grad is a zeros_like of the flattened activation, scattered per shard through x_grad.narrow(...).view_as(x_shard) and unflattened on the way out. The test compared only the loss and never called backward, so that path was not exercised and the copy reshape makes for a non-contiguous input was never checked for putting the gradient back where it came from. shards=2 loss_equal=True x_grad_equal=True max|dx|=0.000e+00 max|dW|=0.000e+00 shards=4 loss_equal=True x_grad_equal=True max|dx|=0.000e+00 max|dW|=0.000e+00 Also asserts the parameter gradients, which is the third comparison TestTiledMLPInputLayout makes. Signed-off-by: alanhuangyoo --- tests/unit/ulysses_alst/test_tiled_compute.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/tests/unit/ulysses_alst/test_tiled_compute.py b/tests/unit/ulysses_alst/test_tiled_compute.py index 5e670c830bac..5a4fc14c2e6b 100644 --- a/tests/unit/ulysses_alst/test_tiled_compute.py +++ b/tests/unit/ulysses_alst/test_tiled_compute.py @@ -435,10 +435,17 @@ def test_transposed_input_matches_a_contiguous_copy(self, y = torch.randint(0, vocab_size, (batch_size, seqlen)) model = self.make_model(hidden_dim, vocab_size, dtype) - losses = [] + losses, param_grads = [], [] for x in (strided, contiguous): model.zero_grad() loss = TiledFusedLogitsLoss.apply(self.loss_fn, model, x, y, None, shards, list(model.parameters()), "sum") + # The backward scatters into a zeros_like of the same flattened activation, so the + # gradient path runs through the flatten under test as well as the forward. + loss.backward() losses.append(loss) + param_grads.append([p.grad.detach().clone() for p in model.parameters()]) torch_assert_close(losses[0], losses[1]) + torch_assert_close(strided.grad, contiguous.grad) + for grad_a, grad_b in zip(*param_grads): + torch_assert_close(grad_a, grad_b)