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

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

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

  • #3506 của @RPalmr — đã đóng, không merge

Đá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
38/100
Loại issue
Lỗi
Độ rõ ràng
Khá rõ ràng
Mức độ hoạt động
Sôi nổi
Công nghệ
python, pytorch
Lĩnh vực
distributed-systems, machine-learning, performance

Hướng nghiên cứu

Bắt đầu từ bộ phát hiện padding của PyTorch DotProductAttention được giới thiệu trong commit 4cd705b75394563c0246bdddfa5d3148106c9285, sau đó kiểm tra get_cu_seqlens_on_cp_rank và đường dẫn zero phần đuôi dQ/dK/dV của THD. Tái hiện các trường hợp pad_between_seqs luân phiên bằng benchmark được cung cấp và so sánh lượng công việc của kernel cũng như thời gian chạy. Được xem là hoàn tất khi metadata tương đương về giá trị tránh được công việc padding không cần thiết, trong khi hành vi số học và hoạt động an toàn với graph vẫn được giữ nguyên.

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

Mô tả

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.

Ngôn ngữ chính
Python
Star
3.6k
Fork
862
Merge trung bình
3 ngày 17 giờ
Pull request đã merge (30 ngày)
55

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.