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

`DotProductAttention` with `fp8_dpa=True` does not update its backward scales when no Linear-family module is in the autocast region, and dQ/dK lose most of their norm

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

@sha7doww đang làm issue này rồi.

Từ ngày 29/10/2026.

  • #3780 của @sha7doww — đ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
48/100
Loại issue
Lỗi
Độ rõ ràng
Đặc tả rõ ràng
Mức độ hoạt động
Sôi nổi
Công nghệ
python, pytorch
Lĩnh vực
backend, machine-learning

Hướng nghiên cứu

Read FP8GlobalStateManager.reduce_and_update_fp8_tensors and the first-FP8-module handling in the Linear-family backward paths, then compare them with attention/dot_product_attention/backends.py and dot_product_attention.py. Run the provided reproduction on a supported GPU, first with dpa_only and then with dpa_update, to confirm the backward-scale behavior. Done means the issue’s expected behavior is met and the reproduction shows dQ and dK norms recovering after scale updates, or the documentation and warning clearly explain the Linear-family requirement.

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

Mô tả

Describe the bug

With DelayedScaling(fp8_dpa=True), if DotProductAttention is the only Transformer Engine module inside te.autocast (for example in a model whose projections are plain torch.nn.Linear), the backward scales in dpa.fp8_meta["scaling_bwd"] stay at 1.0 and the backward tensors are quantized with scale 1.0. With a cotangent of randn * 1e-5, close to the dO magnitudes we see in training, dQ and dK lose about 70–100 % of their norm in the script below, with no error or warning.

In the module files we read on main (dfc4158), the backward amax reduction and scale update, FP8GlobalStateManager.reduce_and_update_fp8_tensors(forward=False), is called from the backward of Linear, LayerNormLinear, LayerNormMLP and GroupedLinear, and only in the module that claimed the "first FP8 module" flag during its forward (through check_fp8_reduce_and_update, or is_first_fp8_module() directly in GroupedLinear). DotProductAttention neither claims the flag nor makes the call (attention/dot_product_attention/backends.py, dot_product_attention.py), so when no Linear-family module runs in the region, nothing updates the backward scales.

Results of the script below (one H20, 6 iterations; TE 2.20.2 and TE 2.13.0 print identical numbers; the FP8 setups ran on FusedAttention sub-backend 2, the bf16 reference on sub-backend 1):

setup max backward scale, after each backward (iterations 0–5) ‖dQ‖ ‖dK‖ ‖dV‖
bf16 reference (no FP8) – 1.30–1.35e-3 1.30–1.35e-3 1.52–1.55e-3
FP8, DotProductAttention only 1.0 every time 0 in iterations 0 and 4, else 2.6–3.9e-4 0 in iterations 0 and 4, else 2.5–4.0e-4 1.32–1.35e-3
same, plus reduce_and_update_fp8_tensors(forward=False) after each backward 5.1e9–8.2e9 1.31–1.33e-3 1.32–1.34e-3 1.52–1.55e-3
a te.Linear producing Q in the same region 6.9e9–8.5e9 1.06–1.10e-3 0.68–0.70e-3 1.24–1.25e-3

The gradient norms of the last two rows are for iterations 1–5: iteration 0 has dQ = dK = 0 in every FP8 setup, as expected while the first backward uses the initial scales of 1.0. The last row is a control for the update trigger only: its Q comes out of a randomly initialized Linear, so its norms are not comparable with the reference.

We first noticed this on TE 2.10 in a small GPT-2-style training setup whose projections are plain torch.nn.Linear and whose attention is DotProductAttention: over 20 AdamW steps, the tracked layer's Q and K projection weights stayed exactly unchanged without the manual update (relative change 0), and changed by 0.22 relative with it.

Steps/Code to reproduce bug

Run the script at the end of this issue on one GPU. It runs four setups in fresh processes and prints, after each backward, the backward scales and the gradient norms of Q, K and V. The relevant pieces:

recipe = DelayedScaling(fp8_format=Format.HYBRID, amax_history_len=16, amax_compute_algo="max", fp8_dpa=True,
                        fp8_mha=False)
dpa = te.DotProductAttention(num_attention_heads=H, kv_channels=D, attn_mask_type="causal", qkv_format="bshd")

and, for the third row, after each backward:

FP8GlobalStateManager.reduce_and_update_fp8_tensors(forward=False)

Expected behavior

Either DotProductAttention takes part in the mechanism the Linear-family modules use (claims the first-FP8-module flag in its forward and runs the backward update when it holds it), or the fp8_dpa documentation states that a Linear-family module must run in the same autocast region, ideally with a warning when none does.

Environment overview

  • Environment location: a container on a GPU cluster
  • Method of Transformer Engine install: 2.20.2 from PyPI (transformer_engine and transformer_engine_cu13 wheels, transformer_engine_torch built from its sdist against the PyTorch below); 2.13.0 built from source at the v2.13 tag (2.13.0+28777046)

Environment details

  • OS version: Linux x86_64
  • PyTorch version: 2.13.0+cu130
  • Python version: 3.10.10
  • Transformer Engine version: 2.20.2 (also 2.13.0)
  • CUDA version: 13.0
  • CUDNN version: 9.20.0 (nvidia-cudnn-cu13 9.20.0.48)

Device details

  • GPU model: NVIDIA H20 (compute capability 9.0), driver 575.57.08

