💓 博客主页：瑕疵的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_SHARD
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

model = FSDP(model)  # 这行少了关键参数！

正确姿势（直接上代码）：

# 正确姿势：必须指定sharding_strategy为SHARD_GRAD_OP
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import ShardingStrategy  # 必须导入

model = FSDP(
    model,
    sharding_strategy=ShardingStrategy.SHARD_GRAD_OP,  # 关键！分片梯度和优化器状态
    device_id=torch.cuda.current_device(),  # 确保设备一致
)

为什么这个参数救命？
SHARD_GRAD_OP把梯度和优化器状态分片到不同GPU，内存占用从23G降到15G。我直接在训练脚本加了这行，GPU内存从“爆满”变“平稳”。
（左：错误配置，右：正确配置）

避坑总结：

1. 别省略sharding_strategy：FSDP不是“开箱即用”，必须指定策略。默认值=坑。
2. 检查模型结构：如果模型有nn.ModuleList或自定义层，提前用auto_wrap_policy包裹。
3. 设备对齐：加device_id避免跨设备错误。
4. 测试小模型：先用100个样本跑FSDP，别一上来就上全量数据。

我踩坑后才明白：FSDP的文档写得像天书，但实际就一行代码搞定。现在代码里都加了sharding_strategy，再也不用凌晨三点盯着报错日志。记住，FSDP的配置比模型结构更重要——别让报错毁了你的咖啡时间。