前言

GPT的启示

ChatGPT石破天惊，GPT-4的问世又引发了进一步的轰动，GPT-5即将到来……它们的影响远远超出了大家的预期和想象。有观点认为，ChatGPT是通用人工智能的Singularity（奇点）。无独有偶，2022年，微软发布了一篇论文Singularity：Planet-Scale，Preemptive and Elastic Scheduling of AI Workloads，介绍其全局分布式调度服务，作者包括Azure的CTO Mark Russinovich。而Azure与OpenAI合作，重新设计了超级计算机，在Azure云上为OpenAI训练超大规模的模型。该论文将Azure的全局分布式调度服务命名为Singularity，可见其深意。

GPT-3模型的参数量是1750亿个，研发这个规模的大型模型，是一个极其复杂的系统工程，涵盖了算法、算力、网络、存储、大数据、框架、基础设施等多个领域。微软2020年发布的信息称，其计划为OpenAI开发的专用超级计算机包括28.5万个CPU、1万个GPU。市场调查机构TrendForce的报告则指出，ChatGPT需要3万个GPU。在OpenAI官网和报告中都提到，GPT-4项目的重点之一是开发一套可预测、可扩展的深度学习栈和基础设施（Infrastructure）。与之对应的是，在OpenAI研发团队的6个小组中，有5个小组的工作涉及AI工程和基础设施。

OpenAI没有提供GPT-4的架构（包括模型大小）、硬件、数据集、训练方法等内容，这非常令人遗憾，但是我们可以从微软发布的论文入手，来研究GPT-4这座冰山在水下的那些深层技术栈。从论文可以看出，GPT使用的底层技术并没有那么“新”，而是在现有技术基础之上进行的深度打磨，并从不同角度对现有技术进行了拓展，做到工程上的极致。比如Singularity在GPU调度方面，就有阿里巴巴AntMan的影子。再比如Singularity从系统角度出发，使用CRIU完成任务抢占、迁移的同时，也巧妙解决了弹性训练的精度一致性问题。

AI的黄金时代可能才刚刚开启，各行各业的生产力革命也会相继产生。诚然，OpenAI已经占据了领先位置，但是接下来的AI赛道会风起云涌，中国企业势必会在其中扮演极其重要的角色，也会在深度学习栈和基础设施领域奋起直追。然而，“弯道超车”需要建立在技术沉淀和产品实力之上，我们只有切实地扎根于现有的分布式机器学习技术体系，并对其进行深耕，才能为更好的创新和发展打下基础。大家都已在路上，没有人直接掌握着通向未来的密码，但面对不可阻挡的深层次的信息革命和无限的发展机遇，我们必须有所准备。

复杂模型的挑战

为了降低在大型数据集上训练大型模型的计算成本，研究人员早已转向使用分布式计算体系结构（在此体系结构中，许多机器协同执行机器学习方法）。人们创建了算法和系统，以便在多个CPU或GPU上并行化机器学习（Machine Learning，ML）程序（多设备并行），或者在网络上的多个计算节点并行化机器学习训练（分布式并行）。这些软件系统利用最大似然理论的特性来实现加速，并随着设备数量的增加而扩展。理想情况下，这样的并行化机器学习系统可以通过减少训练时间来实现大型模型在大型数据集上的并行化，从而实现更快的开发或研究迭代。

然而，随着机器学习模型在结构上变得更加复杂，对于大多数机器学习研究人员和开发人员来说，即便使用TensorFlow、PyTorch、Horovod、Megatron、DeepSpeed等工具，编写高效的并行化机器学习程序仍然是一项艰巨的任务，使用者需要考虑的因素太多，比如：

• 系统方面。现有系统大多是围绕单个并行化技术或优化方案来构建的，对组合多种策略的探索不足，比如不能完全支持各种混合并行。

• 性能方面。不同并行化策略在面对不同的算子时，性能差异明显。有些框架未考虑集群物理拓扑（集群之中各个设备的算力、内存、带宽有时会存在层级的差距），使用者需要依据模型特点和集群网络拓扑来选择或者调整并行化策略。

• 易用性方面。很多系统需要开发人员改写现有模型，进行手动控制，比如添加通信原语，控制流水线等。不同框架之间彼此割裂，难以在不同并行策略之间迁移。

• 可用性方面。随着训练规模扩大，硬件薄弱或设计原因会导致单点故障概率随之增加，如何解决这些痛点是个严重问题。

