15 大模型训练 内存优化

先看GPU结构，我们常说显存的时候，说的一般就是Global memory

训练的过程中，我们为了反向传播过程，必须将中间的结果（激活值）存储下来。

在训练的过程中，那些会消耗内存呢？

• model weights
• optimizer sates
• intermediate activation values

对于有N层的神经网络来说，内存的消耗是O(N)的。

检查点技术

在前向传播的时候，只选择保留部分数值，当进行反向传播时，所需要的中间值会进行重计算。

这样虽然会增减计算成本，但是也大大减少了内存占用。

模型并行

将模型进行拆分

数据并行

将minibatch 划分成更小的micobatch，训练每个batch的时候，每个工作节点获得一个microbatch，

梯度更新

各个节点之间计算出来的梯度要统一，可以使用 all-reduce或者 使用一个参数服务器用来统一更新各个节点之间的梯度。

为了加快训练，可以使得参数传递和计算过程互相掩盖

READING LIST:

• ZeRO
• Beyond Data and Model Parallelism for Deep Neural Networks
• GSPMD: General and Scalable Parallelization for ML Computation Graphs