fix(batch-invariant): full-reduction mean returned the sum

Under batch-invariant mode, `x.mean()` and `torch.mean(x)` reach the
`aten::mean.dim` override (`mean_batch_invariant`) with an empty dim list.
The multi-dim branch divided the sum by the product over that empty list
(1), so every full-reduction mean returned the sum; `dim=None` raised.
Autograd was unaffected (the `mean` derivative divides by numel itself),
but every forward value of a full mean was wrong.

Treat an empty or None dim as all dimensions, as aten::mean does.

Signed-off-by: wdykas <wdykas@nvidia.com>

Author: wdykas

megatron/core/transformer/custom_layers/batch_invariant_kernels.py

@@ -582,6 +582,9 @@ def _log_softmax_batch_invariant(input, dim, _half_to_float):
 def mean_batch_invariant(input, dim, keepdim=False, dtype: torch.dtype | None = None):
     """Batch-invariant replacement for `aten::mean.dim` over one or more dimensions."""
     assert dtype is None or dtype == torch.float32, f"unsupported dtype: {dtype}"
+    # A full reduction (`x.mean()`, `torch.mean(x)`) arrives with an empty dim list.
+    if dim is None or len(dim) == 0:
+        dim = list(range(input.ndim))
     if len(dim) == 1:
         return mean_dim(input, dim[0], keepdim=keepdim)
     else:

tests/unit_tests/transformer/test_te_layers_batch_invariant.py

@@ -19,6 +19,7 @@
 from megatron.core.transformer.custom_layers.batch_invariant_kernels import (
     HAVE_DEEPGEMM_BF16,
     assert_te_supports_batch_invariant_attention,
+    mean_batch_invariant,
     set_batch_invariant_mode,
     te_supports_batch_invariant_attention,
 )
@@ -65,6 +66,29 @@ def test_batch_invariant_rejects_incompatible_te(monkeypatch):
         assert_te_supports_batch_invariant_attention()
 
 
+@pytest.mark.parametrize("dim", ([], (), None))
+@pytest.mark.parametrize("keepdim", (False, True))
+def test_mean_batch_invariant_full_reduction(dim, keepdim):
+    """An empty (or None) dim list reduces over every dimension, as aten::mean does."""
+    x = torch.arange(24, dtype=torch.float32).reshape(2, 3, 4)
+    out = mean_batch_invariant(x, dim, keepdim=keepdim)
+    torch.testing.assert_close(out, x.mean(dim=(0, 1, 2), keepdim=keepdim))
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
+def test_full_mean_under_batch_invariant_mode():
+    """`x.mean()` / `torch.mean(x)` dispatch to aten::mean.dim with dim=[]."""
+    x = torch.randn(4, 8, device="cuda", requires_grad=True)
+    ref = x.detach().sum() / x.numel()
+    with set_batch_invariant_mode(True):
+        out = x.mean()
+        out_fn = torch.mean(x)
+        out.backward()
+    torch.testing.assert_close(out.detach(), ref)
+    torch.testing.assert_close(out_fn.detach(), ref)
+    torch.testing.assert_close(x.grad, torch.full_like(x, 1.0 / x.numel()))
+
+
 # ============================================================================
 # Batch-Invariant test helpers
 # ============================================================================