总之，在将分布式并行训练系统应用于复杂模型时，“开箱即用”通常会导致低于预期的性能，求解最优并行策略成为一个复杂度极高的问题。为了解决这个问题，研究人员需要对分布式系统、编程模型、机器学习理论及其复杂的相互作用有深入的了解。

本书是笔者在分布式机器学习领域学习和应用过程中的总结和思考，期望能起到抛砖引玉的作用，带领大家走入/熟悉分布式机器学习这个领域。

本书的内容组织

PyTorch是大家最常用的深度学习框架之一，学好PyTorch可以很容易地进入分布式机器学习领域，所以全书以PyTorch为纲进行穿插讲解，从系统和实践的角度对分布式机器学习进行梳理。本书架构如下。

• 第1篇 分布式基础

本篇首先介绍了分布式机器学习的基本概念、基础设施，以及机器学习并行化的技术、框架和软件系统，然后对集合通信和参数服务器PS-Lite进行了介绍。

• 第2篇 数据并行

数据并行（Data Parallelism）是深度学习中最常见的技术。数据并行的目的是解决计算墙，将计算负载切分到多张卡上。数据并行具有几个明显的优势，包括计算效率高和工作量小，这使得它在高计算通信比的模型上运行良好。本篇以PyTorch和Horovod为主对数据并行进行分析。

• 第3篇 流水线并行（Pipeline Parallelism）

当一个节点无法存下整个神经网络模型时，就需要对模型进行切分，让不同设备负责计算图的不同部分，这就是模型并行（Model Parallelism）。从计算图角度看，模型并行主要有两种切分方式：层内切分和层间切分，分别对应了层内模型并行和层间模型并行这两种并行方式。业界把这两种并行方式分别叫作张量模型并行（简称为张量并行，即Tensor Parallelism）和流水线模型并行（简称为流水线并行，即Pipeline Parallelism）。

张量模型并行可以把较大参数切分到多个设备，但是对通信要求较高，计算效率较低，不适合超大模型。在流水线模型并行中，除了对模型进行层间切分外，还引入了额外的流水线来隐藏通信时间、充分利用设备算力，从而提高了计算效率，更合适超大模型。

因为流水线并行的独特性和重要性，所以对这部分内容单独介绍。本篇以GPipe、PyTorch、PipeDream为例来分析流水线并行。

• 第4篇 模型并行

目前已有的深度学习框架大都提供了对数据并行的原生支持，虽然对模型并行的支持还不完善，但是各个框架都有自己的特色，可以说百花齐放，百家争鸣。本篇介绍模型并行，首先会对NVIDIA Megatron进行分析，讲解如何进行层内分割模型并行，然后学习PyTorch如何支持模型并行。

• 第5篇 TensorFlow分布式

本篇学习TensorFlow如何进行分布式训练。迄今为止，在分布式机器学习这一系列分析之中，我们大多以PyTorch为纲，结合其他框架/库来穿插完成。但是缺少了TensorFlow就会觉得整个世界（系列）都是不完美的，不仅因为TensorFlow本身有很强大的影响力，更因为TensorFlow分布式博大精深，特色鲜明，对于技术爱好者来说是一个巨大宝藏。

本书面向的读者

本书读者群包括：

• 机器学习领域内实际遇到大数据、分布式问题的人，不但可以参考具体解决方案，也可以学习各种技术背后的理念、设计哲学和发展过程。

• 机器学习领域的新人，可以按图索骥，了解各种框架如何使用。

• 其他领域（尤其是大数据领域和云计算领域）想转入机器学习领域的工程师。

• 有好奇心，喜欢研究框架背后机理的学生，本书也适合作为机器学习相关课程的参考书籍。

如何阅读本书

本书源自笔者的博客文章，总体来说是按照项目解决方案进行组织的，每一篇都是关于某一特定主题的方案集合。大多数方案自成一体，每个独立章节中的内容都是按照循序渐进的方式来组织的。

行文

• 本书以神经网络为主，兼顾传统机器学习，所以举例往往以深度学习为主。

• 因为本书内容来源于多种框架/论文，这些来源都有自己完整的体系结构和逻辑，所以本书会存在某一个概念或者问题以不同角度在前后几章都论述的情况。

• 解析时会删除非主体代码，比如异常处理代码、某些分支的非关键代码、输入的检测代码等。也会省略不重要的函数参数。

• 一般来说，对于类定义只会给出其主要成员变量，某些重要成员函数会在使用时再进行介绍。