Additional context

  • Workaround: call FP8GlobalStateManager.reduce_and_update_fp8_tensors(forward=False) after each backward, as in the script's third setup.
  • Possibly related: #1190 / #1254, where the backward amax reduction was skipped under activation recompute because the first-module flag was lost. Here no module claims the flag at all.
  • Not tested: main at runtime (main was only read), fp8_mha=True, and recipes other than delayed scaling.
Full script 
#!/usr/bin/env python3
"""DotProductAttention(fp8_dpa=True) as the only Transformer Engine module in the autocast region: its backward
delayed-scaling scales stay at 1.0, and the backward loses most of dQ and dK.

Setups, each in a fresh process: N iterations of forward + backward on new random q, k, v (bf16, bshd [2, 1024, 6, 128],
causal) with a cotangent of randn * 1e-5.
  ref         the same DPA outside any FP8 autocast (bf16): the gradient norms to expect
  dpa_only    the DPA alone inside te.autocast
  dpa_update  as dpa_only, plus FP8GlobalStateManager.reduce_and_update_fp8_tensors(forward=False) after each backward
  linear_dpa  a te.Linear computes q inside the same autocast region (it claims the first-FP8-module flag)
Prints, per iteration and after each backward, the DPA's backward scales (min, max) and the gradient norms of q, k, v.
usage: python3 te_dpa_bwdscale_repro.py [--iters 6]
"""
import argparse
import importlib.metadata as md
import os
import subprocess
import sys

ap = argparse.ArgumentParser()
ap.add_argument("--iters", type=int, default=6)
ap.add_argument("--arm", default=None)
a = ap.parse_args()
ARMS = ("ref", "dpa_only", "dpa_update", "linear_dpa")
if a.arm is None:
    rc = 0
    for arm in ARMS:
        print(f"\n===== {arm} =====", flush=True)
        env = dict(os.environ, NVTE_FP8_DPA_BWD="1", NVTE_DEBUG="1", NVTE_DEBUG_LEVEL="1")
        rc = rc or subprocess.run([sys.executable, os.path.abspath(__file__), "--iters", str(a.iters), "--arm", arm],
                                  env=env).returncode
    sys.exit(rc)

import torch
import transformer_engine as T
import transformer_engine.pytorch as te
from transformer_engine.common.recipe import DelayedScaling, Format
try:
    from transformer_engine.pytorch.quantization import FP8GlobalStateManager
except ImportError:  # older TE
    from transformer_engine.pytorch.fp8 import FP8GlobalStateManager


def fp8_region(recipe):
    if hasattr(te, "autocast"):
        return te.autocast(enabled=True, recipe=recipe)
    return te.fp8_autocast(enabled=True, fp8_recipe=recipe)


def cudnn_pkg():
    for p in ("nvidia-cudnn-cu13", "nvidia-cudnn-cu12"):
        try:
            return f"{p} {md.version(p)}"
        except md.PackageNotFoundError:
            pass
    return "cudnn package not found"


B, S, H, D = 2, 1024, 6, 128
torch.manual_seed(0)
print(f"te {T.__version__} | torch {torch.__version__} | {cudnn_pkg()} | {torch.cuda.get_device_name()} | arm {a.arm}",
      flush=True)
recipe = DelayedScaling(fp8_format=Format.HYBRID, amax_history_len=16, amax_compute_algo="max", fp8_dpa=True,
                        fp8_mha=False)
dpa = te.DotProductAttention(num_attention_heads=H, kv_channels=D, attn_mask_type="causal", qkv_format="bshd")
lin = te.Linear(H * D, H * D, bias=False, params_dtype=torch.bfloat16) if a.arm == "linear_dpa" else None


def bwd_scales():
    st = getattr(dpa, "fp8_meta", {}).get("scaling_bwd")
    if st is None:
        return "none"
    return f"[{float(st.scale.min()):.3g}, {float(st.scale.max()):.3g}]"


def norm(t):
    return float(t.float().norm()) if t is not None else float("nan")


for it in range(a.iters):
    rnd = dict(device="cuda", dtype=torch.bfloat16, requires_grad=True)
    x, q, k, v = (torch.randn(B, S, H * D, **rnd), torch.randn(B, S, H, D, **rnd), torch.randn(B, S, H, D, **rnd),
                  torch.randn(B, S, H, D, **rnd))
    g = torch.randn(B, S, H * D, device="cuda", dtype=torch.bfloat16) * 1e-5
    if a.arm == "ref":
        o = dpa(q, k, v)
    else:
        with fp8_region(recipe):
            if lin is not None:
                q = lin(x).view(B, S, H, D)
                q.retain_grad()
            o = dpa(q, k, v)
    o.backward(g)
    if a.arm == "dpa_update":
        FP8GlobalStateManager.reduce_and_update_fp8_tensors(forward=False)
    print(f"it {it}: bwd scales {bwd_scales()} |dq| {norm(q.grad):.3e} |dk| {norm(k.grad):.3e} |dv| {norm(v.grad):.3e}",
          flush=True)
print(f"RESULT {a.arm}: bwd scales {bwd_scales()} last |dq| {norm(q.grad):.3e} |dk| {norm(k.grad):.3e} "
      f"|dv| {norm(v.grad):.3e}", flush=True)
Ngôn ngữ chính
Python
Star
3.6k
Fork
883
Merge trung bình
4 ngày 20 giờ
Pull request đã merge (30 ngày)
71

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 NVIDIA/TransformerEngine

Tất cả issue của NVIDIA/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.