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

Numerical accuracy issue with Triton MXFP4 kernel with respect to quantization

Đang mở
  • Có gì không ổn?
#665 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

Chưa có ai nhận issue này.

Đá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
55/100
Loại issue
Lỗi
Độ rõ ràng
Đặc tả rõ ràng
Mức độ hoạt động
Ít trao đổi
Công nghệ
python
Lĩnh vực
machine-learning, performance

Hướng nghiên cứu

Start in transformer_engine/pytorch/triton_kernels/cast_transpose.py, reading te_cast_transpose_mxfp4_triton and _cast_transpose_triton_mxfp4 alongside the attached reproduction script. Compare Triton and native outputs with shuffle_columnwise_data enabled and disabled. Done means the Triton columnwise FP4 data changes with the flag and matches the native preshuffled layout, while the existing scale-swizzle behavior remains correct.

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

Mô tả

bug

Summary

The Triton MXFP4 cast-transpose path (te_cast_transpose_mxfp4_triton, selected by
NVTE_USE_CAST_TRANSPOSE_TRITON=1) does not honor the quantizer's shuffle_columnwise_data /
shuffle_rowwise_data flags. It applies the scale swizzle (with_gemm_swizzled_scales) but
stores the packed FP4 data in the plain transposed layout, ignoring the data-preshuffle flags
entirely.

The native/HIP path (tex.quantize → te_mxfp4::cast_transpose_mxfp4_shuffled) does apply the data
preshuffle. So when a recipe sets shuffle_columnwise_data=True (the real MXFP4 training layout),
the Triton kernel's columnwise operand is unshuffled and does not match the preshuffled layout
the AITER a4w4 GEMM consumes (bpreshuffle=True).

Because the columnwise operand feeds the backward weight-gradient (wgrad) GEMM, the FP4 values
are correct but their ordering is wrong. This produces structurally-wrong weight gradients that
compound over many steps into a slow grad-norm blow-up (not an immediate NaN).

Reproduction

Attaching a script which helps in reproduction of bug. triton_mxfp4_ignores_columnwise_shuffle.py

The attached script reproduces the bug on a single GPU (random bf16 tensor, no training). It
quantizes the same tensor with both kernels and compares the columnwise FP4 data with the shuffle
flag on and off: the native output changes when the flag flips (it preshuffles), while the Triton
output stays byte-identical (it ignores the flag).

Expected vs. actual

  • Expected: with shuffle_columnwise_data=True, the Triton kernel's columnwise FP4 data is
    preshuffled identically to the native kernel.
  • Actual: the Triton kernel emits identical columnwise data whether the flag is True or False —
    i.e. it ignores the flag and leaves the data unshuffled.

Environment

  • GPU: AMD Instinct MI355X (gfx950), ROCm 7.x
  • Image (reproduces the bug): rocm/primus:pytorch-2.12.0-rocm7.14.0a20260608_te-2.14.0.dev0-e6ede467_v26.4

Root cause (code pointers)

In transformer_engine/pytorch/triton_kernels/cast_transpose.py, the launcher only ever reads the
scale-swizzle flag; the data-shuffle flags are never consulted:

# te_cast_transpose_mxfp4_triton
shuffle_scales = out_metadata["quantizer"].with_gemm_swizzled_scales
SHUFFLE_ROWWISE_SCALING = shuffle_scales and USE_ROWWISE_SCALING
SHUFFLE_COLWISE_SCALING = shuffle_scales and USE_COLWISE_SCALING
# quantizer.shuffle_columnwise_data / shuffle_rowwise_data are never read.

Inside the kernel (_cast_transpose_triton_mxfp4), the columnwise branch stores FP4 data with a
plain transposed store and then swizzles only the scales:

if USE_COLWISE:
    packed_col, bs_col = _mxfp4_quantize_32x32_block(tl.trans(x_chunk), MXFP4_BLOCK_SIZE)
    tl.store(colwise_fp4_ptr + ..., packed_col, mask=col_fp4_mask)  # plain transpose, no data preshuffle
    if SHUFFLE_COLWISE:
        ...  # shuffles SCALES only

The native C++/HIP path reads the same flags and applies the data preshuffle in-kernel, so the two
kernels diverge only when a data-shuffle flag is set.

Impact

With NVTE_USE_CAST_TRANSPOSE_TRITON=1 and a recipe that sets shuffle_columnwise_data=True, wgrad
reads a misordered FP4 operand. In a Flux-12B MXFP4 run (SR off) this turns a converging run
(loss ~0.63, grad norm ~0.2) into a diverging one (grad norm climbing to ~58 over ~1000 iters),
while the native kernel converges under identical settings.

Additional: the Triton path is also slower

Independently of the correctness bug above, the Triton cast-transpose is structurally slower than the
native kernel, for these reasons:

  1. Two quantizes plus an in-kernel transpose per tile. The columnwise operand is produced by
    physically transposing each tile (tl.trans) and re-running the full quantizer, reusing none of
    the rowwise work. The native kernel emits both the rowwise and columnwise layouts in a single
    fused pass.
  2. Expensive per-element encoding. The E8M0 scale is computed with transcendentals
    (log2/floor/exp2) instead of reading the exponent bits directly, and the E2M1 data uses a
    dependent chain of tl.where comparisons instead of the gfx950 hardware packed-FP4 cast.
  3. Not autotuned. The kernel uses a fixed tile size and default launch parameters with generic
    codegen, so it under-uses the GPU and cannot adapt to different tensor shapes. The native kernel
    is hand-tuned for gfx950.
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.