Reshape instead of view in TiledFusedLogitsLoss

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.

tests/unit/ulysses_alst/test_tiled_compute.py

@@ -401,3 +401,51 @@ 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, 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)