feat(moe+ssm): expert-LoRA GEMM stack and hybrid-recompute guard on the d12f6c8c base

[taufeeque9]

Written by Claude (on Taufeeque's behalf).

Why this PR exists

NeMo-RL's farai/main no longer carries a standalone Megatron-LM submodule: after the upstream sync, megatron-core flows in through the Megatron-Bridge fork's nested submodule pin, currently upstream d12f6c8c (Aug 2026). Our fork's features — which every Nemotron-3-Super hackable-env RL run depends on — lived only as feature branches on top of a February 2026 upstream base, so they silently fell out of the stack that main consumes. Concretely, on main today:

• expert-LoRA training with the frozen torch grouped GEMM backend (moe_expert_gemm_backend: torch) doesn't exist — the run configs are refused by a guard in nemo-rl (AlignmentResearch/nemo-rl#60);
• expert_fc1_act selective recompute doesn't exist — mcore rejects it at config validation;
• full-granularity recompute on the hybrid (Mamba) stack — the fix that unlocked 32k+ context (validated on B200: selective OOMs at 32k, full recompute trains 16k/32k/48k) — exists upstream but without the frozen-embedding requires_grad guard, so adapter-only (LoRA) training would silently receive zero gradients: re-entrant checkpointing only attaches a grad_fn when a tensor input requires grad, and a frozen embedding provides none.

This PR ports the full stack onto exactly the base the Bridge pins, so the features re-enter main's supply chain with a minimal delta.

What's in it

Nine commits, each carrying (cherry picked from ...) provenance to the original fork commits:

1. perf(moe): add frozen torch grouped GEMM backend — runs frozen expert base weights through torch._grouped_mm outside TE autograd (memory/perf for frozen-base + expert-LoRA).
   2–3. weight-view validation + zero-token grouped-split test contracts.
2. perf(moe): checkpoint expert FC1 activation (expert_fc1_act) — restructured around upstream's new offload managers; the non-recompute branch is upstream's code verbatim.
3. perf(moe): fuse shared expert LoRA projection.
   6–7. grouped-offsets dtype + BF16 fused-input-rounding test fixups.
4. fix(recompute): keep chunk outputs differentiable with a frozen embedding — the two original SSM commits collapsed to just this guard; upstream had absorbed the rest (hybrid full recompute now lives in models/hybrid/hybrid_block.py → core/recompute.py).
5. style pass matching repo black/pylint conventions.

Two guards added beyond the originals, both preventing silent no-ops on paths that didn't exist at the old base: expert_fc1_act is rejected under use_transformer_engine_op_fuser, and moe_expert_gemm_backend='torch' is rejected with moe_single_grouped_weight=True (explicit error instead of an AttributeError; supporting that combination needs TE ≥ 2.14 work).

Gating

Do not merge before GPU validation (the port is statically verified only). Named items: the FC1-recompute parity/call-count test, the torch grouped-GEMM path end-to-end (SM100+), the deeper-nested activation-offload branch, FP8 + frozen-embedding interaction with the guard. Plan: nemo-rl topology benchmark (selective + full arms, 16k/32k) on the B200 reservation.

Chain

This PR → AlignmentResearch/megatron-bridge-contrib#4 (bumps the nested pin to this branch) → one-line Bridge-pin bump in nemo-rl, after which the guards in nemo-rl#60 pass instead of refusing.

Note: farai/main was created for this PR at d12f6c8c — the exact upstream commit the Bridge pin consumes — establishing the fork's internal integration branch per our convention (the fork's main stays a pure upstream mirror).

🤖 Generated with Claude Code

[taufeeque9]

Written by Claude.

The fork-base scaffolding (manifest, tooling, CI prune) is in #2, which should merge first. Once it's in, this PR's merge ref will carry the manifest and the fork-base check will apply here too.

[taufeeque9]

Written by Codex.

Superseded by #3, which keeps only the validated frozen grouped-GEMM backend and the full-recompute frozen-input correctness guard. The fused expert-LoRA and expert_fc1_act selective-recompute patches were intentionally pruned.