Zellige: Moldable Sequence Placement for Mixed Image-Video DiT Training

2026-08-02

AI Summary

This work addresses load imbalance and communication redundancy in multi-GPU training of Diffusion Transformers (DiT) with heterogeneous-length image and video sequences. To overcome the limitations of conventional fixed-group parallelism, the authors propose a malleable sequence placement system that jointly optimizes per-sequence parallel strategies and GPU assignments through a two-stage anchor-and-fill scheduling approach and a holistic-chunk fused attention mechanism. The system integrates hardware-aware performance modeling, a two-stage sequence planner, an aggregated attention engine, and dynamic parallel configuration selection. Experiments demonstrate speedups of 1.12–1.48× on 16×A800 GPUs and 1.27–1.54× on 32×A6000 GPUs, with planning errors below 0.32%.

Abstract

High-quality video generation requires training Diffusion Transformers (DiTs) jointly on image and video data, posing a mixed-length sequence training problem across GPUs. Existing systems rely on data parallelism (DP), context parallelism (CP), or their combination; we model these designs as disjoint-group placement and prove that they face a fundamental tradeoff between inter-group load imbalance and intra-group communication redundancy. We present Zellige, a moldable sequence placement system that jointly selects each sequence's parallelism configuration and participating ranks. Zellige consists of three components: a hardware profiler that estimates the execution time and memory consumption of candidate placements, a two-stage planner that balances compute-heavy anchor sequences and packs lighter filler sequences into the remaining capacity, and a coalesced attention engine that efficiently executes whole sequences alongside distributed-attention shards. Across 21 plans, the hardware profile predicts step makespan and peak allocated memory with mean absolute percentage errors of $3.4%$ and $1.5%$, respectively. The two-stage planner solves each batch in 33--119 ms, significantly faster than a joint-placement reference that optimizes all sequences together, while their modeled makespans differ by at most $0.32%$. In end-to-end evaluations, Zellige outperforms KnapFormer by $1.12$--$1.48\times$ on 16 A800 GPUs and $1.27$--$1.54\times$ on 32 A6000 GPUs.

Problem

• mixed-length sequence training
• Diffusion Transformers
• load imbalance
• communication redundancy
• sequence placement

Innovation

• moldable sequence placement
• Diffusion Transformer
• mixed-length sequence training
• two-stage planner
• coalesced attention