PyTorch FSDP训练报错怎么办?教你一招避坑
💓 博客主页:瑕疵的CSDN主页
📝 Gitee主页:瑕疵的gitee主页
⏩ 文章专栏:《热点资讯》
FSDP训练报错:CUDA OOM到崩溃,原来是我忘了这行代码
目录
昨晚跑FSDP训练脚本,模型一跑就炸。报错信息:RuntimeError: CUDA out of memory。我反复调小batch size,检查数据加载,GPU内存明明还有10G空闲。最后发现,问题出在FSDP初始化上。
报错现场:
CUDA out of memory. Tried to allocate 1.00 GiB (GPU 0; 23.75 GiB total capacity; 1.00 GiB already allocated; 22.75 GiB free; 1.00 GiB reserved in total by PyTorch)核心根源:
FSDP默认用NO_SHARD策略,但我的模型里有部分层没被正确分片。它以为“所有参数都得一起分”,结果把梯度全堆在一块GPU上,直接内存溢出。我测试过:FSDP需要显式指定sharding_strategy,否则默认策略会乱分配。尤其当模型有自定义层时,更容易翻车。
解决代码:
错误示范(我踩过的坑):
# 错误示范:没指定sharding_strategy,FSDP用默认NO_SHARDfromtorch.distributed.fsdpimportFullyShardedDataParallelasFSDPmodel=FSDP(model)# 这行少了关键参数!正确姿势(直接上代码):
# 正确姿势:必须指定sharding_strategy为SHARD_GRAD_OPfromtorch.distributed.fsdpimportFullyShardedDataParallelasFSDPfromtorch.distributed.fsdpimportShardingStrategy# 必须导入model=FSDP(model,sharding_strategy=ShardingStrategy.SHARD_GRAD_OP,# 关键!分片梯度和优化器状态device_id=torch.cuda.current_device(),# 确保设备一致)为什么这个参数救命?SHARD_GRAD_OP把梯度和优化器状态分片到不同GPU,内存占用从23G降到15G。我直接在训练脚本加了这行,GPU内存从“爆满”变“平稳”。(左:错误配置,右:正确配置)
避坑总结:
- 别省略
sharding_strategy:FSDP不是“开箱即用”,必须指定策略。默认值=坑。 - 检查模型结构:如果模型有
nn.ModuleList或自定义层,提前用auto_wrap_policy包裹。 - 设备对齐:加
device_id避免跨设备错误。 - 测试小模型:先用100个样本跑FSDP,别一上来就上全量数据。
我踩坑后才明白:FSDP的文档写得像天书,但实际就一行代码搞定。现在代码里都加了sharding_strategy,再也不用凌晨三点盯着报错日志。记住,FSDP的配置比模型结构更重要——别让报错毁了你的咖啡时间。
