float8: work around Triton fp8 store miscompile in compiled training

[vkuzo]

Summary:
Compiled float8 training (e.g. FSDP2 + torch.compile) could intermittently
produce NaN/Inf losses. Root cause is a Triton miscompile: when a kernel
contains min.NaN/max.NaN PTX instructions, Triton mis-lowers a neighboring
transposed, vectorized 1-byte (fp8) store, writing garbage bytes that decode as
NaN/Inf. The computed value is correct (the contiguous store of the same value
is clean); only the transposed store is corrupted. Filed upstream as
triton-lang/triton#11111.

The bug became visible after PyTorch inductor PR
pytorch/pytorch#186933, which changed the Triton codegen
for min/max from the tl.where-based triton_helpers.{minimum,maximum} to
tl.{minimum,maximum}(a, b, tl.PropagateNan.ALL). The two forms are numerically
identical, but only the PropagateNan.ALL form lowers to the min.NaN/max.NaN
instructions that trip the store bug.

Timeline: nightlies were GOOD through 2026-06-13 and BAD starting 2026-06-14,
bracketing #186933. The failure is flaky (~25%/run) because inductor autotunes
the fused kernel's launch config and only a subset of configs trip the store
miscompile; the miscompile itself is deterministic per config.

This change reverts the inductor min/max codegen back to the numerically-identical
triton_helpers (tl.where) form, which avoids emitting min.NaN/max.NaN.
It is applied from convert_to_float8_training so it only affects users of the
float8 training product, not every import torchao. The reverted form is always
correct, so this is a numeric no-op even after Triton fixes the underlying bug.

Test Plan:
Repro'd the FSDP2 float8 + torch.compile NaN in a standalone script; verified the
config that deterministically NaN'd now produces finite losses with the workaround
applied, and confirmed convert_to_float8_training (but not bare import torchao)
flips the codegen.

Co-Authored-By: Claude Opus 4.8 noreply@anthropic.com

[vkuzo]

Stack from ghstack (oldest at bottom):

• -> float8: work around Triton fp8 store miscompile in compiled training #4652

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/4652

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure, 1 Pending

As of commit c96f599 with merge base 22705bd ():

NEW FAILURE - The following job has failed:

• Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch torchvision --index-url htt... / linux-job (gh)
  RuntimeError: Command docker exec -t 78de069b0ce3081776d8859ac27ecc1a560ce3d07ba531d821fdbfe8c803aabc /exec failed with exit code 134

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[andrewor14]

@claude review

[claude]

Claude finished @andrewor14's task in 2m 2s —— View job

---

Code Review

• Read changed files and understand the workaround
• Analyze correctness & side effects of the monkeypatch
• Check scope, idempotency, and thread-safety concerns
• Post review feedback

Overall this is a well-reasoned, well-documented workaround for a nasty, flaky, silent-NaN bug. The root-cause writeup in the module docstring is excellent and the defensive try/except around the import is the right instinct. I have a few substantive concerns worth discussing before merge, plus some minor notes.

---

Substantive

1. The blast radius is process-global, not float8-scoped — and the comments imply otherwise.
The comment in float8_linear_utils.py:112-116 and the docstring both say the patch "only affects users who actually use the float8 training product." That's true at the granularity of "did you call convert_to_float8_training", but once called, TritonOverrides.minimum/maximum is mutated process-globally. Every subsequent torch.compile in that process — including non-float8 models — now uses the reverted tl.where codegen. Since the two forms are numerically identical this is a numeric no-op everywhere, but it's a broader change than the phrasing suggests. Suggest tightening the wording to be explicit that the codegen change is global once triggered (just gated behind the float8 entry point).

2. Silent no-op on internal drift reintroduces a silent-NaN bug. _inductor_patch.py:47 logs at logger.debug and returns if TritonOverrides can't be imported (or, see #3, if the attributes have moved). Given the failure mode this guards against is silent NaN/Inf losses in training, a silent debug-level no-op is risky: a future torch refactor that moves these internals would silently drop the workaround and the flaky NaN returns with zero signal. Consider logger.warning on the failure path so the regression is at least discoverable in logs. Fix this →

3. The try/except only guards the import, not attribute existence/signature. If TritonOverrides still imports but minimum/maximum have been renamed, moved to a different overrides class, or changed signature (e.g. gained a dtype arg), the current code will unconditionally overwrite them with a 2-arg staticmethod and set _PATCHED = True — masking the drift and potentially producing broken codegen. Consider asserting the attributes exist and are callable before replacing (e.g. getattr(TritonOverrides, "minimum", None) check), and only then patching. This makes the patch fail-safe against the more likely form of drift.

4. Interaction with inductor's compile cache. The monkeypatch changes codegen output but is invisible to inductor's FXGraph/on-disk cache key. A kernel compiled + cached before the patch was applied (e.g. warm on-disk cache from a prior BAD run, or a compile that happened earlier in the same process before convert_to_float8_training) could still be served from cache with the buggy codegen. In the normal flow the patch runs before any float8 kernel compiles, so this is likely fine — but it's worth a sentence in the docstring noting the workaround assumes no stale cached artifact, and that users hitting the bug should clear their inductor cache. Was this exercised with a warm TORCHINDUCTOR_CACHE_DIR?

Minor

• Thread-safety of _PATCHED: the check-then-set isn't atomic, but convert_to_float8_training is realistically not called concurrently, and a double-apply is harmless (idempotent overwrite). Not worth locking — noting for completeness.
• Reduction path: the patch only covers pointwise minimum/maximum. Amax-style max reductions go through a different helper (triton_helpers.*) that isn't affected by #186933, so this is the correct/sufficient target — just confirming the scope is intentional.
• CI failure: the one new failure (test-nightly, exit code 134 / likely SIGABRT) is worth a glance to confirm it's unrelated infra flakiness vs. anything the patch could touch.
• Verification depth: the test plan is a manual standalone repro (reasonable given the flaky, autotune-dependent nature). Consider whether a lightweight unit test asserting convert_to_float8_training flips TritonOverrides.minimum/maximum to the triton_helpers form (and bare import torchao does not) is worth adding as a regression guard against Headers and gitignore #2/Move test.py #3 above.

---

Note: I reviewed statically — the sandbox blocked running Python, so I couldn't confirm TritonOverrides.minimum/maximum signatures against the installed torch. That's exactly why concern #3 matters.