结论先行:显存不够时用梯度累积模拟大batch size,效果和真大batch接近,显存占用不变。
一、梯度累积原理
每步用小batch前向反向,不立即更新权重,累积N步梯度后再统一更新。等效于batch size乘以N,但显存只按小batch分配。
二、怎么用
每N步optimizer.step()一次,然后optimizer.zero_grad()。注意学习率要等效放大,或者用线性warmup补偿。
三、累积步数怎么选
目标全局batch除以单卡实际batch。例如目标batch 128,单卡只能跑16,累积8步。
四、和梯度检查点配合
显存实在不够,梯度累积加gradient checkpointing双管齐下,能在24GB显存上跑7B模型训练。
补充:梯度累积步数不要太大,超过16步会导致学习率调度不准。如果需要更大batch,不如多开一张卡用DDP。
小显存大batch训练可在GPU算力平台选择4090实例。




