这里回顾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;