[albertvillanova]

This PR makes the Triton loss kernels compile correctly on torch < 2.11, where torch_compile=True currently trains with a zero loss and NaN gradients.

Fix #7687.

Motivation

Before torch 2.11, Inductor's analysis of which tensors a Triton kernel mutates passes the kernels' scalar arguments as constexpr or Python float. The .to(tl.float32) casts added in #7600 raise on those, so Inductor assumes every input is mutated and miscompiles the kernels. On torch 2.10.0, SFT with torch_compile=True logs loss=0.0 and grad_norm=nan at every step; with this PR it matches eager.

Changes

• Cast the scalar arguments with tl.cast(..., tl.float32) instead of .to(tl.float32) in chunked_logprob and chunked_divergence

The kernel compile tests pass on torch 2.7.0, 2.10.0 and 2.14.1, and on 2.6.0 apart from the fullgraph test, which is a separate torch 2.6 limitation. #7642 marks these tests xfail below torch 2.11; whichever of the two PRs merges second removes that xfail.

---

Note

Medium Risk
Touches core Triton training loss/divergence math under torch.compile, but the change is a narrow casting fix with tests across multiple torch versions; incorrect compile behavior was the main risk before this fix.

Overview
Fixes broken torch.compile training on PyTorch before 2.11 for SFT and other paths that use the chunked Triton loss kernels.

Scalar kernel args (beta, logit_scale, softcap, inv_t, etc.) are now promoted to fp32 with tl.cast(..., tl.float32) instead of .to(tl.float32) in _transform, JSD/backward paths in chunked_divergence.py, and the logprob backward path in chunked_logprob.py. On older Inductor, compile-time scalars/constexprs cannot use .to, which led to incorrect mutation analysis, miscompiled kernels, zero loss, and NaN gradients; eager behavior is unchanged and compile matches eager again.

^{Reviewed by Cursor Bugbot for commit fad9a21. Bugbot is set up for automated code reviews on this repo. Configure here.}

[bot-ci-comment]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[qgallouedec]

Suggested change

# tl.cast, not .to: before torch 2.11, Inductor's mutation analysis passes them as constexpr/float

let's avoid documenting what we don't do, unless there is strong reason to do so

[qgallouedec]

Thanks! Matches what I saw on #7642: the compile tests gave NaN from 2.6 to 2.9.1 and passed on 2.14.1. Whichever of the two merges second drops the xfail.

[qgallouedec]

Suggested change

# tl.cast, not .to: before torch 2.11, Inductor's mutation analysis passes them as constexpr/float

[qgallouedec]

Suggested change

# tl.cast, not .to: before torch 2.11, Inductor's mutation analysis passes them as constexpr/float

[qgallouedec]

Suggested change

# tl.cast, not .to: before torch 2.11, Inductor's mutation analysis passes them as constexpr/float