Hacktoberfest 2026: những issue maintainer đã đánh dấu cho tháng Mười, đang mở và phù hợp người mới. Xem issue Hacktoberfest

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

Đang mở
  • Có gì không ổn?
#693 0 bình luận 0 reaction 0 người được giao Xem trên GitHub

Maintainer thường phản hồi trong vòng 2 ngày

@zjin-lcf đang làm issue này rồi.

Từ ngày 19/8/2026.

  • #708 của @zjin-lcf — đang mở

Đánh giá

Độ khó
4/5
Thời gian dự kiến
3-5 ngày
Mức phù hợp với người mới
45/100
Loại issue
Lỗi
Độ rõ ràng
Khá rõ ràng
Mức độ hoạt động
Ít trao đổi
Công nghệ
python
Lĩnh vực
backend, performance

Hướng nghiên cứu

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.

Do mô hình lập chỉ mục viết ra từ nội dung của issue.

Mô tả

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.

Ngôn ngữ chính
Python
Star
77
Fork
36
Merge trung bình
3 ngày 16 giờ
Pull request đã merge (30 ngày)
10

Chuẩn bị môi trường

Bắt đầu từ đâu

  1. Đọc hết issue, rồi đọc hướng dẫn đóng góp của dự án.
  2. Bình luận trên issue rằng bạn sẽ nhận — tránh hai người làm cùng một việc.
  3. Fork repository và làm thay đổi trên một nhánh.
  4. Mở pull request có tham chiếu số hiệu của issue.

Issue khác của ROCm/TransformerEngine

Tất cả issue của ROCm/TransformerEngine

Issue tương tự

Thêm issue về Python

Nhận issue mới trong hộp thư của bạn

Bản tóm tắt ngắn những issue GitHub phù hợp với người mới.