@@ -1336,13 +1336,14 @@ def forward(ctx, x: torch.Tensor, weight: torch.Tensor, eps: float, zero_centere
             raise RuntimeError("Batch-invariant RMSNorm requires CUDA tensors.")
         if not _is_supported_dtype_for_bik(x.dtype):
             raise RuntimeError(f"Unsupported dtype for batch-invariant RMSNorm: {x.dtype}")
-        weight_eff = weight + 1.0 if zero_centered_gamma else weight
 
         # We do everything in rmsnorm_batch_invariant manually here so that we can
         # save rsigma in full precision for backward to match the TE behavior.
         x_dtype = x.dtype
         x_fp32 = x.float()
         w_fp32 = weight.to(device=x.device, dtype=torch.float32)
+        if zero_centered_gamma:
+            w_fp32 = w_fp32 + 1.0
         ms = mean_dim(x_fp32 * x_fp32, dim=-1, keepdim=True)
         rsigma = torch.rsqrt(ms + eps)
         out_fp32 = (x_fp32 * rsigma) * w_fp32
@@ -1364,11 +1365,11 @@ def backward(ctx, grad_output: torch.Tensor):
         accumulation and reduction behavior for numerical stability.
         """
         x, weight, rsigma = ctx.saved_tensors
-        w_eff = (weight + 1.0) if ctx.zero_centered_gamma else weight
-
         go_fp32 = grad_output.float()
         x_fp32 = x.float()
-        w_fp32 = w_eff.to(device=x.device, dtype=torch.float32)
+        w_fp32 = weight.to(device=x.device, dtype=torch.float32)
+        if ctx.zero_centered_gamma:
+            w_fp32 = w_fp32 + 1.0
         r = rsigma
         r3 = r * r * r
         D = x.shape[-1]
@@ -1377,7 +1378,7 @@ def backward(ctx, grad_output: torch.Tensor):
         g_w = (go_fp32 * x_fp32 * r).sum(dim=red_dims).to(weight.dtype)
 
         s = (go_fp32 * x_fp32 * w_fp32).sum(dim=-1, keepdim=True)
-        dx = go_fp32 * (w_fp32 * r) - (w_fp32 * r3) * (s * x_fp32) / D
+        dx = go_fp32 * (w_fp32 * r) - r3 * (s * x_fp32) / D
         dx = dx.to(x.dtype)
 
         return dx, g_w, None, None
 
@@ -618,6 +618,57 @@ def test_batch_invariant_kernels_replay():
     )
 
 
+@pytest.mark.parametrize("zero_centered_gamma", [False, True], ids=["gamma", "zero_centered_gamma"])
+def test_batch_invariant_rmsnorm_replays_and_matches_reference(zero_centered_gamma):
+    """Replay the batch-invariant RMSNorm autograd function and check it against fp32 torch.
+
+    The weight is random on purpose. At ``weight == 1`` the input gradient with the weight
+    factor applied twice equals the correct one, and under zero-centred gamma a forward that
+    forgets the ``+ 1`` returns exactly zero at ``weight == 0`` -- so either default init hides
+    one of the two defects this function has had. Only a non-unit, non-zero weight sees both.
+    """
+    from megatron.core.transformer.custom_layers.batch_invariant_kernels import (
+        BatchInvariantRMSNormFn,
+    )
+
+    seeded()
+    rows, hidden, eps = 4096, 2048, 1e-5
+    x = torch.randn(rows, hidden, device="cuda", dtype=torch.bfloat16)
+    weight = torch.randn(hidden, device="cuda", dtype=torch.bfloat16)
+    grad_output = torch.randn(rows, hidden, device="cuda", dtype=torch.bfloat16)
+
+    def run(x, weight):
+        return BatchInvariantRMSNormFn.apply(x, weight, eps, zero_centered_gamma)
+
+    assert_replays_bit_exact(
+        run,
+        (x.clone().requires_grad_(True), weight.clone().requires_grad_(True)),
+        replays=3,
+        grad_outputs={"out": grad_output},
+        what=f"bik rmsnorm[zero_centered_gamma={zero_centered_gamma}]",
+    )
+
+    # fp32 reference, the same arithmetic TE's RMSNorm defines. Its gradients come from
+    # autograd, so the two defects (a dropped ``+ 1`` in forward, a duplicated weight factor in
+    # the input gradient) each show up as a mismatch here and nowhere in the replay above.
+    x_bik = x.clone().requires_grad_(True)
+    w_bik = weight.clone().requires_grad_(True)
+    out_bik = BatchInvariantRMSNormFn.apply(x_bik, w_bik, eps, zero_centered_gamma)
+    out_bik.backward(grad_output)
+
+    x_ref = x.float().requires_grad_(True)
+    w_ref = weight.float().requires_grad_(True)
+    w_eff = w_ref + 1.0 if zero_centered_gamma else w_ref
+    out_ref = x_ref * torch.rsqrt((x_ref * x_ref).mean(dim=-1, keepdim=True) + eps) * w_eff
+    out_ref.backward(grad_output.float())
+
+    # The function computes in fp32 and rounds once to the input dtype, so compare after the
+    # same rounding of the reference; the only remaining difference is reduction order.
+    torch.testing.assert_close(out_bik, out_ref.to(out_bik.dtype))
+    torch.testing.assert_close(x_bik.grad, x_ref.grad.to(x_bik.dtype))
+    torch.testing.assert_close(w_bik.grad, w_ref.grad.to(w_bik.dtype))
+
+
 def test_mask_routing_padding_replays():
     from megatron.core.transformer.moe.inference_routing_mask_kernel import mask_routing_padding