机器学习系统

本书系统地介绍了机器学习系统的设计原则和实践经验，侧重于介绍机器学习的原理、神经网络和优化器、自动差分算法、机器学习系统编程模型、控制流和数据流，异构硬件加速器的原理和编程、数据流图编译器前端、数据流图编译器后端、数据准备和增强、模型部署相关技术、分布式训练、弹性训练、联合训练和评估平台、调试和优化工具、数据隐私和安全等。在讲授的过程中，本书将根据MindSpore的自身特点，在各个章节突出讨论MindSpore的优势点，从而将MindSpore并列为与TensorFlow，PyTorch的三大框架。

麦络 爱丁堡大学信息学院助理教授，博士生导师。2018年于帝国理工学院获得博士学位，谷歌博士奖学金获得者。主要研究方向为分布式系统、机器学习和数据管理，当前研究工作专注于构建大规模、自适应和可信的机器学习系统，受到谷歌、微软、华为、腾讯和阿里巴巴等多家知名科技公司的资助。在计算机顶级会议OSDI、NSDI、USENIX ATC、CoNEXT、VLDB、ECCV和NeurIPS发表多篇论文。获得CoNEXT会议最佳论文入围奖，ACM Multimedia 2017最佳开源论文奖。

董豪 北京大学计算机学院助理教授，博士生导师，2019年于帝国理工获得博士学位。主要研究方向为计算机视觉、机器人和具身智能，当前研究工作围绕智能机器人的自主决策与泛化交互。担任CVPR 2023领域主席、AAAI 2023高级程序委员、中国科技核心期刊Machine Intelligence Research副编委等，在NeurIPS、ICLR、ICCV、ECCV、IROS等顶级国际会议和期刊发表30余篇论文，Deep Reinforcement Learning：Fundamentals, Research and Applications作者。获得ACM MM最佳开源软件奖，新一代人工智能产业技术创新战略联盟 OpenI 启智社区优秀开源项目、Springer Nature中国作者高影响力研究精选等。

本书分为三大部分。基础篇覆盖机器学习框架使用者所需要了解的核心系统知识和相关编程案例；进阶篇覆盖了机器学习框架开发着所需要理解的核心知识和相关实践案例；拓展篇详细讨论了多种类的机器学习系统，从而为广大机器学习从业者提供解密底层系统所需的基础知识。

缘起

我在2020年来到了爱丁堡大学信息学院，爱丁堡大学是AI（Artificial Intelligence，人工智能）研究的发源地之一，很多学生慕名而来学习机器学习技术。爱丁堡大学拥有许多出色的机器学习课程（如自然语言处理、计算机视觉、计算神经学等），同时也拥有一系列关于计算机系统的基础课程（如操作系统、编程语言、编译器、计算机体系架构等）。但是当我在教学的过程中问起学生：机器学习是如何利用计算机系统实现计算加速和部署的？许多学生投来疑惑的眼神。这促使我思考在爱丁堡大学乃至于其他世界顶尖大学的教学大纲中，是不是缺一门衔接机器学习和计算机系统的课程。

我的第一反应是基于一门已有的课程进行拓展。那时，加州大学伯克利分校的AI Systems（人工智能系统）课程较为知名。这门课描述了机器学习系统的不同研究方向，以研读论文为主进行教学。可惜的是，许多论文已经无法经受住时间的检验。更重要的是，这门课缺乏对于知识的整体梳理，未能形成完整的知识体系架构。学习完这门课程，学生未能形成对于从头搭建机器学习系统的明确思路。我调研了其他学校，华盛顿大学曾短期开过Deep Learning Systems（深度学习系统）课程，这门课程讲述了机器学习程序的编译过程。而这门课程以讲述Apache TVM（Tensor Virtual Machine，张量虚拟机）深度学习编译器为主要目的，对于机器学习系统缺乏完整的教学。另外，斯坦福大学的课程Machine Learning Systems Design（机器学习系统设计），课程设计人的研究领域以数据库为主，因此该课程专注于数据清洗、数据管理、数据标注等主题。

