Need graph-safe multi-input TransformerBlock support for diffusion-style models with PP and CUDA graphs
- 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
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