Hacktoberfest 2026: le issue che i maintainer hanno segnato per ottobre, aperte e adatte ai principianti. Sfoglia le 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

Aperta
  • Qualcosa non va?
#3,778 0 commenti 0 reazioni 0 assegnatari Vedi su GitHub

I maintainer di solito rispondono entro 2 giorni

@sha7doww ci sta già lavorando.

Dal 29/10/2026.

  • #3780 di @sha7doww — aperta

Valutazione

Difficoltà
4/5
Tempo stimato
3-5 giorni
Idoneità per principianti
48/100
Tipo di issue
Bug
Chiarezza
Specificata chiaramente
Stato di attività
Attiva
Stack tecnologico
python, pytorch
Ambito
backend, machine-learning

Direzione di ricerca

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.

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

Descrizione

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)
Lingua principale
Python
Stelle
3.6k
Fork
883
Merge medio
4g 20h
PR unite (30g)
71

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

Tutte le issue di NVIDIA/TransformerEngine

Issue simili

Altre issue su Python

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.