Pulse: Training Acceleration for Large Diffusion Models with Automatic Pipeline Parallelism

2026-06-17

AI Summary

This work addresses the severe communication overhead in pipeline-parallel training of large-scale diffusion models on GPU clusters, caused by long-range skip connections in the UNet architecture, which significantly limits training efficiency. To tackle this, the study introduces an automated pipeline parallelism strategy that treats skip connection locality as a primary optimization objective. It employs a skip-aware dynamic programming partitioner to colocate encoder-decoder layers on the same device and locally cache activations, thereby eliminating cross-device communication. This approach is further integrated with an integer linear programming–based bubble-efficient scheduler and a hybrid parallelism tuner to achieve end-to-end co-optimization. On communication-constrained hardware, the proposed method reduces communication volume by 89% and improves training throughput by up to 2.3× compared to state-of-the-art baselines.

Abstract

Diffusion models are now a dominant approach for high-fidelity image and video generation, yet scaling their training across GPU clusters remains challenging. Unlike transformer-only architectures, diffusion backbones commonly adopt UNet-style encoder-decoder structures with heterogeneous layers and long-range skip connections. Under conventional pipeline parallelism, these non-local dependencies force large skip activations and their gradients to traverse multiple pipeline boundaries, making peer-to-peer (P2P) communication a dominant bottleneck and substantially reducing pipeline efficiency. In this paper, we present PULSE, an automatic pipeline-parallel training strategy that makes skip locality a first-class optimization objective. PULSE eliminates skip-induced communication by collocating skip-connected encoder-decoder layers on the same device and caching skip activations locally for later use in backpropagation. To realize this placement while maintaining high pipeline utilization, PULSE co-designs: (1) a skip-aware dynamic-programming partitioner that balances heterogeneous stage workloads under symmetric collocation constraints, (2) an ILP-based schedule synthesizer that generates bubble-efficient wave schedules for the resulting stage-to-device mapping, and (3) a hybrid parallelism tuner that selects pipeline/data-parallel degrees and microbatch sizes under memory and network constraints. Our extensive experiments show that the volume of communication can be reduced by 89 percent, and the training throughput can be increased by up to 2.3x on communication-bound hardware, compared with state-of-the-art parallelism strategies.

Problem

• diffusion models
• pipeline parallelism
• skip connections
• communication bottleneck
• training acceleration

Innovation

• pipeline parallelism
• diffusion models
• skip connections
• automatic partitioning
• communication optimization