• 本书在描述类之间关系和函数调用流程上使用了UML类图和序列图，但是因为UML规范过于繁杂，所以本书没有完全遵循其规范。对于图例，如果某图只有细实线，则可以根据箭头区分是调用关系还是数据结构之间的关系。如果某图存在多种线条，则细实线表示数据结构之间的关系，粗实线表示调用流程，虚线表示数据流，虚线框表示列表数据结构。

版本

各个框架发展很快，在本书写作过程中，笔者往往会针对某一个框架的多个版本进行研读，具体框架版本对应如下。

• PyTorch：主要参考版本是1.9.0。

• TensorFlow：主要参考版本是2.6.2。

• PS-Lite：master版本。

• Megatron：主要参考版本是2.5。

• GPipe：master版本。

• PipeDream：master版本。

• torchgpipe：主要参考版本是0.0.7。

• Horovod：主要参考版本是0.22.1。

深入

在本书（包括博客）的写作过程中，笔者参考和学习了大量论文、博客和讲座视频，在此对这些作者表示深深的感谢。具体参考资料和链接请扫描封底二维码获取。如果读者想继续深入研究，除论文、文档、原作者博客和源码之外，笔者有如下建议：

• PyTorch：推荐OpenMMLab＠知乎，Gemfield＠知乎。Gemfield是PyTorch的万花筒。

• TensorFlow：推荐西门宇少（DeepLearningStack＠cnblogs）、刘光聪（horance-liu＠github）。西门宇少兼顾深度、广度和业界前沿。刘光聪的电子书《TensorFlow内核剖析》是同领域最佳，本书借鉴颇多。

• Megatron：推荐迷途小书僮＠知乎，其对Megatron有非常精彩的解读。

• 总体：推荐张昊＠知乎、OneFlow＠知乎，既高屋建瓴，又紧扣实际。

• 刘铁岩、陈薇、王太峰、高飞的《分布式机器学习：算法、理论与实践》非常经典，强烈推荐。

因为时间和精力原因，笔者没能对国内的深度学习框架进行分析，非常遗憾。对于有兴趣的读者，推荐以下优秀框架：一流科技OneFlow，华为MindSpore，旷视天元。也强烈推荐这几个框架的代表人物：袁进辉＠知乎、金雪锋＠知乎、许欣然＠知乎。每次读他们的文章，都会让我对机器学习领域有更深入的认识。另外要专门感谢李舒辰博士，他在机器学习方面具备深厚功力，和他的讨论总是让我受益匪浅。

致谢

首先，感谢我生命中遇到的各位良师：许玉娣老师、刘健老师、邹艳聘老师、王凤珍老师、栾锡宝老师、王金海老师、童若锋老师、唐敏老师、赵慧芳老师，董金祥老师……师恩难忘。童若锋老师是我读本科时的班主任，又和唐敏老师一起在我攻读硕士学位期间对我进行悉心指导。那时童老师和唐老师刚刚博士毕业，两位老师亦师亦友，他们的言传身教让我受益终生。

感谢我的编辑黄爱萍在本书出版过程中给我的帮助。对我来说，写博客是快乐的，因为我喜欢技术，喜欢研究事物背后的机理。整理出书则是痛苦的，其难度远远超出了预期，从整理到完稿用了一年多时间。没有编辑的理解和支持，这本书很难问世。另外，因为篇幅所限，笔者博客中的很多内容（比如DeepSpeed、弹性训练、通信优化、数据处理等）未能在书中体现，甚是遗憾。

感谢童老师、孙力哥、媛媛姐、文峰同学，以及袁进辉、李永（九丰）两位大神在百忙之中为本书写推荐语，谢谢你们的鼓励和支持。

最后，特别感谢我的爱人和孩子们，因为写博客和整理书稿，我牺牲了大量本应该陪伴她们的时间，谢谢她们给我的支持和包容。也感谢我的父母和岳父母帮我们照顾孩子，让我能够长时间在电脑前面忙忙碌碌。

本书资源下载

扫描本书封底二维码，可以获取本书的参考资料、链接和代码。

由于笔者水平和精力都有限，而且本书的内容较多、牵涉的技术较广，谬误和疏漏之处在所难免，很多技术点设计的细节描述得不够详尽，恳请广大技术专家和读者指正。可以将意见和建议发送到我的个人邮箱RossiLH＠163.com，或者通过博客园、CSDN、掘金或微信公众号搜索“罗西的思考”与我进行交流和资料获取。我也将密切跟踪分布式机器学习技术的发展，吸取大家意见，适时编写本书的升级版本。

柳浩

2023年5月