[QUESTION] Do tp overlap support thd, whose sequence length is flexible?

[wplf]

Hi, thank you for great works.

I have a question about tp-overlap.
The function below make a buffer for args.seq_length * args.micro_batch_size. Do this support thd format?

def _initialize_tp_communicators():
    """ initializing the communicators with user buffers for high-performance tensor-model-parallel 
        communication overlap """

    try:
       import yaml

       import transformer_engine
       from transformer_engine.pytorch import module as te_module

    except ImportError:
       raise RuntimeError("Tensor Parallel Communication/GEMM Overlap optimization needs 'yaml' and "
             "'transformer_engine' packages") 

    args = get_args()

    if args.tp_comm_overlap_cfg is not None:
       with open(args.tp_comm_overlap_cfg,"r") as stream:    
          ub_cfgs = yaml.safe_load(stream)
    else:
       ub_cfgs = {}

    input_shape = [(args.seq_length * args.micro_batch_size) // args.context_parallel_size , args.hidden_size]

    #We create a MPI process group, which is needed to bootstrap the pipelined 
    #tensor-model-parallel communication overlap
    torch.distributed.new_group(backend='nccl')
    te_module.base.initialize_ub(shape = input_shape, tp_size = args.tensor_model_parallel_size, 
                                 use_fp8 = (args.fp8 is not None) , ub_cfgs = ub_cfgs,)

Follow this question, I have found that after TP/SP mlp layer, the output shape is exactly seqlen, args.hidden_size. So how does that works for qkv_proj hidden_dim * 3/ tp_size and mlp hidden_dim * 2 / tp_size?

[Advait251206]

The TP overlap implementation was originally designed assuming a static (sequence_length × micro_batch_size, hidden_size) layout, so support for truly dynamic THD (variable sequence length) depends on how the buffers are managed.

From the code you shared:

input_shape = [
    (args.seq_length * args.micro_batch_size) //
    args.context_parallel_size,
    args.hidden_size,
]

te_module.base.initialize_ub(
    shape=input_shape,
    tp_size=args.tensor_model_parallel_size,
    ...
)

the user buffers (UBs) are allocated once during initialization using the maximum configured sequence length (args.seq_length). There is no mechanism here to resize or reallocate the buffers on a per-batch basis.

Therefore, if your THD batches contain varying sequence lengths, the TP overlap implementation does not dynamically adjust the communication buffers. Instead, one of the following must happen:

• The implementation pads shorter sequences to the configured args.seq_length.
• The buffers are simply oversized, with only a portion of the allocated memory being actively used.
• True dynamic shapes are not currently supported by the overlap path.

So, unless there have been recent changes in Transformer Engine's UB implementation, TP overlap should be considered optimized for fixed-size communication buffers, not arbitrary THD layouts.

---

Why doesn't the hidden dimension appear in initialize_ub()?

Your second question is a good observation.

The buffer is initialized as

(seqlen * micro_batch, hidden_size)

even though operations like

• QKV projection produce 3 * hidden_size / TP
• MLP expansion produces 2 * hidden_size (or larger, depending on the architecture)

At first glance, this seems inconsistent.

The key point is that the UB shape represents the communication tensor, not the GEMM output tensor.

For example, consider the QKV projection:

Input
-------------------
[S, H]
        │
        ▼
ColumnParallelLinear
        │
        ▼
[S, 3H/TP]

Each TP rank computes only its local shard of the projection. The communication overlap is concerned with moving the input activations (or other communicated tensors), whose hidden dimension remains H, rather than allocating space for the expanded projection output.

Similarly, for the MLP:

Input
[S, H]
    │
    ▼
ColumnParallelLinear
    │
[S, 2H/TP]
    │
GELU/SwiGLU
    │
RowParallelLinear
    │
[S, H]

The intermediate expansion (2H/TP, 4H/TP, etc.) remains local to each TP rank and is not the tensor being staged in the overlap communication buffer.

In other words:

• Communication buffers are sized according to the communicated activation tensor.
• Expanded GEMM outputs remain local to each rank and are handled by the GEMM kernels themselves.

That's why initialize_ub() only needs

[S, hidden_size]

instead of

[S, 3 * hidden_size / TP]

or

[S, 2 * hidden_size]

---

Does this work with THD?

If by THD you mean variable token counts per batch (e.g., packed sequences with different effective lengths), then the current initialization suggests that TP overlap is still based on a statically sized communication buffer.

Whether it works correctly depends on how the runtime handles the unused portion of the buffer:

• If the actual token count is ≤ args.seq_length × micro_batch_size, the communication kernels can operate on a subset of the allocated buffer.
• If the runtime expects the communication tensor to exactly match the configured shape, then padding (or another fixed-length representation) is required.

Without additional logic to recreate or resize the UB, there doesn't appear to be native support for fully dynamic THD shapes in the TP overlap path.

---

So, in summary:

• The UB allocation is based on the communication tensor shape, not the intermediate GEMM output shape.
• QKV (3H/TP) and MLP expansion (2H, 4H, etc.) are local computations and therefore do not determine the UB dimensions.
• Based on the initialization code, TP overlap appears to assume a fixed maximum sequence length, so true variable-length THD support is unlikely unless Transformer Engine internally supports dynamic slicing over a preallocated buffer.

---

If this answer helped or pointed you in the right direction, I'd appreciate it if you could mark it as the accepted answer so it's easier for others with the same issue to find.

Also, if you found my contribution useful, I'd appreciate it if you could check out my GitHub profile, follow me, and star any repositories you find interesting.

GitHub: https://github.com/Advait251206