NVIDIA / NVIDIA/Megatron-LM

Need graph-safe multi-input TransformerBlock support for diffusion-style models with PP and CUDA graphs

Open
#5,620 5 comments 0 reactions 0 assignees View on GitHub
community-request enhancement waiting-on-customer
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 3h
Merged PRs (30d)
272

Description

**Is your feature request related to a problem? Please describe.**

Diffusion-style transformer models often require auxiliary conditioning tensors, e.g. timestep embeddings for AdaLN/DiT blocks, text/image conditioning, or other per-layer inputs. These tensors need to be passed through every transformer layer and, with pipeline parallelism, may also need to cross PP stage boundaries with gradients.

Tag @NVIDIA/mcore-oncall.

Today, custom models can subclass `TransformerLayer` and override `get_layer_static_inputs()`, but end-to-end support is still difficult:

- `TransformerBlock.forward()` has a fixed layer-call signature and does not provide a generic way to pass custom Tensor kwargs such as `condition_emb` to all layers.
- TE CUDA graph capture requires all graph inputs to be Tensor/None and needs matching static sample kwargs.
- Pipeline schedules and tensor-shape helpers are primarily centered around a single hidden-state activation tensor. Multi-input or named activation transport requires custom hooks or packing workarounds.
- Recomputing conditioning modules independently on each PP stage can change model semantics for DiT-style models where a shared conditioning module should feed all layers.

This affects diffusion language models, DiT/AdaLN models, and multimodal models that are not simple GPT-style single-hidden-stream models.

**Describe the solution you'd like**

It would be useful to have an official graph-safe multi-input extension path for transformer blocks, for example:

1. Allow `TransformerBlock.forward()` to pass a generic `extra_layer_kwargs: dict[str, Tensor]` or named Tensor kwargs through to every layer.
2. Extend `get_layer_static_inputs()` / `TECudaGraphHelper` to support custom static Tensor kwargs in a documented way.
3. Provide a graph-safe packing/unpacking convention for non-Tensor metadata and Tensor-only CUDA graph inputs.
4. Provide PP support for named/list activation tensors, including tensor shape metadata and backward gradient propagation.
5. Add a small example or test for a DiT/AdaLN-style block with inputs like:
- `hidden_states`
- `attention_mask`
- `packed_seq_params` / THD metadata
- `condition_emb`

**Describe alternatives you've considered**

Current workarounds include:
- Copying `TransformerBlock.forward()` and manually adding custom kwargs.
- Packing auxiliary tensors into `hidden_states` for PP transport.
- Recomputing conditioning modules on every PP stage.
- Disabling PP or CUDA graphs for diffusion-style models.

These workarounds are brittle and make it harder to keep custom models compatible with future Megatron-Core / Megatron-Bridge versions.

**Additional context**

Megatron-Bridge diffusion examples provide useful local patterns, but they do not appear to demonstrate an end-to-end TP + PP + CP + TE CUDA graph path for a multi-input DiT-style transformer.

Contributor guide

Open the contributing guide

Research direction

Start by reading TransformerBlock.forward(), get_layer_static_inputs(), and TECudaGraphHelper, then inspect the pipeline schedule and tensor-shape helpers mentioned in the issue. Compare them with the Megatron-Bridge diffusion examples. Done should include a documented graph-safe multi-input path and an end-to-end DiT/AdaLN example or test covering TP, PP, CP, CUDA graphs, and gradient transport.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
distributed-systems, machine-learning, performance
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.