当时觉得比较合适的是微软亚洲研究院的AI Systems课程。这门课程讲述了机器学习系统背后的设计理念。但是当我准备将其教授给本科生的时候，我发现这门课对于机器学习系统核心设计理念讲解得很浅，同时这门课要求学生有大量计算机系统的背景知识，实际上它更适合博士生。上述课程的共同问题是：其课程教学方式都以研读相关论文为主，因此教授的内容都是高深和零散的，而不是通俗易懂，知识脉络清晰的，这给学生学习机器学习系统造成了极大的困难。

回首2020年，我们已经拥有了优秀的操作系统、数据库、分布式系统等基础性教材。同时，在机器学习相关算法方面也有了一系列教材。然而，无论在国内外，我很难找到一本系统性讲述机器学习系统的教材。因此，许多公司和高校实验室不得不花费大量的人力和物力从头培养学生和工程师，使他们加强对于机器学习底层基础设施的认识。这类教材的缺乏已经制约了高校的人才培养，不利于高校培养出符合业界与学界和时代发展需求的人才了。因此，我开始思考：我们是不是应该推出一本有关机器学习系统的教科书呢？

开端

带着写书的构想，我开始和朋友沟通。大家都非常认可编写这类书的巨大价值，但是现实的情况是：很少有人愿意做这么一件费力的事情。我当时的博士后导师也劝我：你现在处在教师职业生涯的初期，追求高影响力的学术论文是当务之急，写一本书要耗费大量的时间和精力，最后可能也无法出版面世。而我和同行交流时也发现：他们更愿意改进市面上已经有的教科书，即做有迹可循的事情，而不是摸着石头过河，做从无到有的事情。特别是对于机器学习系统这个快速发展、频繁试错的领域，能不能写出经受住时间检验的书也是一个未知数。

考虑到写作的巨大挑战，我将写书的构想藏于心中，直到一次回国和MindSpore的架构师金雪锋聊天。和雪锋的相识大约是在2019年的圣诞节，雪锋来伦敦访问，他正在领导MindSpore的开发（当时MindSpore 1.0还没有发布）。而对于机器学习系统的开发，我也有很深的兴趣。我在2018年也和好友一起从头搭建一个机器学习框架（类似于PyTorch），虽然最终资源不足无疾而终，不过许多的思考成就了我之后发表的多篇机器学习系统论文。和雪锋聊起来，我们都对AI系统开发之难深有同感。我们共同的感慨就是：找到懂机器学习系统开发的人太难了。现今的学生都一心学习机器学习算法，很多学生对于底层系统的运作原理理解得很浅，也不够重视。而当他们在实际中应用机器学习技术时才意识到系统的重要性，那时想去学习，却没有了充沛的学习时间。我对雪锋苦笑道：“我是准备写一本机器学习系统教材的，但是可能还要等三四年才能完成。”雪锋说：“我也有这个想法啊。你要是写的话，我能帮助到你吗？”

这句话点醒了我。传统的图书写作，往往依赖于一两个教授将学科十余年的发展慢慢总结整理成书。这种模式类似于传统软件开发的瀑布流方式。可是，在科技世界，这已经变了！软件的发展从传统的瀑布流进化到如今的开源敏捷开发。而图书的写作为什么还要停留在传统方式呢？MXNet开源社区编写的专注于深度学习算法的图书Deep Dive into Deep Learning就是一个很好的例子啊。我因此马上找到当年一起创立TensorLayer开源社区的小伙伴——北京大学的董豪。我们一拍即合，说干就干。雪锋也很高兴我和董豪愿意开始做这件事，也邀请了他的同事干志良来帮忙。我们终于开始图书的写作了！

经过几轮的讨论，我们将书名定为《机器学习系统：设计和实践》。我们希望通过教给学生机器学习系统设计原理，同时也为学生提供大量的系统实现经验分享，让他们在将来工作和科研中遇到实际问题时知道该如何分析和解决。

社区的构建

考虑到机器学习系统本身就是一个不断发展并且孕育细分领域的学科。我从一开始就在思考：如何设计一个可扩展（Scalable）的社区架构保证这本书的可持续发展呢？因为我专注于大规模软件系统，故决定借鉴几个分布式系统的设计要点构建社区。

