Distributed Training and Inference

• Concise summary in the introduction and related work here: Aji 2017
• For a modern overview, see section 3.4 of Liu et al 2024 - Understanding LLMs: A Comprehensive Overview from Training to Inference
• A good overview is in section 3.3.2 of the The Llama 3 technical report Uses tensor parallelism, pipeline parallelism, context parallelism and data parallelism.
• Verbraeken et al 2019 - A Survey on Distributed Machine Learning
• Chaha et al 2018 - A Hitchhiker’s Guide On Distributed Training of Deep Neural Networks Older, but thorough at the time
• Targeted Overviews
  • Ouyang & Dong 2020 - Communication optimization strategies for distributed deep neural network training: A survey
  • Gao et al 2022 - Deep Learning Workload Scheduling in GPU Datacenters: Taxonomy, Challenges and Vision
  • Liu et al 2024 - Acceleration for Deep Reinforcement Learning using Parallel and Distributed Computing: A Survey
  • Liang et al 2024 - Communication-Efficient Large-Scale Distributed Deep Learning: A Comprehensive Survey
• Blog posts
  • Parameter server Warning: contains some conceptual errors
  • An Overview of Pipeline Parallelism and its Research Progress

Used as the optimizer in both data-parallel and model-parallel training (usually the method can be used for both).

• Parameter server
• Hogwild
• Seide et al 2014 - 1-Bit Stochastic Gradient Descent and its Application to Data-Parallel Distributed Training of Speech DNNs
• Aji & Heafield 2017 - Sparse Communication for Distributed Gradient Descent Has a good summary of related work
• LAMB optimizer: You et al 2020 - Large Batch Optimization for Deep Learning: Training BERT in 76 minutes Distributed on a single TPU Pod. Reduces training time from 3 days to 76 minutes on a TPU Pod.

• Li et al 2020 - PyTorch Distributed: Experiences on Accelerating Data Parallel Training

Pipeline parallel, which partitions groups of layers across different accelerators in a pipeline, can be thought of as a special case of model parallel. This is used in GPipe, etc. Tensor slicing (or tensor parallel) is another type of model parallelization, used in Megatron-LM. Megatron-Turning NLG uses both types.

• Krizhevsky 2014 - One weird trick for parallelizing convolutional neural networks
• Shazeer et al 2018 - Mesh-TensorFlow: Deep Learning for Supercomputers Distributes tensor operations (matrix multiplication) across devices through by tiling across devices. Used in T5.
• Huang et al 2018 - GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism Partitions groups of layers across different accelerators in a pipeline, resulting in linear speedup across multiple accelerators. Works because not much communication overhead between layers. Megatron-Turning NLG uses this (page 6) along with other methods.
• Shoeybi et al 2019 - Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism This paper introduced what is now called tensor parallelism. Proposes a particular split across devices for both the feedforward and attention layers in the Transformer. Used to train a 8.3B LM. blog github
• Rajbhandari et al 2019 - ZeRO: Memory Optimizations Toward Training Trillion Parameter Models Allows training up to 100B LM.
• Lepikhin et al 2020 - GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding Trains a Sparsely-Gated Mixture-of-Experts model up to 600B LM.
• Ren et al 2021 - ZeRO-Offload: Democratizing Billion-Scale Model Training Code and tutorial
• Rajbhandari et al 2021 - ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning
• Zheng et al 2022 - Alpa: Automating Inter- and Intra-Operator Parallelism for Distributed Deep Learning
• Burham et al 2022 - Pathways: Asynchronous Distributed Dataflow for ML Used to train PaLM
• Qi et al 2024 - Zero Bubble Pipeline Parallelism
• Lin et al 2024 - Infinite-LLM: Efficient LLM Service for Long Context with DistAttention and Distributed KVCache
• Shyam et al 2024 - Tree Attention: Topology-aware Decoding for Long-Context Attention on GPU clusters

• Lin et al 2024 - Infinite-LLM: Efficient LLM Service for Long Context with DistAttention and Distributed KVCache

• Wang et al 2022 - TopoOpt: Co-optimizing Network Topology and Parallelization Strategy for Distributed Training Jobs

• PyTorch Distributed
• DeepSpeed, Github Highly optimized distributed training for PyTorch. See their related papers like ZeRO-Offload
• Mesh-Tensorflow
• Megatron-LM Distributed training for language models. Used for example by Bloom.
• HuggingFace Accelerate Easy distributed training for PyTorch models (blog post introduction) I believe this is essentially GPipe
• Alpa paper ICML 2022 Tutorial Developed at Berkeley
• PyTorch Fully Sharded Data Parallel paper This is what Chris T. used for pre-training experiments

• MLSys

• Large-Scale ML
• GPU Deep Learning
• Systems & ML