Hacktoberfest 2026: le issue che i maintainer hanno segnato per ottobre, aperte e adatte ai principianti. Sfoglia le issue Hacktoberfest

[PyTorch][Attention] THD P2P context-parallel regression when padded cu_seqlens are value-equal but not object-identical

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

I maintainer di solito rispondono entro 2 giorni

Nessuno ha ancora preso questa issue.

  • #3506 di @RPalmr — chiusa senza merge

Valutazione

Difficoltà
4/5
Tempo stimato
3-5 giorni
Idoneità per principianti
38/100
Tipo di issue
Bug
Chiarezza
Abbastanza chiara
Stato di attività
Attiva
Stack tecnologico
python, pytorch
Ambito
distributed-systems, machine-learning, performance

Direzione di ricerca

Inizia dal rilevatore di padding di PyTorch DotProductAttention introdotto nel commit 4cd705b75394563c0246bdddfa5d3148106c9285, quindi esamina get_cu_seqlens_on_cp_rank e il percorso THD di azzeramento della coda dQ/dK/dV. Riproduci i casi alternati di pad_between_seqs con il benchmark fornito e confronta il lavoro del kernel e i tempi di esecuzione. Il lavoro è completato quando i metadati equivalenti nei valori evitano lavoro di padding non necessario, mantenendo intatti il comportamento numerico e l’operatività sicura per i grafi.

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

Descrizione

attention

Describe the bug

Commit 4cd705b75394563c0246bdddfa5d3148106c9285 introduces a PyTorch DotProductAttention performance regression under this combination:

  • qkv_format="thd"
  • P2P context parallelism
  • pad_between_seqs=None (automatic detection)
  • both regular and padded cumulative sequence-length tensors are provided
  • the padded and unpadded tensors are distinct objects but have identical relevant values, so there is no actual padding between sequences

This is observable during ordinary eager training; CUDA graph capture does not need to be enabled.

The new automatic detection uses tensor object identity as a proxy for padding semantics:

if cu_seqlens_q_padded is cu_seqlens_q:
    pad_between_seqs = False
elif cu_seqlens_q_padded is not None or cu_seqlens_kv_padded is not None:
    pad_between_seqs = True

Thus independently allocated but value-identical tensors are classified as pad_between_seqs=True. With P2P context parallelism, that selects the path that repeatedly calls get_cu_seqlens_on_cp_rank, rather than the cheaper no-inter-sequence-padding path.

The same commit also adds THD dQ/dK/dV tail-zeroing operations. They launch arange/compare/masked-fill work even when the valid endpoint already equals the tensor endpoint and there is no tail to clear.

Steps/Code to reproduce bug

  1. Create a BF16 DotProductAttention module with qkv_format="thd" and a four-rank P2P context-parallel group.

  2. Provide independently allocated cumulative sequence-length tensors with identical values:

    cu_seqlens = torch.tensor([0, sequence_length], dtype=torch.int32, device="cuda")
cu_seqlens_padded = cu_seqlens.clone()

assert cu_seqlens_padded is not cu_seqlens
assert torch.equal(cu_seqlens_padded, cu_seqlens)

  3. Run repeated attention forward/backward calls, alternating these two cases in the same process:

    • automatic detection: pad_between_seqs=None
    • known-correct metadata: pad_between_seqs=False

  4. Discard warmup and compare steady-state timings. A single attention forward/backward call shows a small direct overhead. The impact becomes much larger in an attention-heavy training schedule where the branch is exercised repeatedly and interacts with context-parallel stream scheduling.

We also performed a controlled source-level reverse experiment on Transformer Engine 2.18.0+27486e03. All arms used the same process, allocation, inputs, configuration, and byte-identical compiled Transformer Engine extensions; only the Python attention hunks from the cited commit differed. Each arm used 50 post-warmup iterations.

Variant Mean iteration time Median Delta vs. stock
Stock 595.084 ms 589.000 ms —
Revert padding detector only 560.958 ms 556.800 ms -5.735%
Revert gradient zero-fill only 577.764 ms 574.950 ms -2.911%
Revert both 554.254 ms 549.650 ms -6.861%

An ABBA repetition of stock and the full reverse patch measured a 6.276% aggregate iteration-time improvement with the reverse patch. The stock and reverse-patched order drift was 1.534% and 0.861%, respectively.

Across all ranks in two Nsight Systems trials, the full reverse patch reduced the five-step trace span by 7.411% on average, with every paired rank faster. Over five steps it removed, per rank:

  • 2,400 helper-generated kernel launches associated with get_cu_seqlens_on_cp_rank
  • 900 masked-fill kernels from the new gradient tail-zeroing blocks

The source reversal reduced main-stream kernel work by 31.702 ms/rank and main-stream gaps by 237.154 ms/rank over the captured five-step window. These are separate trace observations, not additive wall-time attribution. Numerical-health checks remained clean.

Expected behavior

When the padded and unpadded cumulative sequence-length tensors have equal relevant values, automatic detection should not select the inter-sequence-padding path solely because they are different Python objects.

Could the API carry graph-safe padding metadata explicitly, or otherwise avoid using object identity as the semantic proxy? The tail-zeroing work could also be gated when metadata establishes that no gradient tail exists, while preserving CUDA graph compatibility.

The current workaround for callers that know there is no inter-sequence padding is to pass pad_between_seqs=False explicitly.

Environment overview

  • Environment location: containerized bare-metal system
  • Transformer Engine: 2.18.0+27486e03
  • Installation: preinstalled container package

Environment details

  • Python: 3.12.3
  • PyTorch: 2.13.0a0+8145d630e8.nv26.6.54250401
  • CUDA reported by PyTorch: 13.3
  • cuDNN: compiled against 9.23; node-visible runtime 9.21.1

The causal comparison used one unchanged environment, so the cuDNN packaging detail was identical across all variants.

Device details

  • 4x NVIDIA H100 80GB HBM3

Additional context

The detector is the primary contributor. After removing the zero-fill blocks, reverting the detector still improved iteration time by 4.069%. Once the detector was corrected, removing zero-fill added another 1.195%. The effects overlap on the same context-parallel critical path and therefore should not be added independently.

Lingua principale
Python
Stelle
3.6k
Fork
851
Merge medio
5g 1h
PR unite (30g)
52

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.