[论文解读] 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 Applications被引用 4

一句话总结

本文对分布式深度学习中的模型并行性进行了全面的文献综述，将模型并行性分为算子内（intra-operator）和算子间（inter-operator）两类，识别出通信开销和算子图划分不优等关键挑战，并分析了在百亿参数级Transformer模型（如Megatron和PaLM）中的现代应用场景。研究指出，结合算子内与算子间策略的混合并行性对于扩展大规模模型至关重要，专用硬件使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.

研究动机与目标

• 系统性地对分布式深度学习系统中使用的模型并行性类型进行分类和定义。
• 识别并分析实现模型并行性过程中的关键技术与架构挑战，包括通信瓶颈和算子图划分不优等问题。
• 研究模型并行性在现代大规模语言模型（LLMs）中的实际应用，特别是多亿参数级Transformer模型。
• 评估不同并行策略在内存使用、执行时间和设备利用率方面的权衡。
• 倡导深度神经网络自动并行化领域的标准化，以实现可复现的基准测试，并推动该领域更广泛的应用。

提出的方法

• 本研究采用两阶段雪球式文献综述方法：首先分析深度神经网络自动并行化技术以构建理论模型，其次研究真实世界的大规模语言模型部署案例。
• 通过算子图对模型并行性进行形式化表达，其中并行性沿两个维度定义：算子内（单个操作内部）和算子间（多个操作之间）。
• 通过基于设备特性、张量大小和通信开销的执行模拟器，评估如FlexFlow和TensorOpt等自动并行化框架的性能。
• 在FlexFlow中采用马尔可夫链蒙特卡洛搜索策略，以模拟性能指标为指导，探索并行化搜索空间。
• 通过收集并筛选包含模型并行性实现细节的论文，分析真实世界的大规模语言模型（如Megatron、Gopher、PaLM、GPT）。
• 分析内容包括硬件利用率、通信模式，以及专用硬件（如TPUv4、A100）在实现高效算子内并行性中的作用。

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

实验结果

研究问题

• RQ1 存在哪些类型的模型并行性？它们在算子图中的定义方式是什么？
• RQ2 在实现模型并行性时，尤其是在分布式系统中，主要的技术与架构挑战是什么？
• RQ3 模型并行性在现代大规模Transformer模型（如Megatron、PaLM和Gopher）中是如何应用的？
• RQ4 在通信、内存和设备利用率方面，算子内并行性与算子间并行性之间的权衡是什么？
• RQ5 如何通过表示和基准的标准化来改进深度神经网络自动并行化领域？

主要发现

• 模型并行性主要分为算子内（单个操作内部）和算子间（多个操作之间）两类，实践中常见的是结合两者的混合策略。
• 算子内并行性由于频繁在设备间进行数据重排，通信开销极高，尤其是在以太网等较慢互连网络上。
• 算子间并行性常因负载不均和计算模式不规则，导致训练期间设备利用率低下。
• 将算子图在设备间最优划分是一个重大挑战，因为模型图结构通常与底层设备拓扑不匹配。
• 现代大模型如PaLM利用专用硬件（如TPUv4）来缓解算子内通信开销，从而在不依赖算子间划分的情况下实现高效的模型并行性。
• Megatron和Gopher等模型通过在节点内采用算子内并行性，结合节点间算子间并行性与数据并行性，突破了单节点的扩展限制。

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