（1）预防单点故障和瓶颈: 现代分布式系统往往采用控制层和数据层分离的设计避免单点故障和瓶颈。那么我们在设计高度可扩展的写作社区的时候也要如此。因此，我们设计了如下分布式机制：编辑决定花最多的时间来寻找优秀的、主动的、负责任的书稿章节负责人。章节负责人可以进一步寻找其他作者共同协作。章节负责人和章节作者进行密切的沟通，按照给定时间节点，全速异步推进。编辑和章节负责人设定了每周讨论同步写作的进展，确保并行完成的章节内容质量能够持续符合编辑和社区的整体预期。

（2）迭代式改进: 深度学习的优化算法------随机梯度下降本质上是在复杂问题中利用局部梯度进行海量迭代，最终找到局部最优解。因此我们利用了同样的思路设计图书质量的迭代提高。我们首先在Overleaf上写作好书籍的初版（类似于初始参数）。接下来，将图书的内容做成标准的Git代码仓库。建立机制鼓励开源社区和广大读者开启GitHub问题（Issue）和拉取请求（Pull Request，PR），持续改进图书质量，而我们设置好完善的图书构建工具、持续集成工具、贡献者讨论会等，就可以让图书的质量持续提高，实现与随机梯度下降（Stochastic Gradient Descent）一样的最优结果。

（3）高可用性: 构建7×24小时在线的写作平台，让图书参与者可以在全球任何时区、任何语言平台下都能参与开发图书，倾听社区的反馈。因此将Git仓库放置在GitHub上，并准备之后在Gitee做好镜像。这样，就搭建了一套高可用的写作平台了。

（4）内容中立: 一个分布式系统要能长久运行，其中的每一个节点都要同等对待，遇到故障才能用统一的办法进行故障恢复。考虑到图书写作中的故障（设计无法经受时间检验，写作人中途不得不退出等）可能来源于方方面面，我们让不同背景的参与者共同完成每一个章节，确保写出中立、客观、包容的内容，并且写作不会因为故障而中断。

现状和未来

机制一旦建立好，写作就自动化地运行起来了，参与者也越来越多，我带过的学生袁秀龙、丁子涵、符尧、任杰、梁文腾也很用心参与编写，董豪邀请了鹏城实验室的韩佳容和赖铖，志良邀请了许多MindSpore的小伙伴进来做贡献，许多资深的机器学习系统设计者也和我们在各个渠道展开讨论，提供了非常多宝贵的写作建议。另外，学界和产业界的反响也很热烈。海外很多优秀的学生（斯坦福大学的孙建凯、卡耐基--梅隆大学的廖培元、剑桥大学的王瀚宸、爱丁堡大学的穆沛），产业界的朋友（英国葛兰素史克公司机器学习团队的肖凯严）都加入了我们的写作。同时，学界的教授（英国伦敦帝国理工学院的Peter Pietzuch教授、香港科技大学的陈雷教授等）也持续给我们提供了写作意见，改进了图书质量。

充分发动了“分布式系统”的力量后，图书的内容得以持续高质量地添加。当我们开源了图书以后，图书的受众快速增长，GitHub上关注度的增长让我们受宠若惊。在社区的推动下，图书的中文版、英文版、阿拉伯语版都已经开始推进。这么多年来，我第一次意识到我在分布式系统和机器学习中学习到的知识，在解决现实复杂问题的时候是如此的有用！

很多时候，当我们面对未知而巨大的困难时，个人的力量真的很渺小。而和朋友、社区一起就变成了强大的力量，让我们鼓起勇气，走出了最关键的第一步！希望我的一些思考，能给其他复杂问题的求解带来一些小小的启发。

截至2022年5月，本书的核心作者和编辑包括麦络、董豪、金雪锋和干志良，谭志鹏策划了本书。本书各章节由以下作者参与编写：导论（麦络、董豪、干志良）、编程模型（赖铖、麦络、董豪）、计算图（韩佳容、麦络、董豪）、AI编译器和前端技术（梁志博、张清华、黄炳坚、余坚峰、干志良）、AI编译器后端和运行时（褚金锦、穆沛、蔡福璧）、硬件加速器（张任伟、任杰、梁文腾、刘超、陈钢、黎明奇）、数据处理（袁秀龙）、模型部署（韩刚强、唐业辉、翟智强、李姗妮）、分布式训练（麦络、廖培元）、联邦学习系统（吴天诚、王瀚宸）、推荐系统（符尧、裴贝、麦络）、强化学习系统（丁子涵）、可解释AI系统（李昊阳、李小慧）、机器人系统（孙建凯、肖凯严）。

