Hacktoberfest 2026: le issue che i maintainer hanno segnato per ottobre, aperte e adatte ai principianti. Sfoglia le issue Hacktoberfest

CP LSE merge switches compiled variants on ROCm, breaking forward reproducibility

Aperta
  • Qualcosa non va?
#693 0 commenti 0 reazioni 0 assegnatari Vedi su GitHub

I maintainer di solito rispondono entro 2 giorni

@zjin-lcf ci sta già lavorando.

Dal 19/8/2026.

  • #708 di @zjin-lcf — aperta

Valutazione

Difficoltà
4/5
Tempo stimato
3-5 giorni
Idoneità per principianti
45/100
Tipo di issue
Bug
Chiarezza
Abbastanza chiara
Stato di attività
Tranquilla
Stack tecnologico
python
Ambito
backend, performance

Direzione di ricerca

Start by locating flash_attn_fwd_softmax_lse_correction and flash_attn_fwd_second_half_softmax_lse_correction and inspecting their @jit_fuser usage. Reproduce the CP=2 forward comparison on ROCm, then evaluate the proposed eager and dynamic compilation options. Done means reference and actor forwards remain bit-identical without switching compiled variants.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Descrizione

Summary

The two CP LSE merge functions use @jit_fuser, which calls torch.compile:

  • flash_attn_fwd_softmax_lse_correction
  • flash_attn_fwd_second_half_softmax_lse_correction

Dynamo may keep multiple compiled variants of the same function in one process. These variants can select different Triton launch configurations. On gfx950, libdevice.log1p is not bit-exact across those configurations. As a result, the reference and actor forwards can disagree even when their inputs and weights are bit-identical.

Symptom

At RL step 0, the reference model and actor are initialized from the same checkpoint, and no optimizer update has happened yet. When evaluated on the same tokens with the same inputs, their forward passes should therefore produce bit-identical logits and log-probabilities.

On 4× MI355X with CP=2, they differ by roughly 1e-5 to 2e-4. The same test is bit-exact on H200 and MI355X with CP=1. With CP=1, the LSE merge is never called.

Root cause

The first mismatch appears in:

layers.0.self_attention.core_attention.flash_attention

The attention inputs and both LSE merge operands are bit-identical. Only the merge result differs. The reference and actor forwards use different compiled variants of the same merge function:

  • one takes a scalar log1p path;
  • the other takes a vectorized log1p path.

On identical inputs, the two paths differ by about 1 ULP on gfx950.

Testing 48 combinations of XBLOCK and num_warps produces exactly two result groups. The split occurs when LLVM vectorizes the kernel body. Among 18 tested elementwise operations, only these depend on the launch configuration:

log1p    14.0% of elements differ
cosh      3.0% of elements differ

Operations such as exp, log, sqrt, tanh, erf, pow, and fma are invariant. Assembly confirms that vectorization changes the multiply-add contraction pattern inside the OCML implementation.

In the real training run:

  • the reference forward exactly matches the scalar-path result;
  • the actor forward exactly matches the vectorized-path result.

This roughly 1-ULP difference is amplified across the model's attention layers into the observed 1e-5 to 2e-4 log-probability mismatch.

Why this appears on ROCm

Both ROCm results are within OCML's documented 2 ULP accuracy bound. The problem is that torch.compile allows the process to switch between two valid but non-bit-identical implementations.

Minimal reproducer

import torch
import triton
import triton.language as tl
from triton.language.extra import libdevice

N = 225280


@triton.jit
def kern(in_ptr, out_ptr, xnumel, XBLOCK: tl.constexpr):
    i = tl.program_id(0) * XBLOCK + tl.arange(0, XBLOCK)
    mask = i < xnumel
    x = tl.load(in_ptr + i, mask)
    tl.store(out_ptr + i, libdevice.log1p(x), mask)


def run(x, XBLOCK, num_warps):
    out = torch.empty_like(x)
    kern[(triton.cdiv(N, XBLOCK),)](x, out, N, XBLOCK=XBLOCK, num_warps=num_warps)
    torch.cuda.synchronize()
    return out


torch.manual_seed(0)
x = torch.rand(N, device="cuda", dtype=torch.float32)

a = run(x, XBLOCK=256, num_warps=4)      # 256 / (4 * 64) = 1 element per thread
b = run(x, XBLOCK=1024, num_warps=4)     # 1024 / (4 * 64) = 4 elements per thread

print(f"{int((a != b).sum())} / {N} elements differ, max abs {(a - b).abs().max():.3e}")

ref = torch.log1p(x.double())
ulp = torch.finfo(torch.float32).eps * ref.abs()
for name, o in (("1 element per thread ", a), ("4 elements per thread", b)):
    err = (o.double() - ref).abs() / ulp
    print(f"  {name}: max {err.max():.3f} ULP, mean {err.mean():.4f} ULP")

Output on MI355X (gfx950), ROCm 7.2.0, torch 2.9.1+rocm7.2.0, Triton 3.6.0:

37933 / 225280 elements differ, max abs 5.960e-08
  1 element per thread : max 0.966 ULP, mean 0.2211 ULP
  4 elements per thread: max 1.394 ULP, mean 0.2594 ULP

The same split occurs for CP=2/4/8/16 and THD shapes.

Proposed fix

Both options prevent these functions from switching between static and dynamic compiled variants.

Option A: remove @jit_fuser

Run these two functions eagerly.

Option B: compile these two functions with dynamic=True

Compile these two functions as shape-dynamic from the start, avoiding the initial static specialization.

Lingua principale
Python
Stelle
77
Fork
36
Merge medio
3g 16h
PR unite (30g)
10

Preparare l'ambiente

Come iniziare

  1. Leggi tutta la issue e poi la guida ai contributi del progetto.
  2. Commenta sulla issue per dire che te ne occupi tu — evita che due persone facciano lo stesso lavoro.
  3. Fai un fork del repository e lavora su un branch.
  4. Apri una pull request che faccia riferimento al numero della issue.

Altre issue di ROCm/TransformerEngine

Tutte le issue di ROCm/TransformerEngine

Issue simili

Altre issue su Python

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.