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

所有评论(0)