Investigate torch.compile numerical differences between grad and no-grad RMSNorm forward passes

Description

NeMo RL DPO training expects the initial preference loss to equal ln(2) (~0.693147) because the policy and reference model are identical before the first optimizer step.

In the NanoV3 30B DPO job, the initial loss is instead approximately 0.692086. This indicates that the policy and reference forward passes produce different outputs despite using identical weights and inputs.

The behavior has been narrowed to the compiled FP32 RMSNorm implementation:

https://github.com/NVIDIA-NeMo/Automodel/blob/ec2f7240af92929e383a863189a4d395b0025cb6/nemo_automodel/components/models/common/utils.py#L237-L243

The reference forward runs under torch.no_grad(), while the training forward runs with gradients enabled. With:

@torch.compile(fullgraph=True, dynamic=True)

PyTorch/Inductor appears to generate reductions with different floating-point accumulation orders for the two execution contexts.

Removing the torch.compile decorator resolves the discrepancy. Treating RMSNorm as an opaque operation, as proposed in NVIDIA-NeMo/Automodel#2971, also resolves the DPO failure.

Reproduction

1. Run the DPO reproducer associated with NVIDIA-NeMo/RL#2662 for one training step.
2. Record the initial train/preference_loss.
3. With the current compiled FP32 RMSNorm implementation, observe that the loss differs from:

math.log(2)  # 0.693147...

4. Repeat either:
   • after removing the torch.compile decorator from _float32_rms_norm_fwd, or
   • using the changes from NVIDIA-NeMo/Automodel#2971.
5. Confirm that the initial loss returns to ~0.693147.

Related issue: https://github.com/NVIDIA-NeMo/RL/issues/2662
Candidate fix: https://github.com/NVIDIA-NeMo/Automodel/pull/2971

Investigation goals

• Produce a minimal AutoModel/PyTorch reproducer that does not depend on NeMo RL.
• Compare identical RMSNorm inputs and weights under:
  • eager grad
  • eager no-grad
  • compiled grad
  • compiled no-grad
• Identify whether the difference comes from:
  • separate Dynamo/Inductor graphs,
  • Triton reduction decomposition,
  • kernel autotuning,
  • PyTorch or Triton version changes, or
  • another execution-context-dependent optimization.
• Determine the affected PyTorch, CUDA, Triton, GPU, dtype, shape, and AutoModel version ranges.
• Verify that the proposed fix preserves forward numerics, backward correctness, CPU support, and expected performance.

Acceptance criteria

• A minimal regression test demonstrates the issue without running a full DPO job.
• Identical inputs and weights produce matching outputs across grad and no-grad contexts within an explicitly defined tolerance.
• The regression test fails with the affected implementation and passes with the fix.
• Forward and gradient parity are covered on CPU and CUDA.
• The supported software-version matrix is documented.
• Any performance impact of the fix is measured.
• NeMo RL DPO starts at ln(2) with the corrected AutoModel version.
• A CI compatibility test catches future compiler or dependency changes before they reach NeMo RL nightlies.

Broader follow-up

This should become a small numerical-compatibility gate at the AutoModel boundary rather than relying on downstream RL workloads to detect compiler- or dependency-induced numerical changes.

Start with _float32_rms_norm_fwd in nemo_automodel/components/models/common/utils.py and reproduce the discrepancy using the linked NeMo RL DPO reproducer. Compare eager and compiled RMSNorm outputs and gradients under grad and no-grad on CPU and CUDA, then narrow the affected version, dtype, shape, and GPU ranges. Done means a minimal regression test fails before the fix, passes afterward, and the documented compatibility and performance checks are covered.