这里回顾dlsys第十五讲，本讲内容是如何训练大模型。

课程主页：

大纲

  • 内存节省技术；
  • 并行和分布式训练；

内存节省技术

机器学习系统的要素

更大的数据集需要更大的模型容量。这反过来又对计算设备提出了要求。机器学习的成功是所有三个要素的结合。 最近的许多进展要求我们将这三者都推向极限。

今天我们将学习两个主题：

  • 如何减少内存消耗，以便我们可以将更大的模型装入单个设备。
  • 如何扩大训练过程。

回顾：GPU内存层次结构

  • 共享内存：每个核64KB；
  • GPU内存（全局内存）：
    • RTX3080 10GB；
    • RTX3090 24GB；
    • A100 40/80GB；

内存消耗的来源

  • 用于训练的典型计算图的简化视图，权重被省略并隐含在梯度步骤中：

内存消耗的来源：

  • 模型权重；
  • 优化器状态；
  • 中间激活值；

仅用于内存节省推理的技术

我们只需要$O(1)$内存来通过循环通过两个缓冲区来计算$N$层深度网络的最终输出。

训练的激活内存成本

因为需要为梯度步骤保持中间值（检查点）。训练$N$层神经网络需要$O(N)$内存。我们将使用以下简化视图来结合梯度和前向计算：

AD中的checkpointing技术

  • 只有checkpoint被标记为彩色（步骤 0）；
  • 重新计算小段中缺失的中间节点（步骤 1、2）；

Sublinear内存开销

对于一个$N$层神经网络，
如果我们隔$K$层存储checkpoint，那么内存开销为：

说明：

  • 第一项为checkpoint开销；
  • 第二项为重新计算的开销；

选择$K=\sqrt N$，内存开销会比较小。

并行和分布式训练

并行训练问题

利用可能分布在多个工作节点上的多个 (GPU) 设备来训练模型。

模型并行

模型并行将部分计算图映射到workers：

分解方式：

对计算图进行分区，将发送/接收对放在边界中。

数据并行

考虑损失函数：

让每个worker访问小批量的$\frac B \kappa$，并运行梯度计算，然后将所有梯度加在一起，每个worker都运行相同的模型副本。

Allreduce抽象

接口：

result = allreduce(float buffer[size])

例子：

通过allreduce进行数据并行训练

同一计算图的多个副本并行运行：

Parameter server抽象

接口：

ps.push(index, gradient)
ps.pull(index)

在服务器上执行权重更新（键值存储）：

通过parameter server进行数据并行训练

在远程服务器上更新结果并将更新后的结果发回。

通信计算重叠

在通过网络发送数据的同时也有许多计算的机会，这样可以节省时间开销：

并行化总结

  • 模型并行：对计算图划分；
  • 数据并行：对数据划分；
  • 在所有情况下，利用机会将计算与通信重叠。

高级并行化方法

有更多方法可以并行化计算图，一些可选的参考读物：

  • ZeRO: Memory Optimizations Toward Training Trillion Parameter Models；
  • Beyond Data and Model Parallelism for Deep Neural Networks；
  • GSPMD: General and Scalable Parallelization for ML Computation Graphs；