Research map / Architectures and efficiency / RNNs and normalization
Pfeife: Automatic Pipeline Parallelism for PyTorch
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Pfeife is a PyTorch torch.compile backend that automatically pipelines models across GPUs by co-optimizing graph slicing and looped 1F1B scheduling with prefetching, using critical-path cost estimation.
Problem
Large ML models do not fit in a single GPU, and existing pipeline parallelism tools need manual cutpoints, model rewriting, or only support sequential/transformer models, making automatic pipelining of general models (e.g., non-sequential Stable Diffusion) hard.
Approach
Capture the data-flow graph via TorchDynamo and jointly search graph slicing (beam search over split points) and looped 1F1B schedule parameters (loop count, loop batch, prefetch count) with a critical-path cost model, finding that uneven slicing with prefetching is often optimal.
Key results
- Pfeife outperforms DeepSpeed and Colossal-AI on ViT-g/14 by up to 22% (4 devices, 4 batches: 151.4 vs 124.5 items/s) and reaches up to 10% higher throughput than Colossal-AI on Llama2-7B.
- Pfeife trains SDXL U-Net (non-sequential) with pipelining, achieving 3-4x the throughput of ZeRO3 and avoiding OOM in FP32 where ZeRO2/ZeRO3 fail.
- Of 49 TorchBench models, 37 run successfully with identical outputs under Pfeife; 11 fail due to torch.compile issues and one produces incorrect results.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use Pfeife instead of DeepSpeed | ViT-g/14 pipeline parallelism · ViT-g/14, 4 devices, 4 micro-batches | throughput (item/s): 124.5 -> 151.4 (+22%) | Helped |
| Use Pfeife instead of DeepSpeed | ViT-g/14 pipeline parallelism · ViT-g/14, 8 devices, 16 micro-batches | throughput (item/s): 297.7 -> 282.2 (-5.5%) | Hurt |
| Use Pfeife instead of Colossal-AI | Llama2-7B pipeline parallelism · Llama2-7B, 8 devices, 8 micro-batches | throughput (item/s): 6.80 -> 7.46 (+10%) | Helped |
| Use Pfeife instead of ZeRO3 | SDXL training · SDXL, 8 devices, FP32 | throughput (item/s): 3.22 -> 6.09 | Helped |
| Add prefetching of forward passes before the first backward pass (ablation) | Pfeife pipeline schedule · ViT-g/14 pipeline parallelism | latency: 8% speedup | Helped |
| Add a looped schedule to a schedule with prefetching (ablation) | Pfeife pipeline schedule · ViT-g/14 pipeline parallelism | latency: by 9% | Helped |
| Use uneven slicing with prefetching instead of uniform weight or time slicing (ablation) | Pfeife pipeline schedule · ViT-g/14, 2 and 4 devices | latency | Helped |
| Use uneven slicing without prefetching (ablation) | Pfeife pipeline schedule · ViT-g/14 pipeline parallelism | latency | Hurt |
| Use Pfeife’s predicted split points rather than the optimal split points (ablation) | Pfeife pipeline schedule · ViT-g/14, 2 devices | latency | No clear effect |
| Use Pfeife’s cost model estimate rather than actual latency for schedule selection (ablation) | Pfeife pipeline schedule · ViT-g/14, 2 and 4 devices | latency: 30-40ms | No clear effect |
| Use pipeline parallelism rather than ZeRO2 data parallelism | SDXL training · SDXL, 8 devices, FP16 | throughput (item/s): 19.8 -> 12.1 | Hurt |
| Use pipeline parallelism rather than ZeRO2 data parallelism | SDXL training · SDXL, 2 devices, FP16 | throughput (item/s): 3.78 -> 3.72 | No clear effect |
Surprising findings
- Unbalanced (uneven) slicing combined with prefetching is often optimal, contrary to the common belief that work should be distributed evenly across devices.
- Prefetching is necessary to take advantage of uneven slicing.
Limitations
- Linearizes the computation graph via TorchDynamo topological order, approximating models with parallel branches
- 11 TorchBench models fail under torch.compile and one produces incorrect results
- Underperforms DeepSpeed for ViT with 8 devices and batch size 16
Experimental setup
Benchmarks: TorchBench ViT-g/14 throughput Llama2-7B throughput Stable Diffusion XL throughput
Metrics: throughput (items/s) latency (ms) peak memory speedup
Baselines: DeepSpeed Colossal-AI ZeRO-2 ZeRO-3 Alpa PiPPy
Builds on: PyTorch torch.compile / TorchDynamo DAPPLE Megatron-LM looped 1F1B PipeDream GPipe
Closest papers by meaning
- PyTorch: An Imperative Style, High-Performance Deep Learning Library (NeurIPS 2019 · similarity 0.89)
- Synergistic Tensor and Pipeline Parallelism (NeurIPS 2025 · similarity 0.87)
- KernelBench: Can LLMs Write Efficient GPU Kernels? (ICML 2025 · similarity 0.86)
- Nesterov Method for Asynchronous Pipeline Parallel Optimization (ICML 2025 · similarity 0.85)
- PipeOffload: Improving Scalability of Pipeline Parallelism with Memory Optimization (ICML 2025 · similarity 0.85)
- AsyncMesh: Fully Asynchronous Optimization for Data and Pipeline Parallelism (NeurIPS 2026 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.