Research map / Large language models / Reasoning and chain of thought
Rethinking Memory and Communication Costs for Efficient Data Parallel Training of Large Language Models
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.
Proposes PaRO, a set of partially redundant data-parallel partitioning strategies (PaRO-DP) plus a topology-aware collective (PaRO-CC) and a cost-model guideline for choosing strategies, reporting up to 266% speedup over ZeRO-3 on heterogeneous GPU clusters.
Problem
Distributed data-parallel training of LLMs is bottlenecked by slow inter-group (inter-node/switch) communication, and existing basic strategies (ZeRO, MiCS, FSDP HYBRID_SHARD) offer limited memory-communication trade-off options, especially for partial-parameter and PEFT training.
Approach
Partition parameters, gradients and optimizer states at three granularities (none, intra-group, global) so that frequent collectives stay within fast intra-group links while accepting limited memory redundancy; pair this with a ring-of-rings collective (PaRO-CC) that separates intra- and inter-group traffic, and a throughput-prediction formula to pick the partition scheme.
Key results
- PIIG achieves up to 266% throughput improvement over ZeRO-3 on LLaMA-65B with partial-parameter training (Ψ′=Ψ/16).
- PaRO-CC reduces all-gather time by 25.2%-39.7% versus Ring from 16 to 128 GPUs and raises Megatron-TP throughput by 17% (129 to 152 samples/s).
- PIIG achieves near-linear scalability (speedup ratio 0.92 from 16 to 128 GPUs) versus 0.63 for ZeRO-2 and ZeRO-3.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use PaRO strategy NII instead of ZeRO-2 | LLaMA-7B full-parameter training · LLaMA-7B full-parameter training | throughput: 67% | Helped |
| Use PaRO strategy PIIG instead of ZeRO-3 | LLaMA-7B full-parameter training · LLaMA-7B full-parameter training | throughput: 128% | Helped |
| Use PaRO strategy PIIG instead of ZeRO-3 | LLaMA-65B full-parameter training · LLaMA-65B full-parameter training | throughput: 86% | Helped |
| Use PaRO strategy PNIG instead of ZeRO-2 | LLaMA-7B partial-parameter training · LLaMA-7B partial-parameter training (Ψ′ = Ψ/16) | throughput: 15% | Helped |
| Use PaRO strategy PNII instead of ZeRO-2 | LLaMA-7B partial-parameter training · LLaMA-7B partial-parameter training (Ψ′ = Ψ/16) | throughput: 21% | Helped |
| Use PaRO strategy PIII (MiCS) instead of ZeRO-3 | LLaMA-65B partial-parameter training · LLaMA-65B partial-parameter training (Ψ′ = Ψ/16) | throughput: 264% | Helped |
| Use PaRO strategy PIIG instead of ZeRO-3 | LLaMA-65B partial-parameter training · LLaMA-65B partial-parameter training (Ψ′ = Ψ/16) | throughput: 266% | Helped |
| Use PaRO strategy PING instead of ZeRO-3 | LLaMA-65B PEFT · LLaMA-65B PEFT (Ψ′ = 3Ψ/1000) | throughput: 44% | Helped |
| Add PaRO-CC to Ring collective communication | Collective all-gather · Communication volume 1GB, 16–128 GPUs | communication time: 25.2%-39.7% | Helped |
| Add PaRO-CC | Megatron-TP end-to-end training · 4 DGX nodes | throughput: from 129 to 152 samples/second | Helped |
| Increase PaRO group/ring scale from 16 to 128 GPUs | PaRO-DP on LLaMA-7B · LLaMA-7B scalability | speedup ratio: 0.92 | Helped |
| Increase ZeRO scale from 16 to 128 GPUs | ZeRO-2 and ZeRO-3 on LLaMA-7B · LLaMA-7B scalability | speedup ratio: 0.63 | Mixed |
Surprising findings
- The authors observe that in PEFT with LLaMA-65B, PaRO's ING strategy improves throughput by 44% over ZeRO-3 while using more memory (47.71 GB vs 29.54 GB).
- Inter-group reduce-scatter needed for optimizer state can be avoided only under certain partition combinations; 13 of 27 partition combinations were eliminated by the authors' insight about optimizer state granularity.
Limitations
- Gains depend on heterogeneous networks with slow inter-group links; limited applicability to homogeneous networks.
- Extremely large models would need combination with tensor or pipeline parallelism, reducing ease of use.
- Partition strategies trade extra memory for communication, so some strategies OOM at large model sizes.
Experimental setup
Datasets: C4 (RedPajama)
Metrics: training throughput (samples/sec) peak GPU memory (GB) collective communication time (ms) loss convergence
Baselines: ZeRO-1 ZeRO-2 ZeRO-3 MiCS ZeRO++ FSDP HYBRID_SHARD (FSDP-hsz) Megatron-TP Ring all-gather
Builds on: ZeRO MiCS ZeRO++ FSDP Megatron-LM
Closest papers by meaning
- Communication-Efficient Language Model Training Scales Reliably and Robustly: Scaling Laws for DiLoCo (NeurIPS 2025 · similarity 0.89)
- HALoS: Hierarchical Asynchronous Local SGD over Slow Networks for Geo-Distributed Large Language Model Training (ICML 2025 · similarity 0.88)
- Benchmarking Optimizers for Large Language Model Pretraining (NeurIPS 2026 · similarity 0.88)
- Revisiting Parameter Server in LLM Post-Training (ICLR 2026 · similarity 0.88)
- EDiT: A Local-SGD-Based Efficient Distributed Training Method for Large Language Models (ICLR 2025 · similarity 0.88)
- ZeCO: Zero-Communication Overhead Sequence Parallelism for Linear Attention (NeurIPS 2025 · similarity 0.88)
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.