一、单机单卡训练
1. 基本概念
- 芯片的内存
- 模型训练时的内存主要消耗在：
神经元的中间值
参数值（包括权重、梯度值和优化器的状态值）
2. 优化方法
单卡速度主要由单块AI加速芯片的运算速度、数据IO来决定。
- 优化方法：可以用激活重计算这个方法来进行优化。
- 具体操作为：
在网络中标记少量的算子 ，前向计算只保留这些被标记的算子的输出结果（激活值）， 其余前向算子的输出结果直接被释放，这样就可以极大减少激活值消耗的内存。当反向更新梯度需要前向算子的输出时，利用被标记的算子重新计算获取。
激活重计算是一个以时间换空间的策略。
通过类似于检查点的机制，通过标记部分算子，存储他们的输出值，其他算子的输出值就直接忽略，在反向传播时，利用标记算子，来继续前项传递，求出其余算子的前项输出来进行梯度调整。反向阶段完成后，优化器进行权重调整。
神经元的中间值主要指的是前向算子的output，这些output在训练结束后是不需要保留的，所以称为中间值。如果是推理网络，算子执行完，输出给下一个当完输入后，就可以释放了。但训练网络不同，前向算子的输出要一直保留到反向网络执行时，因为前向算子的输出还是对应反向算子的输入。即前向网络中的中间值的声明周期延长到了反向网络阶段了。
二、分布式训练
1. 基本概念
为什么要分布式训练？
- 大模型训练需要的内存超过了单个机器的内存
- 大模型训练的数据太大
- 想要减少训练时间
在多台机器上面训练模型。假设有两台机器，每台机器上有2个GPU，那么我们就启动了4个进程。
进程之间用如下方式通信： - peer-to-peer：一个进程向另一个进程发送数据
- collective：一组进程执行分散、聚集、all-reduce、广播等操作
2. 分类
2.1 数据并行
将数据分割成几部分，分配到不同的设备上。这相当于沿批次维度对训练过程进行并行。在反向传播之后，每张卡上的梯度和权值、优化器的状态值是一样的，即各张卡之间进行的是完全一样的计算。因此可以将这部分的计算分散到各张卡，每张卡仅保存部分的梯度和状态值，在不同的设备上的模型参数保持同步。
缺点：每个GPU都有全部的模型参数。
2.2 模型并行
将模型分割成几部分，分配到不同的设备上。
2.2.1 流水线并行
将模型按层分割成若干块，分配到不同的设备上。在前向传递过程中，每个设备将中间的激活传递给下一个阶段。在后向传递的过程中，每个设备将输入张量的梯度传回给前一个流水线阶段。
缺点：会有设备参与计算的冒泡时间，造成浪费。
2.2.2 张量并行
将一个张量沿特定维度分成N块，每个设备只持有张量的1/N。
将设要计算C = AB，将B沿着列分割成[B0 B1 B2 B3]，每个设备上有一列。然后将A与每个设备的B的一列相乘，得到[AB0 AB1 AB2 AB3]，每个设备仍持有一部分的结果。然后将张量串联起来，得到最后的结果。
3. 小结
三、参考
- https://colossalai.org/zh-Hans/docs/concepts/distributed_training
- https://blog.csdn.net/RogersStar/article/details/123617394
- https://zhuanlan.zhihu.com/p/432289008