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

[Feature Request] Support dynamic token counts in TP communication/GEMM overlap

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

Đá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
45/100
Loại issue
Tính năng
Độ 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
ai-infra-agents, machine-learning, performance

Hướng nghiên cứu

The issue is about TP communication/GEMM overlap in Transformer Engine for variable token counts. Start by reading the prototype description and the linked THD layout documentation. Look at the UserBuffer initialization and overlap paths in the codebase, focusing on bulk, pipeline, and ring-exchange communicators. Understand how tensor views and communication counts are currently computed. The goal is to modify these to use an active shape instead of the fixed buffer size, ensuring the changes work for BF16 and maintain compatibility with TP size divisibility.

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

Mô tả

Problem

Our use case is packed-sequence training with THD-format inputs and dynamically sized microbatches. In TE's THD layout, T is the total number of packed tokens in a microbatch. As the microbatch size and sequence lengths change, the flattened activations used by TP communication/GEMM overlap have shape [T, hidden_size], and T can differ from one microbatch to the next.

The TP-overlap UserBuffer is initialized with a fixed shape. In the v2.11 code on which my prototype is based, overlap paths use the registered buffer size for tensor views and communication counts. As a result, they cannot directly operate on a smaller [T, hidden_size] tensor while retaining an allocation registered for the maximum token count.

I would like to support this workload without allocating and registering new UserBuffers for different microbatch sizes or padding every microbatch to the maximum token count. I have a prototype based on a fixed-capacity allocation and a per-microbatch active shape, described below for design discussion.

This is related to #1303, which asked whether TP overlap supports variable sequence length. This issue adds a concrete workload and an implementation proposal.

Proposed API

# Distributed setup and overlap configuration are omitted here.
# Existing initialization uses the maximum token capacity.
te.initialize_ub(
    [max_tokens, hidden_size],
    tp_size,
    quantization_modes=[te.UserBufferQuantizationMode.NONE],
    dtype=torch.bfloat16,
)

# Before processing a microbatch with fewer tokens:
te.set_ub_active_shape([active_tokens, hidden_size])
# Run the model with tensors sized for active_tokens.

The prototype exposes set_ub_active_shape() on the PyTorch UserBuffer manager. It also exposes set_buffer_active_shape() and active/capacity shape getters on individual CommOverlap and CommOverlapP2P communicators. The manager updates its unquantized bulk, pipeline, and ring-exchange communicators; quantized and external UserBuffers are left at their capacity shape.

Implementation approach in the prototype

Keep allocation capacity separate from the active shape. initialize_ub() registers a buffer sized for the maximum token count. Each communicator stores that capacity shape and a separate active shape. set_buffer_active_shape() updates the logical shape and tensor views; it does not allocate, free, or re-register the underlying UserBuffer. The active data occupies a contiguous prefix of the registered allocation.

Use active sizes throughout the overlap paths. For bulk AG/RS, buffer views, local-chunk offsets, copy-size checks, and communication counts use the active element count rather than the registered capacity. For pipeline and ring exchange, the P2P layout rebuilds its per-rank chunk views and byte offsets using the active first dimension. Reduce-scatter retains its extra P2P chunks within the original allocation.

Constrain the initial behavior. The active shape must be two-dimensional and nonzero, fit within the registered capacity, keep the hidden dimension unchanged, and have a first dimension divisible by TP size. The current manager API applies one active shape to its eligible communicators. Participating TP ranks must therefore use the same active token count before the corresponding overlap operations. The prototype targets unquantized, 2-byte/BF16 UserBuffers for bulk, pipeline, and ring-exchange overlap; dynamic FP8/quantized UserBuffers are outside this initial scope.

The prototype is based on TE v2.11 and has not yet been ported or validated against current main.

Questions for maintainers

  1. Would Transformer Engine maintainers be interested in supporting TP communication/GEMM overlap when the token count varies between microbatches, as in the THD packed-sequence workload above?
  2. If so, would registering a UserBuffer at maximum capacity and updating its active shape for each microbatch be an acceptable direction for an upstream contribution?

If the overall direction makes sense, I would also appreciate guidance on whether a manager-level update is the right API and whether BF16 bulk, pipeline, and ring-exchange paths are a useful first scope.

I have a prototype implementation and can port it to current main and submit a focused PR if this direction makes sense.

Ngôn ngữ chính
Python
Star
3.6k
Fork
851
Merge trung bình
4 ngày 15 giờ
Pull request đã merge (30 ngày)
51

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.