[Connor-XY]

With deterministic_mode, the TE operation-fuser grouped MLP must not accumulate the routing-probability gradient with atomics. TE's CuTe DSL fused grouped MLP multiplies the activation by the routing probabilities inside the fused kernel, and the cuDNN grouped-GEMM backward it calls accumulates the probability gradient across CTAs with atomics. When FC2 has no bias, FC2(p * act) == p * FC2(act), so this change passes unit probabilities to the fused operations and scales the FC2 output instead. The probability gradient then becomes an ordinary PyTorch reduction. FC2 bias and NCCL-EP zero-copy output buffers cannot use this rewrite, so they raise a RuntimeError in deterministic mode instead of silently running the atomic path. Without deterministic_mode nothing changes, including when only PyTorch's global determinism flag is set.

This was split out of #7708 (deterministic DSv4 sparse attention) because it touches only the MoE grouped MLP and does not depend on the DSv4 changes. It is based on main.

Tests:

• test_grouped_mlp.py: CPU tests drive the real TEGroupedMLP._fused_forward with a stand-in fused operation. They check unit probabilities reaching the fused op, exact outputs and input/probability gradients against an independent formula with byte-identical replay, the FC2-bias and zero-copy errors, and that the torch flag alone leaves bias, probabilities and zero-copy output storage untouched. The existing fixture config gains the deterministic_mode field.
• test_moe_kernels.py::test_te_grouped_mlp_replays_on_uneven_experts gains an operation-fuser variant (TE >= 2.14, GB200 lane). It uses the weighted squared-ReLU activation, which the operation fuser accepts without TE's CuTe DSL selector variable, and checks bit-exact replay of the deterministic path on uneven expert loads. In BF16, TE runs the grouped operations unfused, so this case checks the end-to-end wiring; it passes with or without this change. The fused CuTe DSL kernel is selected only through an environment variable that TE reads when it is imported, so no unit test here can enable it.

Validation:

• On four GB300 ranks with the normal repository conftest: test_grouped_mlp.py (61 passed; 9 skipped because the installed TE lacks native single grouped bias or the fused unary grouped MLP), both replay variants and test_manifest.py pass. Black, isort, Ruff, Pylint, the kernel determinism coverage gate and the copyright check pass.
• A one-off GB300 check exercised the real fused kernel (TE 2.17, MXFP8, SwiGLU with 32-wide GLU interleaving, 8 uneven experts, hidden 2048, FFN 4096, the CuTe DSL selector set before start-up, side-stream contention). Without this change, the probability gradient took 8 and 6 distinct values over 12 replays (deterministic mode off and on). With this change and deterministic_mode, all 12 replays were bit-identical; outputs, input gradients and all 16 weight gradients were bit-identical in every arm.

TE compatibility: TE builds that include NVIDIA/TransformerEngine#3407 (TE main since 2026-09-04; not in 2.17 through 2.19) check determinism inside the fused grouped MLP. With deterministic algorithms requested, they compute a deterministic probability gradient for scaled SReLU (cuDNN frontend >= 1.28), but raise for GLU activations, because TE cannot tell that the scales passed here are all ones (it recognizes unit scales only for single-group fusions). On those builds, deterministic SwiGLU MoE training needs the operation fuser disabled until TE skips that gradient when the scales do not require grad.

Performance: only deterministic mode is affected. The fused activation runs with unit scales, and one elementwise multiply scales the FC2 output. That multiply keeps the FC2 output (tokens x hidden) for the probability gradient and briefly allocates an FP32 product of the same size. This has not been benchmarked.

🤖 Generated with Claude Code

[github-actions]

This PR has been automatically converted to draft because all PRs must start as drafts.

When you are ready for review, click Ready for Review to begin the review process. This will:

1. Add the oncall reviewer (optional reviewer)
2. Add required review teams based on your changes

See the contribution guide for more details.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[Connor-XY]

/ok to test 85762f9

[Connor-XY]

/ok to test ba72913

[yaox12]

/review model=claude mode=strict

[svcnvidia-nemo-ci]

Strict review: PR #7769, deterministic operation-fuser grouped-MLP probability gradients

Findings: CRITICAL: 1 · IMPORTANT: 0 · SUGGESTION: 0. Risk: moderate. The rewrite is correct for the default BF16 combine path. It breaks one configuration that the config validation currently accepts.

Most impactful finding

