[REGRESSION] MoE step-time regression in core_r0.18.0 vs core_v0.18.0 (deepseekv3_proxy CE te→native; nemotron3 TE grouped-MLP/SReLU bump)

**Describe the regression**

Two MoE release recipes regress in step time between `core_v0.18.0` and the
release branch `core_r0.18.0`. `core_r0.18.0` is a clean fast-forward of the tag
(merge-base = `core_v0.18.0`, +8 commits), so the surface is small and the two
cases have **distinct root causes**:

- **Case A — `deepseekv3_proxy_flex_*` (4 recipes):** the recipe itself was
changed from TE-fused cross-entropy to native, which is the slower path.
- **Case B — `nemotron3_super_release_gb200*` (2 recipes):** the recipe is
**byte-identical** across the two refs, so its regression is entirely due to the
bundled Transformer Engine bump landing on the grouped-MLP / SReLU / NVFP4
path this model exercises.

cc @NVIDIA/mcore-oncall

---

### Case A — `deepseekv3_proxy_flex_*` (cross-entropy fusion `te` → `native`)

The only functional Megatron change in the range is the cherry-pick of #5115
("Disable TE cross entropy loss fusion") + #5162 ("Move TE cross entropy guard to
training args"). These:

1. hard-`assert` against `--cross-entropy-loss-fusion --cross-entropy-fusion-impl te`
(`megatron/training/arguments.py`), and
2. flip every affected release recipe from `te` → `native`.

The 4 deepseekv3 recipes changed **only** that one line:

```yaml
# tests/functional_tests/test_cases/mixtral/deepseekv3_proxy_flex_tp1pp4emp16etp1cp1_release/model_config.yaml
--cross-entropy-loss-fusion: true
- --cross-entropy-fusion-impl: te
+ --cross-entropy-fusion-impl: native
```

TE-fused CE was the recommended perf path (it sat in the MoE performance-flags
list in `megatron/core/transformer/moe/README.md`, which this range also edits to
`native`). Moving to `native` is a **deliberate throughput trade for stability**,
so a regression here is expected — but it should be quantified against the release
dashboard, and if the delta is large we may want a stable fused-CE option rather
than a blanket ban. Affected recipes:

- `mixtral/deepseekv3_proxy_flex_tp1pp4emp16etp1cp1_release{,_sm}`
- `mixtral/deepseekv3_proxy_flex_tp2pp2emp16etp1cp1_gb_200_release{,_sm}`

Note deepseekv3 is bf16, so the NVFP4 TE changes below do **not** apply to it; its
only other exposure is the grouped-MLP metadata/pointer refactors (TE #3075 /
#3076 / #3001).

### Case B — `nemotron3_super_release_gb200*` (TE grouped-MLP / SReLU / NVFP4 bump)

`nemotron3_super_release_gb200/model_config.yaml` is **identical** on both refs
(and already uses `--cross-entropy-fusion-impl native`), and the NGC base image is
unchanged (`nvcr.io/nvidia/pytorch:26.04-py3` on both → same PyTorch/CUDA/NCCL).
The only thing that moved for this recipe is the TE pin (#5518,
`release_v2.16.post`):

```
transformer-engine 4220403e → b9d690e0 # NVIDIA/TransformerEngine
```

That TE range is **9 commits, all on the grouped-MLP / grouped-GEMM / NVFP4 /
SReLU path** — i.e. nemotron3's expert compute hot path. Recipe flag → TE commit:

| nemotron3 flag | TE commit |
| --- | --- |
| `--squared-relu` + `--use-fused-weighted-squared-relu` | **#2981 "GGEMM+srelu kernels for MxFP8 Nemotron"** — adds grouped-MLP SReLU fusion and **wires ScaledSReLU recompute into grouped MLP** |
| `--fp4-format e2m1` / `--fp4-recipe nvfp4` + `te_quant.cfg` | **#3048 "Enable NVFP4 fused grouped MLP"**, **#2972 "4over6 NVFP4 recipe"** |
| `--moe-grouped-gemm true` | #3075 / #3076 / #3001 (grouped-split metadata + GPU pointer prep) |
| `--enable-cuda-graph` / `--cuda-graph-modules "[mamba attn]"` | **#3038 "Make `modules.GroupedLinear` graph-safe"** |

**Prime suspect:** #2981's "Wire ScaledSReLU recompute in grouped MLP." Recompute
trades activation memory for extra backward FLOPs — a classic step-time
regression — and it lands exactly on nemotron3's fused-SReLU grouped-GEMM path.
Runner-up: #3048 may switch nemotron3 onto a newly-enabled NVFP4 fused
grouped-MLP path that did not exist on the tag TE.

**Previous performance**

`iteration-time` for the two release recipes at `core_v0.18.0`
(TE `4220403e`, deepseekv3 on `--cross-entropy-fusion-impl te`). Exact numbers
are in the `megatron-core-release-runs` W&B project — to be attached.

**New performance**

`iteration-time` regressed at `core_r0.18.0` (TE `b9d690e0`, deepseekv3 on
`native`). Numbers to be attached from the same dashboard. For nemotron3, the
recipe also logs `mem-allocated-bytes` / `mem-max-allocated-bytes`: **time↑ with
mem↓ is the recompute signature** and would confirm #2981.

**To Reproduce**

- **Case A (isolate CE):** on `core_v0.18.0`, run
`deepseekv3_proxy_flex_tp1pp4emp16etp1cp1_release` with
`--cross-entropy-fusion-impl te` vs `native`. That delta is the CE switch; any
residual vs `core_r0.18.0` is the TE bump.
- **Case B (isolate TE):** the nemotron3 recipe is frozen, so bisect the 9 TE
commits `4220403e..b9d690e0` against it (seed at #2981), or compare the memory
metrics above.

**Environment**

- Previous Megatron-LM commit: `ba7b5ebce12af60627a80985792a1449ce45f46c` (`core_v0.18.0`)
- New Megatron-LM commit: `458c8d0ecafdf6d9e36771600d62ade27f2a67b7` (`core_r0.18.0`)
- Transformer Engine: `4220403e831d29e93868f7793693ea83f6b8b05b` → `b9d690e042b1c4e455214e7dab65d6d3512c05d6`
- NGC base image: `nvcr.io/nvidia/pytorch:26.04-py3` (unchanged → PyTorch/CUDA/NCCL unchanged)

**Proposed fix**

- **Case A:** intended stability trade — quantify the throughput cost and decide
whether a vetted fused-CE path should remain available instead of a hard reject.
- **Case B:** confirm which TE commit regresses nemotron3 (bisect / memory
signature). If ScaledSReLU recompute (#2981) is default-on where it is not
needed, gate it off for this config; if #3048's newly-enabled NVFP4 fused path
is the cause, evaluate reverting or tuning it.

Start with the affected model_config.yaml recipes, megatron/training/arguments.py, and the MoE performance-flags README entry. Run the specified deepseekv3 TE-versus-native comparison, then bisect Transformer Engine commits 4220403e..b9d690e0 against the frozen nemotron3 recipe while checking iteration-time and memory metrics in the release dashboard. Done means the responsible change is isolated and a validated mitigation or quantified stability trade is documented.

Written by the indexing model from the issue text.