perf(moe): add frozen grouped GEMM for adapter training

[taufeeque9]

Written by Codex.

Why this change

Nemotron-3 Super LoRA freezes the 128 local expert base weights. On B200, the pinned Transformer Engine grouped path launched roughly one small GEMM per expert, leaving substantial launch overhead. torch._grouped_mm executes the frozen base branch as one SM100 grouped CUTLASS kernel while the ordinary Bridge LoRA branch remains unchanged.

This path reduced the measured logprob-plus-update cycle from 21.640 s to 14.684 s on the production-shaped benchmark: 32.1% less time and 47.4% more throughput, with the same peak allocated memory.

B200 throughput evidence

The comparison used the same gb215 node, step-75 checkpoint, fixed production-mixed token batch, and TP4/PP1/EP4/ETP1 topology. Both arms used BF16 LoRA, ordinary all-to-all dispatch, a 32k dynamic token cap, and 121,088 input tokens per iteration.

Metric	Transformer Engine	torch._grouped_mm	Change
Update time	17.176 s	11.599 s	-32.5%
Update throughput	7,050 tok/s	10,440 tok/s	+48.1%
Logprob time	4.458 s	3.078 s	-31.0%
Logprob + update cycle	21.640 s	14.684 s	-32.1%
Cycle throughput	5,596 tok/s	8,246 tok/s	+47.4%
Update peak allocated	117.721 GiB	117.721 GiB	unchanged
Update peak reserved	123.465 GiB	123.504 GiB	+0.039 GiB

Artifacts:

• TE reference: Schmidt job 102251, tf-super-deepep-ab-v1-a4e9b214c0/super-topology-benchmark.json, tp4-alltoall arm.
• Torch candidate: Schmidt job 102265, tf-super-torch-gmm-v1-a518a324d7/super-topology-benchmark.json.
• Kernel profile: Schmidt job 102277, tf-super-torch-gmm-profile-v1-a518a324d7.

The parity gate reported zero relative-L2 and max-absolute error for both forward output and input gradient, including experts with zero assigned tokens. The runtime packed 80 grouped FC1/FC2 modules and relocated 52.5 GiB of frozen weights without adding checkpoint entries.

Why existing behavior remains safe

• transformer_engine remains the default. No model uses the new path unless its config explicitly selects moe_expert_gemm_backend="torch".
• The torch path accepts only grouped, frozen, bias-free BF16 expert weights. FP8, FP4, trainable base weights, non-MoE models, TE single-grouped weights, and the TE operation fuser fail during configuration or preparation.
• Existing Parameter objects and state-dict keys are preserved. Each expert parameter becomes a view into one non-persistent contiguous buffer.
• Tests cover TE-versus-torch output and input-gradient parity, empty batches, zero-token experts, checkpoint-schema preservation, and recovery after a middle expert weight is replaced.
• Downstream NeMo-RL tests compare ordinary Bridge LoRA over TE and torch base GEMMs, including output, input-gradient, and adapter-gradient parity.

Full-recompute correctness

The PR also carries the small frozen-input guard needed by PP1 adapter-only full recompute. Re-entrant checkpointing otherwise sees a frozen embedding output with requires_grad=False and can detach every checkpointed chunk from the adapters inside it.

A direct regression now freezes the block and input, installs a trainable low-rank adapter inside a fully checkpointed hybrid layer, runs backward, and requires nonzero gradients on every adapter parameter. The production 32k and 48k NeMo-RL runs use this full-recompute path; their recorded gradients are nonzero and their losses move.

Scope deliberately omitted

• Fused expert LoRA was measured at 53.137 s versus 53.192 s for ordinary LoRA and used about 1.16 GiB more peak memory, so it is not included.
• Selective expert_fc1_act recompute was faster than whole-MoE recompute at short context, but the 32k selective arm OOMed because it did not cover the 40 Mamba mixers. The tested 32k and 48k recipes use full recompute instead.

Validation status

• Black, isort, Ruff, Python compilation, and git diff --check pass.
• Downstream NeMo-RL PR #60 passes pre-commit, lockfile, fork-base, build-deploy, and GPU-gate checks with this Core pin.
• A fresh recursive Bridge clone resolves this exact Core commit.
• End-to-end current-image B200 validation is provided by NeMo-RL job 113785, which uses this exact Core head with moe_expert_gemm_backend: torch, PP1 full recompute, and ordinary LoRA. It migrated the step-75 adapter checkpoint and completed optimizer steps through at least step 31; W&B run as6tqeze records nonzero gradient norms and changing rewards. The isolated pytest job 113381 timed out while building its SIF and produced no test result, so it is not counted as evidence.

Commit signing

All commits carry DCO sign-off. Cryptographic signing was unavailable because this machine has no configured GPG or SSH signing identity.