[Paper Review] Model Parallelism on Distributed Infrastructure: A Literature Review from Theory to LLM Case-Studies

Felix Brakel, Uraz Odyurt+1 more|arXiv (Cornell University)|Mar 6, 2024

Auction Theory and Applications4 citations

TL;DR

This paper provides a comprehensive literature review on model parallelism in distributed deep learning, categorizing it into intra-operator and inter-operator types, identifying key challenges like communication overhead and suboptimal partitioning, and analyzing modern use cases in billion-parameter Transformer models such as Megatron and PaLM. It highlights that hybrid parallelism combining intra- and inter-operator strategies is essential for scaling large models, with specialized hardware enabling efficient intra-operator parallelism in models like PaLM.

ABSTRACT

Neural networks have become a cornerstone of machine learning. As the trend for these to get more and more complex continues, so does the underlying hardware and software infrastructure for training and deployment. In this survey we answer three research questions: "What types of model parallelism exist?", "What are the challenges of model parallelism?", and "What is a modern use-case of model parallelism?" We answer the first question by looking at how neural networks can be parallelised and expressing these as operator graphs while exploring the available dimensions. The dimensions along which neural networks can be parallelised are intra-operator and inter-operator. We answer the second question by collecting and listing both implementation challenges for the types of parallelism, as well as the problem of optimally partitioning the operator graph. We answer the last question by collecting and listing how parallelism is applied in modern multi-billion parameter transformer networks, to the extend that this is possible with the limited information shared about these networks.

Motivation & Objective

• To systematically categorize and define the types of model parallelism used in distributed deep learning systems.
• To identify and analyze the key technical and architectural challenges in implementing model parallelism, including communication bottlenecks and suboptimal partitioning of operator graphs.
• To investigate real-world applications of model parallelism in modern large language models (LLMs), particularly multi-billion parameter Transformers.
• To evaluate the trade-offs between different parallelism strategies in terms of memory usage, execution time, and device utilization.
• To advocate for standardization in DNN auto-parallelization to enable reproducible benchmarking and broader access to the field.

Proposed method

• The study uses a two-phase snowballing literature review approach: first analyzing DNN auto-parallelization techniques for theoretical modeling, and second studying real-world LLM deployments.
• Model parallelism is formalized through operator graphs, where parallelism is defined along two dimensions: intra-operator (within a single operation) and inter-operator (across multiple operations).
• The paper evaluates auto-parallelization frameworks like FlexFlow and TensorOpt using execution simulators that estimate execution time based on device characteristics, tensor sizes, and communication overhead.
• A Markov Chain Monte Carlo search strategy is used in FlexFlow to explore the parallelization search space, guided by simulated performance metrics.
• The study analyzes real-world LLMs (e.g., Megatron, Gopher, PaLM, GPT) by collecting and filtering papers with detailed implementation insights on model parallelism.
• The analysis includes comparing hardware utilization, communication patterns, and the role of specialized hardware (e.g., TPUv4, A100) in enabling efficient intra-operator parallelism.

Figure 1 . Megatron-NLG compared to other large language models (source: (Smith:2022:UDSM, 29 ) ).

Experimental results

Research questions

• RQ1 What types of model parallelism exist, and how are they defined in terms of operator graphs?
• RQ2 What are the primary technical and architectural challenges in implementing model parallelism, especially in distributed systems?
• RQ3 How is model parallelism applied in modern large-scale Transformer models such as Megatron, PaLM, and Gopher?
• RQ4 What are the trade-offs between intra-operator and inter-operator parallelism in terms of communication, memory, and device utilization?
• RQ5 How can the field of DNN auto-parallelization be improved through standardization of representations and benchmarks?

Key findings

• Model parallelism is primarily categorized into intra-operator (within a single operation) and inter-operator (across multiple operations), with hybrid strategies combining both being common in practice.
• Intra-operator parallelism incurs extremely high communication costs due to frequent data shuffling across devices, especially over slower interconnects like Ethernet.
• Inter-operator parallelism often leads to low device utilization during training due to load imbalance and irregular computation patterns.
• The optimal partitioning of operator graphs across devices is a major challenge, as the structure of the model graph rarely aligns well with the underlying device topology.
• Modern LLMs like PaLM use specialized hardware (e.g., TPUv4) to mitigate intra-operator communication costs, enabling efficient model parallelism without relying on inter-operator partitioning.
• Models such as Megatron and Gopher use a combination of intra-operator parallelism within nodes and inter-operator + data parallelism across nodes to scale beyond single-node limits.

Figure 2 . Overview of scaling up NNs within the NN compute infrastructure.