Skip to content

Synergistic Tensor and Pipeline Parallelism

🕒 Published (v1): 2025-01-01 · Source: NeurIPS · Venue: NeurIPS 2025 · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

This paper proposes STP, a hybrid parallelism schedule that simultaneously reduces both tensor-parallelism (TP) communication bubbles and pipeline-parallelism (PP) idle bubbles during distributed LLM/MLLM training. It decomposes Transformer layers into fine-grained computation units and interleaves forward and backward passes into "braided execution blocks" that overlap TP All-Reduce communication with PP computation. On Qwen2-VL (a multimodal LLM), STP achieves up to 16.7% throughput improvement over existing schedules.

Problem

Hybrid TP+PP parallelism—the dominant approach for training large LLMs and MLLMs—suffers from two independent inefficiency sources: TP introduces All-Reduce communication bubbles (growing to 27.5% of total time at TP=8, seq=6144), and PP introduces pipeline idle bubbles. Prior work addresses only one at a time (e.g., kernel-level overlap for TP, or Zero Bubble scheduling for PP), and their combination can worsen overall efficiency (ZB-V's backward decoupling exposes previously-overlapped TP bubbles, making it no better than 1F1B-I in practice).

Method

Braided Execution Blocks: Each Transformer layer is decomposed into four fine-grained units—Pre-Attn, Attn, Pre-MLP, MLP—and their forward/backward counterparts. Forward units from one microbatch are interleaved with backward units from another microbatch within the same PP stage, so TP All-Reduce operations (which cannot be overlapped within a single pass due to data dependencies) run concurrently with weight-gradient computation from the paired backward pass. Residual connections are fused into Attn/MLP units to eliminate additional data dependencies.

"V"-Shape PP Schedule: Building on these blocks, the schedule implements a "V"-shape dataflow across virtual stages to balance peak activation memory across devices (fixing the stage-0 bottleneck in 1F1B-I). The schedule proceeds through warm-up (maximize in-flight microbatches, activate weight-gradient separation), steady (deactivate weight separation; alternate F&B for chunk 1 then chunk 0), and cool-down phases.

Enhanced Variant: An activation offloading variant targets memory-constrained settings, offloading chunk-0 activations (long lifespan) to CPU during the steady phase via PCIe, reducing peak memory by 10–19.2% with negligible throughput loss.

Key Contributions

  • Synergistic reframing: jointly optimize TP and PP bubbles at the scheduling (software) level rather than via kernel modifications.
  • Braided execution blocks that achieve near-complete TP bubble elimination by overlapping All-Reduce with PP backward computation.
  • "V"-shape PP schedule that reduces the PP bubble rate while maintaining balanced memory footprints across stages.
  • Theoretical analysis showing TP bubble cost scales as O(p) (in stages) rather than O(m) (in microbatches) under the proposed schedule, versus 2m·T_AR for 1F1B-I and 4m·T_AR for ZB-V.
  • Activation offloading variant for limited-memory scenarios.

Results

  • LLM (Qwen2-12.1B, 16Ă—A800, TP=8, PP=2, seq=6144): up to +12.2% throughput over 1F1B-I.
  • LLM (Qwen2-26.3B, 32Ă—A800): consistent improvements over 1F1B-I and ZB-V across all TP/PP/seq configurations.
  • MLLM (Qwen2-VL-14.9B, 16Ă—A800, TP=8, PP=4, balanced ViT/LM FLOPs): +11.7% over 1F1B-I; +16.7% in PP=2 with imbalanced ViT/LM workloads.
  • MLLM (28.8B, 32Ă—A800, TP=8, PP=4): +8.5% throughput over 1F1B-I.
  • Enhanced variant (16Ă—H20, TP=2, PP=8): 92.86% MFU vs. 92.09% for 1F1B-I, with peak memory reduced from 76 GB to 68 GB.
  • ZB-V consistently underperforms or matches 1F1B-I due to exposed TP bubbles from backward decoupling.

Limitations

  • Higher peak activation memory than baselines (e.g., 3p·Ma vs. 2p·Ma for ZB-V), requiring the offloading variant in memory-constrained settings.
  • Performance gains diminish at large PP sizes (PP=8) because explicit inter-stage pipeline communication cannot be overlapped with computation.
  • Theoretical analysis excludes multimodal models; MLLM ViT/LM imbalance requires special virtual-stage assignment heuristics.
  • Offloading variant effectiveness is hardware-dependent (PCIe Gen 5 on H20 vs. A800 where gains are limited).
  • Source code built on Megatron-Core; portability to other frameworks not demonstrated.

Relevance to Vision-Language Models

Multimodal LLMs such as Qwen2-VL present a qualitatively different parallelism challenge compared to LLMs: the ViT encoder and language model have mismatched hidden sizes and layer counts, breaking the assumption of uniform virtual-stage execution times. This paper directly confronts that asymmetry—assigning ViT to its own virtual stage and calibrating FLOPs balance—and demonstrates that synergistic TP+PP scheduling yields larger gains on MLLMs (up to 16.7%) than on pure LLMs (12%), making it particularly relevant for teams training or fine-tuning large VLMs at scale. For researchers tracking VLMs, this work offers a practical, software-level (no custom CUDA kernels) training acceleration method compatible with the standard Megatron-Core stack.