最后，我们非常欢迎新成员的加入以扩展内容并提升图书质量，感兴趣的读者可以通过图书的GitHub社区1 联系我们。我们非常期待和大家一起努力，编写出一本推动业界发展的机器学习系统图书！

麦络

英国爱丁堡

2023年2月

1 可参考网址为: \url{https://github.com/openmlsys/。

基础篇

第1章 导论 003

1.1 机器学习应用 003

1.2 机器学习框架的设计目标 004

1.3 机器学习框架的基本组成原理 005

1.4 机器学习系统生态 006

1.5 本书结构和读者对象 007

第2章 编程模型 009

2.1 机器学习系统编程模型的演进 009

2.2 机器学习工作流 011

2.2.1 环境配置 011

2.2.2 数据处理 012

2.2.3 模型定义 013

2.2.4 损失函数和优化器 014

2.2.5 训练及保存模型 015

2.2.6 测试和验证 016

2.3 定义深度神经网络 017

2.3.1 以层为核心定义神经网络 017

2.3.2 神经网络层的实现原理 021

2.3.3 自定义神经网络层 022

2.3.4 自定义神经网络模型 023

2.4 C/C++编程接口 024

2.4.1 在Python中调用C/C++函数的原理 025

2.4.2 添加C++编写的自定义算子 025

2.5 机器学习框架的编程范式 030

2.5.1 机器学习框架编程需求 030

2.5.2 机器学习框架编程范式现状 030

2.5.3 函数式编程案例 031

2.6 总结 032

2.7 拓展阅读 032

第3章 计算图 033

3.1 设计背景和作用 033

3.2 计算图的基本构成 034

3.2.1 张量和算子 035

3.2.2 计算依赖 037

3.2.3 控制流 038

3.2.4 基于链式法则计算梯度 041

3.3 计算图的生成 043

3.3.1 静态生成 043

3.3.2 动态生成 046

3.3.3 动态图和静态图生成的比较 048

3.3.4 动态图与静态图的转换和融合 049

3.4 计算图的调度 051

3.4.1 算子调度执行 051

3.4.2 串行与并行 052

3.4.3 数据载入同步与异步机制 053

3.5 总结 054

3.6 拓展阅读 055

进阶篇

第4章 AI 编译器和前端技术 059

4.1 AI编译器设计原理 059

4.2 AI编译器前端技术概述 061

4.3 中间表示 062

4.3.1 中间表示的基本概念 062

4.3.2 中间表示的种类 063

4.3.3 机器学习框架的中间表示 065

4.4 自动微分 072

4.4.1 自动微分的基本概念 072

4.4.2 前向与反向自动微分 074

4.4.3 自动微分的实现 077

4.5 类型系统和静态分析 081

4.5.1 类型系统概述 081

4.5.2 静态分析概述 082

4.6 常见前端编译优化方法 083

4.6.1 前端编译优化简介 083

4.6.2 常见编译优化方法介绍及实现 083

4.7 总结 085

第5章 AI 编译器后端和运行时 086

5.1 概述 086

5.2 计算图优化 088

5.2.1 通用硬件优化 088

5.2.2 特定硬件优化 090

5.3 算子选择 091

5.3.1 算子选择的基础概念 091

5.3.2 算子选择的过程 095

5.4 内存分配 095

5.4.1 Device内存概念 096

5.4.2 内存分配 096

5.4.3 内存复用 098

5.4.4 常见的内存分配优化手段 099

5.5 计算调度与执行 101

5.5.1 单算子调度 101

5.5.2 计算图调度 102

5.5.3 交互式执行 106

5.5.4 下沉式执行 110

5.6 算子编译器 110

5.6.1 算子调度策略 111

5.6.2 子策略组合优化 112

5.6.3 调度空间算法优化 114

5.6.4 芯片指令集适配 115

5.6.5 算子表达能力 116

5.6.6 相关编译优化技术 117

5.7 总结 117

5.8 拓展阅读 118

第6章 硬件加速器 119

6.1 概述 119

6.1.1 硬件加速器设计的意义 119

6.1.2 硬件加速器设计的思路 119

6.2 硬件加速器基本组成原理 120

6.2.1 硬件加速器的架构 120

6.2.2 硬件加速器的存储单元 121

6.2.3 硬件加速器的计算单元 122

6.2.4 DSA芯片架构 124

6.3 加速器基本编程原理 125

6.3.1 硬件加速器的可编程性 125

6.3.2 硬件加速器的多样化编程方法 128

6.4 加速器实践 132

6.4.1 环境 132

6.4.2 广义矩阵乘法的朴素实现 133

6.4.3 提高计算强度 135

6.4.4 使用共享内存缓存复用数据 138

6.4.5 减少寄存器使用 139

6.4.6 隐藏共享内存读取延迟 140

6.4.7 隐藏全局内存读取延迟 141

6.4.8 与cuBLAS对比 142

6.4.9 小结 143

6.5 总结 143

6.6 拓展阅读 144

第7章 数据处理 145

7.1 概述 146

7.1.1 易用性 146

7.1.2 高效性 147

7.1.3 保序性 147

7.2 易用性设计 147

7.2.1 编程抽象与接口 147

7.2.2 自定义算子支持 151

7.3 高效性设计 153

7.3.1 数据读取的高效性 154

7.3.2 数据计算的高效性 157

7.4 保序性设计 162

7.5 单机数据处理性能的扩展 163

7.5.1 基于异构计算的数据预处理 163

7.5.2 基于分布式的数据预处理 165

7.6 总结 166

第8章 模型部署 168

8.1 概述 168

8.2 训练模型到推理模型的转换及优化 169

8.2.1 模型转换 169

8.2.2 算子融合 170

8.2.3 算子替换 172

8.2.4 算子重排 173

8.3 模型压缩 173

8.3.1 量化 174

8.3.2 模型稀疏 176

8.3.3 知识蒸馏 178

8.4 模型推理 179

8.4.1 前处理与后处理 179

8.4.2 并行计算 180

8.4.3 算子优化 181

8.5 模型的安全保护 186

8.5.1 概述 186

8.5.2 模型混淆 186

8.6 总结 188

8.7 拓展阅读 189

第9章 分布式训练 190

9.1 设计概述 190

9.1.1 设计动机 190

9.1.2 系统架构 191

9.1.3 用户益处 192

9.2 实现方法 192

9.2.1 方法分类 192

9.2.2 数据并行 194

9.2.3 模型并行 194

9.2.4 混合并行 197

9.3 流水线并行 197

9.4 机器学习集群架构 198

9.5 集合通信 200

9.5.1 常见集合通信算子 200

9.5.2 基于AllReduce的梯度平均算法 203

9.5.3 集合通信算法性能分析 205

9.5.4 利用集合通信优化模型训练的实践 206

9.5.5 集合通信在数据并行的实践 207

9.5.6 集合通信在混合并行的实践 208

9.6 参数服务器 210

9.6.1 系统架构 210

9.6.2 异步训练 211

9.6.3 数据副本 212

9.7 总结 212

9.8 拓展阅读 213

拓展篇

第10章 联邦学习系统 217

10.1 概述 217

10.1.1 定义 217

10.1.2 应用场景 217

10.1.3 部署场景 218

10.1.4 常用框架 218

10.2 横向联邦学习 219

10.2.1 云云场景中的横向联邦 219

10.2.2 端云场景中的横向联邦 220

10.3 纵向联邦学习 222

10.3.1 纵向联邦架构 222

10.3.2 样本对齐 223

10.3.3 联合训练 224

10.4 隐私加密算法 225

10.4.1 基于LDP算法的安全聚合 226

10.4.2 基于MPC算法的安全聚合 226

10.4.3 基于LDP-SignDS算法的安全聚合 227

10.5 展望 229

10.5.1 异构场景下的联邦学习 229

10.5.2 通信效率提升 230

10.5.3 联邦生态 230

10.6 总结 231