为何PyTorch DDP模式下模型显存占用为单卡两倍？如何优化？

• PyTorch DDP为了提升分布式梯度同步的效率，默认会为每个GPU上的模型可学习参数，单独申请一块和参数体积完全相同的梯度累加缓冲区，用来跨进程汇总梯度、执行AllReduce同步操作。
• 你测试用的1000x1000 Linear层，float32精度下参数总大小约为4MB，和你日志里打印的单卡初始占用4004352字节完全匹配；DDP初始化后新增的4MB就是梯度缓冲区的占用，所以总显存刚好翻倍。

是，这是DDP的默认设计，属于正常的预期行为。梯度缓冲区可以避免每次迭代临时申请显存、重复拷贝数据，是DDP高性能实现的一部分。

你可以通过以下方法消除或减少这部分额外占用：

1. 启用gradient_as_bucket_view参数
   PyTorch 1.10及以上版本支持该参数，开启后DDP会直接复用参数的.grad属性对应的内存空间存储梯度，不需要额外申请独立的梯度缓冲区，可直接消除这一倍的显存开销。只需修改DDP初始化代码即可：

model = DDP(model, device_ids=[device_id], find_unused_parameters=False, gradient_as_bucket_view=True)

注意该模式下不要手动修改、覆盖参数的.grad属性，否则会导致梯度同步错误。
2. 冻结不需要更新的参数
对于训练全程不需要更新的参数，直接设置param.requires_grad = False，DDP不会为无梯度的参数创建梯度缓冲区，可减少对应部分的显存占用。
3. 使用混合精度训练
配合torch.cuda.amp做混合精度训练，参数和梯度都用FP16精度存储，整体显存占用直接减半，额外的缓冲区开销也会同步降低。
4. 调整梯度分桶配置
如果模型参数分布零散，可以适当调小DDP的bucket_cap_mb参数，减少显存碎片带来的额外占用。

内容的提问来源于stack exchange，提问作者Jacob Stern