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 <alanhuangyoo@gmail.com>

Author: alanhuangyoo

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)