Hacktoberfest 2026：メンテナが10月に向けて印を付けた、オープンで初心者向けの issue。 Hacktoberfest の issue を見る

`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

オープン
  • 何か問題がありますか？
#3,778 コメント 0 件 リアクション 0 件 担当者 0 名 GitHub で見る

メンテナーはふだん 2 日以内に返信

@sha7doww がすでに取り組んでいます。

2026年10月29日 から。

  • #3780 @sha7doww による — オープン

評価

難易度
4/5
見積もり時間
3〜5日
初心者へのやさしさ
48/100
issue の種類
バグ
明瞭さ
明確に書かれている
活発さ
活発
技術スタック
python, pytorch
領域
backend, machine-learning

調査の方向性

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.

索引モデルが issue の本文から書いたものです。

説明

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)
主要言語
Python
スター
3.6k
フォーク
883
平均マージ
4日 20時間
マージ済み PR（30日）
71

環境構築

はじめの一歩

  1. issue を最後まで読み、次にプロジェクトのコントリビューションガイドを読みます。
  2. 着手することを issue にコメントします — 二人が同じ作業をするのを防げます。
  3. リポジトリをフォークし、ブランチを切って変更します。
  4. issue 番号を参照したプルリクエストを送ります。

NVIDIA/TransformerEngine のほかの issue

NVIDIA/TransformerEngine の issue をすべて見る

似ている issue

Python の issue をもっと見る

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。