深度学习框架机制及分布式并行
内容来自B站ZOMI大佬,做个文字版笔记,有兴趣可以观看ZOMI大佬的原视频。
自动微分
微分的分类
微分主要分为三种方式:符号微分、数值微分和自动微分。
- 符号微分:通过求导法则指定表达式变换规则;
- 数值微分:使用有限差分进行近似;
- 自动微分:所有数值计算都由有限的基本运算组成,基本运算的导数表达式是已知的,通过链式法则将数值计算各部分组合成整体。
考虑以下公式。
符号微分
符号微分需要先手工计算微分公式。
公式过长,省略一部分。
再将上述公式编写成代码。
1 |
|
符号微分的缺点显而易见,表达式会严重的膨胀,但计算的数值会比较精确。
通过符号微分无法实现自动微分。
数值微分
数值微分是利用差分方式来求微分。
1 |
|
数值微分虽然容易实现,但计算不精确,计算的复杂度也较高,同时对
自动微分
自动微分则是令计算机求解。
1 |
|
自动微分数值精度高,表达式也不会无限膨胀,但需要存储中间求导结果,占用大量内存资源。
自动微分模式
考虑下列公式。
正常的正向计算过程如下。
Forward Primal Trace | Result |
---|---|
计算机会将原函数转换成一个有向无环图 (DAG)。
前向微分 (Forward Model)
会将每个计算步骤都进行求导,并保留每一步的值。下面是对
Forward Tangent (Derivative) Trace | Result |
---|---|
前向计算过程中,会假设当前计算的输入
后向微分 (Backward Model)
后向微分就是沿着DAG的反方向,利用链式求导法则计算。
Reverse Adjoint (Derivative) Trace | Result |
---|---|
反向计算过程中,会假设
雅可比矩阵
对于函数
前向微分可以在一次计算中通过链式法则得到Jacobian矩阵中的一列,反向微分则是得到Jacobian矩阵的一行。
对于前向微分,设
前向模式的问题在于,每次前向过程只能计算对一个自变量的偏导数,若输入变量数量远大于输出变量数量的话,前向模式就太低效了
对于反向微分,设
自动微分的实现方式
基本表达式法 (Library, LIB)
封装基本的表达式及其微分表达式作为库函数,运行时记录基本表达式和相应的组合关系,利用链式法则对基本表达式的微分结果进行组合。(以autodiff为代表)
操作符重载法 (Operator Overload, OO)
利用语言多态特性,使用操作符重载基本运算表达式,运行时记录基本表达式和相应的组合关系,利用链式法则对基本表达式的微分结果进行组合。(以PyTorch为代表)
源码转换法 (AST)
是语言预处理器、编译器或解释器的扩展,对程序表达进行分析得到基本表达式的组合关系,利用链式法则对基本表达式的微分结果进行组合。(以MindSpore为代表)
实现方式 | 优点 | 缺点 |
---|---|---|
LIB | 实现简单;适用任意编程语言 | 使用库函数编程;无法使用原语言运算表达式 |
OO | 实现简单;语言具备多态性;易用性高,贴合原生语言 | 显式构造Tape数据结构和对Tape的读写;额外数据结构和操作的引入,不利于高阶微分;if ,while 等控制流表达式难以通过操作符重载 |
AST | 丰富数据类型和语言原生操作;无额外数据结构,易于实现高阶微分;微分结果以代码形式存在,方便分布式系统计算 | 代码难以理解,涉及底层编译原理;实现复杂度高,需要扩展语言预处理器、编译器或解释器;容易写出不支持的代码导致错误,需要更强的检查告警系统 |
AI框架
编程范式
命令式编程
Imperative programming,又称define-by-run,以动态图为代表。
前端语言直接驱动后端算子执行,用力表达式会被立即求值。以PyTorch为代表。
命令式编程方便调试,灵活性高。但缺少对算法的统一描述,缺乏编译期的优化。
声明式编程
Declarative programming,又称define-and-run,以静态图为代表。
前端语言中的表达式不直接执行,首先会构建一个完整前向计算过程表示,然后对数据流进行优化后再执行。以TensorFlow为代表。
声明式编程执行之前会得到全程序描述,运行前会有编译优化与极致性能优化。但数据和控制流限制比较强,不方便调试,灵活性较低。
计算图
计算图由两个基本部分构成:
- 基本数据结构:Tensor;
- 基本运算单元:Operator。
计算图利用有向无环图来表示计算逻辑和状态,结点代表Operator,边代表Tensor。其中有一些特殊的操作,例如for
、while
等控制流。还会有特殊的边,例如控制边表示结点间的依赖。
计算图解决了AI系统向下硬件执行计算的统一表示,向上承接AI相关程序表示。
计算图与自动微分
引入中间变量将一个复杂的函数分解成一系列基本函数,再将这些基本函数构成一个计算流图。
在计算图中注册前向计算结点和导数计算结点,其中前向结点接受输入计算输出,反向结点接受损失函数对当前张量操作输出的梯度,计算当前张量操作每个输入的Vector-Jacobian乘积。
AI框架中实现自动微分
表达式追踪
区别于计算图的实现方式,前向计算过程中保留中间计算结果,根据反向模式的原理依次计算出中间导数。
这种方式需要保存大量中间计算结果,但好处是方便跟踪计算过程,以PyTorch为代表。
计算图
将导数的计算也表示成计算图,通过Graph IR来对计算图进行统一表示。
这种方式不便于调试跟踪计算和数学表达过程,但优点是方便进行全局图优化,节省内存,以MindSpore和TensorFlow为代表。
优化Pass:给定前向的数据流图,以损失函数为根节点广度优先遍历前向数据流图,按照对偶结构自动生成求导的数据流图。
计算图优化
主要经过以下几个步骤:
- 常量折叠;
- 常量传播;
- 算子融合;
- 表达式简化;
- 表达式替换;
- 公共子表达式消除;
计算图的单设备算子间调度
计算图准确描述了算子之间的依赖关系,可以根据计算图找到相互独立的算子进行并发调度,提高计算并行性。
计算图的切分与多设备执行
计算图切分:给定一个计算图,将计算图切分后放置到多个设备上,每个设备拥有计算图的一部分。
插入跨设备通信:经过切分的计算图会被分为若干子图,每个子图放置在一个设备上,进行跨设备跨子图数据传输。
图与控制流
控制流主要有三种解决方案:
- 后端对控制流语言结构进行原生支持,计算图中允许数据流和控制流的混合,以TensorFlow为代表;
- 复用前端语言的控制流语言结构,用前端语言中的控制逻辑驱动后端数据流图的执行,以PyTorch为代表;
- 后端对控制流语言结构解析成子图,对计算图进行延伸,以MindSpore为代表。
TensorFlow静态图
静态图向数据流图中添加控制流原语。声明式编程计算前能够获取计算图的统一描述,使得编译期能够全局优化。执行流无需在前端语言与运行时反复切换,可以有更高的执行效率。
优点:
- 向计算图中引入控制流原语有利于编译期得到全计算过程描述,发掘运行时效率提升点;
- 解耦宿主语言与执行过程,加速运行时执行效率。
缺点:
- 控制流原语语义设计首要服务于运行时系统的并发执行模型,与深度学习概念差异很大;
- 对控制流原语进行再次封装,以控制流API的方式提供前端用户使用,导致计算图复杂。
PyTorch动态图
动态图复用宿主语言的控制流语句。框架不再维护计算图,神经网络变成一段Python代码,后端张量计算以库的形式提供使能芯片执行计算。
优点:
- 用户能自由使用前端宿主的控制流语言,即时输出张量计算的求值结果;
- 定义神经网络计算就像是编写真正的程序。
缺点:
- 用户易于滥用前端语言特性,带来复杂的性能问题;
- 执行流会在语言边界来回跳转,带来严重的运行时开销;
- 控制流和数据流被严格地隔离在前端语言和后端语言,跨语言优化困难。
MindSpore静态图
MindSpore的静态图会进行源码解析,对计算图进行展开或转换为子图。这个过程会将计算图可以直接表达的控制流直接展开,例如for
循环展开;不能直接表达的会创建子图进行表示,在运行时动态选择执行的子图,例如if else
。
优点:
- 用户能够一定程度自由地使用前端宿主语言的控制流;
- 解耦宿主语言与执行过程,加速运行时执行效率;
- 计算图在编译期得到全计算过程描述,发掘运行时效率提升点。
缺点:
- 硬件不支持控制方式下,执行流会仍然在语言边界跳转,带来运行时开销;
- 部分宿主的控制流语言不能表示,带有一定约束性。
动静统一
动态图转换为静态图,以MindSpore、TensorFlow的Auto-graph以及PyTorch的JIT为代表。
有两种方式:
- 基于追踪Trace:直接执行用户代码,记录算子调用序列,将该序列保存为静态图,执行中脱离前端语言环境,由运行时按照静态图逻辑执行;
- 优点:能够更广泛地支持宿主语言的各种动态控制流语句;
- 缺点:执行场景受限,只能保留程序有限执行轨迹并线性化,静态图失去了源程序完整的控制结构。
- 基于源码解析:以宿主语言的抽象语法树为输入,转化为内部语法树,经过别名分析,SSA(static
single value assignment),类型推断等Pass,转换为计算图表示;
- 优点:能够更广泛地支持宿主语言的各种动态控制流语句;
- 缺点:
- 后端实现和硬件实现会对静态图表示进行限制和约束,多硬件需要切分多后端执行逻辑;
- 宿主语言的控制流语句并不总是能成功映射到后端运行时系统的静态图表示;
- 遇到过渡灵活的动态控制流语句,运行时会退回到由前端语言跨语言调用驱动后端执行。
分布式并行
集群架构
参数服务器模式 (Parameter Server, PS)
主要分为三个步骤:
- 计算损失和梯度;
- 梯度聚合至参数服务器;
- 参数更新并从重新广播。
这里的参数服务器有三种选择:
- CPU充当参数服务器;
- GPU0充当参数服务器;
- 参数服务器分布在所有GPU上。
并行模式
同步并行
等待全部结点完成本次通信之后才继续下一轮计算。
优点:本地计算和通信同步严格顺序化,能够保证并行的执行逻辑与串行相同。
缺点:本地计算结束更早的工作结点需要等待其他工作结点处理,浪费硬件资源。
异步并行
当前batch迭代结束后与其他服务器进行通信,传输网络模型参数。
优点:执行效率高,除了单机通信时间以外没有任何通信和执行之间的阻塞等待。
缺点:网络模型训练不收敛,训练时间长,模型参数反复使用导致无法工业化。
半同步并行
通过动态限制进度推进范围,有限定的宽松同步障的通信协调并行。会跟踪各结点的进度并维护最慢结点,保证计算最快和最慢结点差距在一个预定范围内。
环同步算法
环同步算法 (Ring All Reduce)主要分为两步:
- Scatter-reduce:以环的形式将当前GPU的数据提供给下一个GPU,一次完整的迭代后,每张GPU上都会存有下一张GPU的完整数据;
- All Gather:同样以环的形式,将完整数据传递给下一张GPU,一次完整迭代后,每张GPU的数据都得到了同步。
只需要遍历两次环,所有的GPU就可以完成通信。在这种算法下,集群每秒处理的样本数量与GPU数量几乎是线性的。
通信系统
硬件实现方式
机器内通信:
- 共享内存;
- PCIe;
- NVLink(直连模式)。
机器间通信:
- TCP/IP网络;
- RDMA网络(直连模式)。
软件实现方式
- MPI:通用接口,可调用OpenMPI,MVAPICH2,Intel MPI等;
- NCCL / HCCL:NCCL英伟达通信协议,HCCL华为通信协议,有GPU通信优化,仅支持集中式通信;
- Gloo:Facebook集体通信库,提供集合通信算法。
通信原语
- 一对多:
- Broadcast:广播模式
- Scatter:分散模式
- 多对一:
- Reduce:归约模式
- Gather:聚合模式
- 多对多:
- All-Reduce:归约后广播
- All-Gather:聚合后广播
- Reduce-Scatter:归约后分散
- All-to-All:将每张卡中位置
i
的数据交给卡i
数据并行
数据并行主要有三种方式:
- Data Parallelism, DP:简单数据并行;
- Distribution Data Parallel, DDP:分布式数据并行;
- Fully Sharded Data Parallel, FSDP:全数据切片并行。
简单数据并行
DP会自动切分训练数据,并将模型任务发送到多个GPU,当每个GPU都完成计算后,DP会累计梯度,这样就完成了一个step的训练。同步时使用的是All-Reduce的集合通讯方式。
数据并行是使用多线程实现的,python的多线程会被GIL约束。
GIL:全局解释器锁,本身并不是Python的特性,是Python的其中一种解释器(CPython)实现的互斥锁,它可以防止多个本地线程同时执行Python字节码,因为CPython的内存管理不是线程安全的。
每台机器都会有一个独立的模型,只是单纯的利用单节点的算力。这样就会导致每台机器都会有自己的梯度,需要将这些梯度进行累积。
梯度累积:将不同设备的梯度求和汇总,更新服务器参数的网络模型,再分发给每一台机器。
梯度累积的时机有两种:
- 同步梯度累积:每个设备计算完梯度后,汇总到参数服务器进行统一更新;
- 异步梯度累积:每个设备单独更新自己的服务器参数。
实际中为了模型更好收敛,通常采用同步梯度累积。
分布式数据并行
DDP采用的是多进程的实现方式,没有了GIL的约束。每个进程并不是同步所有的参数,而是同步梯度误差,减少了通讯的数据。采用了环同步算法来提升同步效率。同步时使用的是All-Reduce的集合通讯方式。
主要过程为以下四步:
- 梯度分桶:每个服务器都有多个桶,不同的梯度归属于不同的桶,用来维护待更新的梯度;
- 梯度逆向排序:确定梯度的更新顺序;
- 跳过某些梯度:跳过一些很久未更新,或时间太慢的梯度;
- 集合通讯:对已经满足更新条件的桶中的梯度进行同步。
全数据切片并行
会将网络模型参数、梯度和优化器的状态都进行并行更新,还会将静态的内存卸载到CPU中,从而节省显存 。
反向计算完成后会采用Reduce-Scatter的方式对模型权重参数进行更新,更新完成后再进行All-Gather将数据同步给每个机器
模型并行
张量并行
张量并行和流水线并行的区别在于:
- 流水线并行是层间并行,按照模型的层切分到不同的机器上;
- 张量并行是层内并行,将层内不同的参数切分到不同的机器上。
张量切分方式
- 行切分:按行切分张量,与之计算的左矩阵被迫按列切分;
- 列切分:按列切分张量,与之计算的左矩阵可以直接与切片计算;
张量切分下的随机性问题
- 参数初始化的随机性
不做处理的情况下,参数初始化后,多个设备的参数会初始化为相同的值,和单设备上的参数在数学上不等价,失去了真正的随机性。
所以切分多卡时,将参数切分到多个卡上后,再修改相应卡的随机性,保证各卡的随机种子不同,这样就保证了参数初始化随机性与单卡是相同的。
- 算子计算的随机性
在Dropout、TruncatedNormal、StandardLaplace等算子中会引入随机种子,在多卡中计算时也要注意使用不同的随机种子。
张量重排
MindSpore中为了张量重排而设计了一种张量切片策略:
- 根据网络模型定义的脚本,通过源码转换得到带有切分策略的计算图;
- 为每个未配置切分策略的算子枚举可行的策略;
- 枚举每条边的重排策略和相应的代价;
- 由已配置策略的算法出发,传播到整张计算图。
流水并行
朴素流水线并行
Naive Pipeline parallelism, PP
同一时刻只有一个设备进行计算,其余设备处于空闲状态,计算设备利用率通常较低。
小批次流水线并行
Mini-batch Pipeline parallelism
将朴素流水线并行的batch再进行拆分,减小设备间空闲状态的时间,可以显著提升流水线并行设备利用率。
流水线并行作为模型并行的一部分,一般不会单独使用,是通过混合张量并行、数据并行等方式共同进行的。
深度学习框架机制及分布式并行