梯度累积与大批次训练技巧

在 GPU 内存有限时，梯度累积是训练大模型的关键技巧——它通过多次前向传播累积梯度，模拟更大的 batch size。

实现

def train_with_accumulation(model, dataloader, optimizer, accum_steps=8):
    model.train()
    optimizer.zero_grad()
    
    for i, batch in enumerate(dataloader):
        with torch.cuda.amp.autocast():
            loss = model(batch) / accum_steps
        loss.backward()
        
        if (i + 1) % accum_steps == 0:
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            optimizer.zero_grad()

有效 Batch Size 计算

有效 Batch Size = micro_batch_size × accum_steps × num_gpus

例如：micro_batch=4, accum_steps=8, 8 GPUs → 有效 Batch Size = 256

梯度累积 vs 数据并行

梯度累积是时间换空间（串行累积），数据并行是空间换时间（并行计算）。两者可组合使用。

相关资源

• 大规模训练技术
• DeepSpeed 框架
• FSDP 文档