[wdykas]

Summary

Stacked on #7302. Opt in to TE device-metadata grouped GEMM for batch-invariant training using moe_use_grouped_tensor=True, reusing the existing 256-row expert padding. Dense and legacy grouped GEMMs keep their restricted workspace.

The grouped path needs the full workspace, with separate allocations for forward/dgrad/wgrad layouts. Restrict it to the numerically validated SM100 + cuBLASLt 13.5.1 combination and TE's GroupedLinear(use_grouped_tensor=...) API. Other GPU/library pairs fail explicitly. Fixed-batch backward replay is tested; gradients across different minibatches are not promised to be identical.

Also includes the isolated SwiGLU rounding fix pushed to #7302; that commit disappears from this diff when the base mirror updates.

Validation

• Qwen3-30B-A3B and NVIDIA-Nemotron-3-Nano-30B-A3B-BF16: two complete NeMo RL GRPO steps each, with exact zero raw token-logprob difference for scoring vs generation and training-forward vs generation (atol=0, rtol=0). Each model checked 64 valid response tokens in both comparisons, including a second refit/rollout.
• TP1/EP4 on four GB200 GPUs; routed expert projections only in MXFP8, first two/last four layers in BF16, shared experts and other modules in BF16. Torch/vLLM-selected inference. No R3 replay.
• Ten grouped forward/backward/graph cases passed, including Nano expert dimensions, changing expert sizes, empty experts and layout-specific dispatch.
• Four EP2 training/inference MoE parity cases, two activation tests, and selective-MXFP8 M2N tests for Torch/vLLM/FlashInfer passed locally.

Pinned test TE: 51317b440be0378f570ca6cc7dec2962f6e3f8df (2.20.0.dev0+51317b4), PyTorch 2.11.0+cu130, cuBLASLt 13.5.1. The full-model smoke tests use short eight-token responses, zero rewards and no model-level CUDA graphs; they are not convergence or throughput benchmarks. Kernel-level graphs are covered separately. NeMo RL integration fixes are on wdykas/RL:fix/minf-mxfp8-zero-logprob-parity.

Existing base-PR CI is not claimed green: its logs include GDN config-default drift, an MLA optimizer assertion, and an M2N illegal access/timeout. The M2N test passes with the local pinned stack; the CI-environment difference remains to be isolated.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[wdykas]

/ok to test 2153480

[wdykas]

Closing in favor of #7303, which now includes these TE device-metadata batch-invariant training changes and their tests.