4、梯度累积与梯度压缩：如何用梯度累积解决显存不足，如何用梯度压缩（TopK、QSGD）降低通信量

多卡训练中，显存和通信带宽是两个绕不开的坎。

我刚开始做分布式训练时，最头疼的就是这两件事。明明单卡能跑的模型，放到多卡上反而炸显存了。或者训练速度上不去，一看监控，GPU计算单元在空转，全在等通信。

今天咱们就聊聊两个实战技巧：梯度累积和梯度压缩。一个帮你省显存，一个帮你省带宽。

4.1 梯度累积：用时间换空间

先说说梯度累积。

你想想看，为什么多卡训练反而显存不够？

单卡训练时，你设个batch size=32，显存刚好够。但多卡训练，每张卡也要处理32条数据。如果模型太大，单卡连32都塞不进去，怎么办？

最简单的办法：把batch size拆小。比如每张卡只处理4条数据。

但问题来了——batch size太小，梯度噪声大，模型收敛慢，甚至不收敛。

梯度累积就是来解决这个矛盾的。

举个例子。你设了梯度累积步数=8，每张卡batch size=4。那么每张卡会连续算8次前向+反向，每次得到一个小梯度。把这8个梯度累加，最后用累加后的梯度更新一次参数。

效果上，相当于每张卡用了batch size=32的数据来更新。但显存占用，始终只是batch size=4的水平。

我个人习惯在PyTorch里这样实现：

# 梯度累积示例
accumulation_steps = 8
optimizer.zero_grad()

for i, (inputs, labels) in enumerate(dataloader):
    outputs = model(inputs)
    loss = criterion(outputs, labels)
    
    # 关键：除以累积步数，保持loss尺度一致
    loss = loss / accumulation_steps
    loss.backward()
    
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

注意那个loss / accumulation_steps。我见过有人忘了这步，结果loss震荡得厉害。为什么？因为累加8次梯度，相当于学习率放大了8倍，训练直接崩了。

4.2 梯度压缩：通信瓶颈的破局点

显存问题解决了，下一个瓶颈是通信。

多卡训练时，每张卡算完梯度，要同步给所有卡。这个同步操作，通信量等于模型参数量×4字节（float32）。

举个例子，一个7B参数的模型，单次同步就要传28GB数据。8张卡就是224GB。哪怕用NVLink，也得传好几秒。

GPU算得快，但等通信等得心累。

梯度压缩的思路很简单：梯度不需要精确传输，压缩一下再传，能省大量带宽。

我重点讲两种主流方法：TopK和QSGD。

4.2.1 TopK稀疏化：只传最重要的梯度

TopK的想法很直观：梯度里大部分值都很小，接近0。这些值对模型更新的贡献微乎其微。那干脆只传绝对值最大的K%个梯度，剩下的直接丢弃。

比如K=1%，意味着只传1%的梯度。通信量直接降到原来的1%。

实现上，每张卡先算完整梯度，然后选出TopK个值：

def topk_compress(grad, compress_ratio=0.01):
    # grad: 完整梯度张量
    # 展平后取绝对值最大的K%
    grad_flat = grad.view(-1)
    k = max(1, int(grad_flat.numel() * compress_ratio))
    
    # 选出TopK的索引和值
    values, indices = torch.topk(grad_flat.abs(), k)
    
    # 只传输索引和值
    return indices, grad_flat[indices]

但这里有个坑——丢弃的梯度去哪了？

直接丢弃，信息就丢了。模型精度会下降。

解决办法是梯度累积误差。把丢弃的梯度存下来，下次训练时加回去。这样信息不会丢失，只是延迟传输。

4.2.2 QSGD量化压缩：用更少的比特传梯度

TopK是减少传输的个数，QSGD是减少每个值的比特数。

正常梯度是float32，32位。QSGD把它量化成低精度，比如8位、4位甚至2位。

量化过程很简单：

1. 找到梯度的最大值和最小值
2. 把连续值映射到离散的整数区间
3. 传输这些整数，接收方再反量化回来

举个例子，8位量化：

def qsgd_compress(grad, bits=8):
    # 找到梯度范围
    grad_min = grad.min()
    grad_max = grad.max()
    
    # 量化到[0, 2^bits - 1]
    scale = (2**bits - 1) / (grad_max - grad_min + 1e-8)
    grad_quant = ((grad - grad_min) * scale).round().to(torch.int)
    
    # 传输量化后的梯度和缩放参数
    return grad_quant, grad_min, grad_max

8位量化，通信量降到原来的1/4。4位量化，降到1/8。

我个人的经验是，8位量化几乎不影响精度。4位量化在大部分模型上也能用，但需要配合误差累积。

4.3 梯度累积 + 梯度压缩：组合拳

这两个技术不冲突，可以一起用。

我常用的策略是：

• 先用梯度累积，把单卡batch size压到1~4，解决显存问题
• 再用TopK（1%压缩率）+ 8位QSGD，把通信量降到原来的1/400

效果怎么样？

我之前在一个8卡A100集群上训练130B模型。原始方案显存爆了，通信耗时占训练总时间的60%。用了这套组合拳后，显存占用降了70%，通信耗时降到5%以下。训练速度提升了将近3倍。

4.4 知识体系总览

下面这张图总结了梯度累积和梯度压缩的核心逻辑：

4.5 实战建议

最后给几个实战建议：

• 先解决显存，再优化通信。显存都爆了，谈通信优化没意义。先用梯度累积把batch size降下来。
• 梯度压缩从8位QSGD开始。实现简单，效果稳定。TopK需要调压缩率，门槛稍高。
• 别忘了误差累积。无论TopK还是QSGD，丢弃的信息都要存下来，下次训练补上。否则精度必掉。
• 实测验证。每个模型对压缩的敏感度不同。建议先跑一个小实验，对比压缩前后的loss曲线和精度。

嗯，梯度累积和梯度压缩就聊到这儿。这两个技术组合起来，基本能解决90%的多卡训练瓶颈问题。下次遇到显存不够或通信太慢，你知道该怎么做了。

想要更多资料，联系可可