1.  Megatron batchsize设置

global_batch_size = micro_batch_size * DP * num_micro_batches

                                 = 每step喂多少* 多少并行 * 串行几次

参数 角色
micro_batch_size 每次进流水线的样本数(影响显存)
DP 数据并行路数据(增加GPU减少M)
M(num_microbatches) 流水线串行轮数(影响bubble rate)
global_batch_size 三者之积,决定每步学到多少数据

num_microbatches是“一个optimizer step内,流水线要跑几轮”

micro_batch_size通常设置为1或比较小的数;num_microbatches会上万, 设置会比较大。

流水线中每个stage需要缓存所有in-flight的激活值用于反向传播。

micro_batch_size和num_micro_batches两个参数控制的是不同资源。

参数 控制什么 影响
micro_batch_size 显存:单次前向/反向的激活值大小 显存占用
num_microbatches(M) 流水线效率 GPU利用率

具例:

micro_batch_size = 1   ,   每次只送1个样本,激活值尽量小,省显存

M = 3XXXX,   串行跑3万多轮,保证bubble 约等于0 。

假设反例:

  •  micro_batch_size 设置比较大,激活值存储很容易OOM,激活值需要至少保存micro_batch_size*seq_len*hidden_size*num_layers。
  • num_micrbatches(M) 比较小,比如等于1,在PP=20的场景。一个step只跑一轮,bubble rate = (20-1)/(1+20-1) = 95%, 这时GPU几乎都在空转。
2.  梯度累积

通常说的梯度累积是W的梯度,X的梯度会在zbv的B阶段处理完就会丢弃。B阶段和W阶段的依赖关系:

阶段 计算公式 依赖
B阶段 dx = dy @ W 上游梯度dy + 权重W
W阶段 dw += dy.T @ X 上游dy + 前向激活值 X

在一个step最后optimizer 执行前,W的梯度会进行num_microbatches轮的累积,每轮还需要按DP维度进行累积。公式中W阶段layer0需要的dy,是B阶段layer0+1层算出的dx;B阶段的layer0+1层dx显存释放,必须要等对应的layer0层的W阶段执行完成。

在loss前向,每个microbatchsize的loss会先除以M,保证w的梯度累积不会过大。

M 个 microbatch 的梯度累加:

W_grad += dL_1/M
W_grad += dL_2/M
W_grad += dL_3/M
...
W_grad += dL_M/M
─────────────────────────
W_grad = (dL_1 + dL_2 + ... + dL_M) / M
       = 平均梯度

一个global_batch_size的所有B和W阶段全部完成,最后执行optimizer step, 更新W参数。

Logo

AtomGit 是由开放原子开源基金会联合 CSDN 等生态伙伴共同推出的新一代开源与人工智能协作平台。平台坚持“开放、中立、公益”的理念,把代码托管、模型共享、数据集托管、智能体开发体验和算力服务整合在一起,为开发者提供从开发、训练到部署的一站式体验。

更多推荐