多卡训练时的内存分配策略总结
作为一个在分布式训练中摸爬滚打多年的工程师，今天想分享一些多卡训练时内存分配的踩坑经验。大家都知道，多卡训练的核心问题就是如何合理分配GPU内存，避免OOM（Out Of Memory）。
1. 基础设置与常见问题
首先，使用PyTorch的分布式训练时，建议先设置torch.cuda.set_per_process_memory_fraction()来限制每个进程使用的显存比例。例如：
import torch
torch.cuda.set_per_process_memory_fraction(0.8)
这样可以避免多个进程争抢显存导致的崩溃。
2. 分布式训练中的显存分配策略
在使用torch.nn.parallel.DistributedDataParallel时，推荐在初始化后立即设置以下参数：
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[0])
然后通过环境变量控制内存分配：
export CUDA_LAUNCH_BLOCKING=1
export NCCL_IB_DISABLE=0
3. 实际调优步骤
- 第一步：设置
NCCL_P2P_DISABLE=1（如果遇到P2P通信问题） - 第二步：调整
torch.cuda.empty_cache()的调用频率，建议在每个epoch结束后清理一次 - 第三步：使用
torch.cuda.memory_summary()查看内存占用情况
4. 避坑指南
千万别在训练过程中随意修改batch size或者模型结构，这会导致内存分配混乱。最好在训练前就确定好参数，然后通过小批量测试来验证。
经验分享完毕，大家在实际使用中还遇到过哪些问题？欢迎在评论区交流！
讨论