Search papers, labs, and topics across Lattice.
This paper introduces Zellige, a moldable sequence placement system designed to optimize the training of Diffusion Transformers (DiTs) on mixed image and video data by addressing the challenges of sequence length variability across GPUs. The authors model existing parallelism strategies as disjoint-group placements and demonstrate that these approaches suffer from inherent tradeoffs, which Zellige effectively mitigates through a combination of a hardware profiler, a two-stage planner, and a coalesced attention engine. Experimental results show that Zellige achieves significant performance improvements, outperforming KnapFormer by up to 1.54 times across various GPU configurations while maintaining low prediction errors in resource usage.
Zellige redefines sequence placement for DiT training, achieving up to 1.54x faster performance than existing methods while optimizing resource utilization.
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.
