Hacktoberfest 2026: los issues que los mantenedores marcaron para octubre, abiertos y aptos para principiantes. Explorar issues de Hacktoberfest

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

Abierto
  • ¿Algo no está bien?
#3,564 2 comentarios 0 reacciones 0 asignados Ver en GitHub

Los mantenedores suelen responder en 2 días

Nadie ha tomado este issue todavía.

Evaluación

Dificultad
4/5
Tiempo estimado
3-5 días
Aptitud para principiantes
45/100
Tipo de issue
Nueva funcionalidad
Claridad
Bien especificado
Estado de actividad
Activo
Stack tecnológico
python, pytorch
Área
ai-infra-agents, machine-learning, performance

Línea de trabajo

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.

Escrito por el modelo de indexación a partir del texto del issue.

Descripción

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.

Lenguaje dominante
Python
Estrellas
3.6k
Forks
844
Merge medio
4 d 42 min
PR fusionados (30 d)
49

Preparar el entorno

Primeros pasos

  1. Lee el issue completo y luego la guía de contribución del proyecto.
  2. Comenta en el issue que vas a ocuparte — evita que dos personas hagan lo mismo.
  3. Haz un fork del repositorio y trabaja en una rama.
  4. Abre un pull request que haga referencia al número del issue.

Más de NVIDIA/TransformerEngine

Todos los issues de NVIDIA/TransformerEngine

Issues similares

Más issues de Python

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.