• [CRITICAL Correctness] The deterministic rewrite corrupts the MXFP8 combine-backward gradient (inline on experts.py#L901).
  • Trigger: deterministic_mode together with moe_flex_dispatcher_backend="ncclep" and moe_combine_bwd_dtype="mxfp8", without zero-copy.
  • Cause: FC2 expects its output gradient to be the opaque MXFP8 carrier, unmodified (experts.py#L738-L744). The new post-FC2 multiply and cast run their backward on those packed bytes first.
  • Effect: expert gradients and the new probability gradient are garbage, with no error.
  • Fix: add one more guard next to the zero-copy guard.

Checked and found correct

• Math: FC2(p·act) == p·FC2(act) holds for a bias-free grouped linear. Padded rows carry zero probabilities, so they contribute nothing either way.
• Gradient dtype: the probability gradient is a per-row fp32 PyTorch reduction (bf16 output promoted by the fp32 multiply). That is deterministic and has the right dtype.
• Off by default: with deterministic_mode=False (including when only the torch global flag is set), bias, probabilities and zero-copy storage are untouched. test_torch_flag_alone_preserves_fused_forward covers this.
• Paged stash: it only redirects saved tensors that carry grouped_tensor_scale_inv, so the multiply's saved tensors stay on the normal autograd path (paged_stash.py#L712-L723).
• Config: deterministic_mode is a ModelParallelConfig field. Enabling the op fuser forces moe_use_grouped_tensor=True (transformer_config.py#L1695-L1696), so the new replay variant reaches _fused_forward.
• Test wiring: squared_relu and is_te_min_version are imported in test_moe_kernels.py. launch_on_gb200 is an existing marker. The manifest registers tests by file, so parametrizing the test does not affect it.
• Existing fixture: the mock in test_fused_forward_caches_ops_and_forwards_expected_arguments now sets deterministic_mode=False. The other _fused_forward callers I checked use real TransformerConfigs.

The memory and compute cost in deterministic mode (FC2 output kept for the probability gradient, a transient fp32 product) and the TE#3407 GLU limitation are documented, intentional trade-offs, so I have not raised them as findings. I could not run tests; everything above comes from reading the source.

Note: formal approval is not available for this request, so this review is posted as a comment only.

Review generated by claude

[svcnvidia-nemo-ci]

[CRITICAL Correctness] The deterministic rewrite corrupts the MXFP8 combine-backward gradient

Failure.

Trigger configuration:

• deterministic_mode=True
• moe_token_dispatcher_type="flex", moe_flex_dispatcher_backend="ncclep"
• use_transformer_engine_op_fuser=True, moe_grouped_gemm=True, add_bias_linear=False
• moe_combine_bwd_dtype="mxfp8", with no zero-copy

Eager ncclep is the validated fp8-wire mode, and eager rejects zero-copy (token_dispatcher.py#L1558-L1563). So output_buffer is None, and neither guard above fires. transformer_config.py#L2137-L2151 accepts this configuration, and determinism.py does not reject ncclep.

Why it breaks:

1. With this wire dtype, _make_fused_ops sets op.ep_mxfp8_carrier_grad = True on FC2, because "the FC2 output grad arrives as an opaque MXFP8 carrier tensor" (L738-L744). The carrier is E4M3 data plus scales packed into a plain tensor's storage (L629-L631).
2. Before this PR, nothing numeric sat between FC2 and nccl_ep_combine, so the carrier reached FC2 unmodified.
3. Now the backward of (output * post_probs.unsqueeze(-1).float()).to(output.dtype) (L932-L933) runs first:
   • the .to backward casts the packed bytes to fp32 as if they were bf16 values;
   • the mul backward scales them by p and hands FC2 a new bf16 tensor;
   • FC2 then decodes that tensor as a carrier.

Result: expert dgrad/wgrad and the new probability gradient sum(carrier * fc2_out) come out as garbage (or NaN/Inf), with no error. Expected: either a correct gradient or an explicit rejection, as in the zero-copy case.

Correction. In TEGroupedMLP._fused_forward, inside if self.config.deterministic_mode: and next to the zero-copy guard, reject this configuration before any op runs. The invariant to restore: whenever FC2 has ep_mxfp8_carrier_grad=True, no autograd op may sit between the FC2 output and the combine.

                if getattr(self.config, 'moe_combine_bwd_dtype', 'bf16') == 'mxfp8':
                    raise RuntimeError(
                        "deterministic operation-fuser grouped MLP does not support "
                        "moe_combine_bwd_dtype='mxfp8': the FC2 output gradient arrives as an "
                        "MXFP8 carrier that must reach FC2 unmodified."
                    )

Using getattr matches L743 and keeps the SimpleNamespace configs in test_grouped_mlp.py, which lack this field, working. A config-time ValueError would also work, but a runtime guard matches the two existing ones. Please verify the code before adopting it; I could not run it.

Preservation.

• Keep the check inside the deterministic_mode branch. Non-deterministic mxfp8-wire runs pass the carrier straight to FC2 and must keep working.
• Do not also gate on moe_dispatch_fwd_dtype. The FC1-input carrier is untouched by this rewrite, so deterministic runs that set only moe_dispatch_fwd_dtype='mxfp8' should stay allowed.
• Deterministic runs with the default bf16 combine-backward dtype keep the post-FC2 scaling.

Acceptance.

• Proposed test: add "mxfp8_combine_bwd" to unsupported in test_deterministic_fused_forward_probability_gradients, setting moe_combine_bwd_dtype="mxfp8" on its config. Expect RuntimeError matching moe_combine_bwd_dtype, and assert not seen.
• Existing test: the unsupported=None case still matches the reference formula with byte-identical replay.
• Existing test: test_torch_flag_alone_preserves_fused_forward still passes, including seen["probs"] is probabilities.
• Source invariant to inspect: with deterministic_mode=False and mxfp8 combine-backward, post_probs stays None.