Fix the batch-invariant RMSNorm forward and its input gradient

[huthvincent]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Fixes two defects in BatchInvariantRMSNormFn, one commit each. Both date from Batch Invariance (#2308) and neither has been touched since, so this is not a regression report.

1. The forward drops zero_centered_gamma. It computes weight_eff = weight + 1.0 if zero_centered_gamma else weight and then never reads it — two lines later it builds w_fp32 from the raw weight. So zero_centered_gamma=True produced output bitwise identical to False, while backward did honour the offset. The sibling fused path in the same file does it correctly (w_fp32 = ln_weight.float() then if zero_centered_gamma: w_fp32 = w_fp32 + 1.0), and the offset is applied after the fp32 cast here to match it.

2. The input gradient multiplies the variance term by the weight twice. For y_i = x_i r w_i with r = (mean_j(x_j^2) + eps)^{-1/2} over the last axis of size D, dr/dx_k = -r^3 x_k / D, so

dL/dx_k = g_k w_k r  -  (r^3 x_k / D) * sum_i g_i x_i w_i

The weight appears once per term: explicitly in the first, and inside s = sum_i g_i x_i w_i in the second. backward computed - w_k r^3 x_k s / D. The weight gradient g_w is correct and unchanged.

Because the extra factor is elementwise w_k, the error is identically zero while every weight is exactly 1 — i.e. at initialisation and never again after the first optimizer step. That is also why test_te_rmsnorm_parity passes: it does compare grad_diff = (x.grad - x_clone.grad).abs().max(), but it never perturbs layer_bik.weight from its ones_ init, so the one test that could catch this is blind at exactly the weight values it uses.

Issue tracking

For PRs from open-source community contributors:

• New features: a linked issue is required.
• Small updates (bug fixes, minor improvements): a linked issue is recommended.

Linked issue: none — 3 changed lines in one file. Happy to open one if you would prefer to track it.

Heads-up on a likely conflict: #4740 rewrites g_w inside this exact backward and carries the dx line through as an unmodified context line, so it will conflict with commit 2. #6634 also touches this file but not this function. Say the word and I will rebase behind either.

Contribution process

Pre-checks

• I have added relevant unit tests — see the note above: the natural place is test_te_rmsnorm_parity, which needs a non-unit weight to be sensitive to commit 2 at all. Happy to add that if you want it in this PR rather than a follow-up.
• I have added relevant functional tests
• I have added proper typing to my code — the touched lines carry no annotations.
• I have added relevant documentation
• I have run the autoformatter.sh on my PR — black==26.3.0 --skip-magic-trailing-comma --skip-string-normalization, isort==5.13.2 and ruff check all clean on the changed file.

How to reproduce, and how we checked it

Where this is live, stated first because it is the main limitation. batch_invariant_backend defaults to "te_native", and that path calls _te_patch_for_batch_invariant(skip_gemm=True, skip_rmsnorm=True), which returns at if skip_rmsnorm: return before rms_cls.forward = _te_rmsnorm_forward_patched. So under the default backend BatchInvariantRMSNormFn is never installed and neither defect can fire. They are live under "deepgemm" or "triton", both validated first-class values, with set_batch_invariant_mode defaulting to backend or "triton"; enable_batch_invariant_mode itself defaulted to "deepgemm" until #6521 merged on 2026-08-13. Once the class is patched, _te_rmsnorm_forward_patched carries no is_batch_invariant_mode_enabled() guard, so every RMSNorm backward in the step goes through it.

Checked on 8xL4 / torch 2.12.0+cu130, plus float64 on CPU. Three independent harnesses were used; none differentiates through mean_dim, because it is a Triton kernel and autograd is not reliable through it — building a reference that way produced a misleading number once and it was discarded.

Commit 1, forward:

dtype=bfloat16  forward(zero_centered_gamma=True) == forward(False): True   (before)
dtype=float32   forward(zero_centered_gamma=True) == forward(False): True   (before)
max|forward(True) - (x * rsqrt(mean(x^2)+eps)) * (weight+1)| = 2.29688 (bf16) / 2.38706 (fp32)  before
                                                            = 0                                after
zero_centered_gamma=False path: bitwise unchanged by the patch

Commit 2, input gradient — analytic backward against a finite difference of this function's own forward, zero_centered_gamma=False:

                        weight=randn        weight=ones (control)
before   fp32           0.9878              2.122e-05
after    fp32           2.599e-05           2.122e-05
before   fp64 (CPU)     1.506833 / 6.0326   4.44e-16 / 5.74e-10
after    fp64 (CPU)     1.776357e-15 / 1.989215e-09

The two fp64 columns are two separately written harnesses. The weight=ones control is what makes this a discriminator rather than a measurement: the extra factor is elementwise w, so the error must vanish at unit weight — and it does, unchanged, while the non-unit case drops to machine epsilon. weight=zeros gives exactly 0.0. The forward agrees with an independent reference to 0.0, which confines the second defect to the backward.

What we did not do. We did not run a training job end to end, and we did not benchmark anything — no performance claim is made anywhere in this PR. transformer_engine is not installed on the machine we tested on, so BatchInvariantRMSNormFn was exercised directly rather than through a patched te.pytorch.RMSNorm.

The strongest objection to this change

The default backend does not install this function, so neither defect can fire in a default run — the audience is whoever explicitly selects deepgemm or triton, plus anyone who adopted batch-invariance before #6521 changed the default. And the second defect is identically zero at initialisation, which is why the existing parity tests and the GRPO train/generation parity runs all pass. A reviewer could reasonably conclude the blast radius is small. Our answer is only that the audience for batch-invariance is people who need bitwise reproducibility, who are the least likely to notice a systematically wrong gradient and the most likely to be misled by it — but that is a judgement about priority, not about correctness, and it is yours to make.

[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.

[huthvincent]

@NVIDIA/mcore-oncall — this has been Ready since yesterday and copy-pr-bot is still waiting on validation, so NVIDIA's runners have not built it yet. Could someone run /ok to test 459ebed when it is convenient?

It is two commits in one file (+3 −2) fixing the batch-invariant RMSNorm forward and its input gradient; the reasoning, the reachability limits and a standalone repro are in the PR body. Happy to rebase onto a newer main first if you would prefer that.