PyTorch 生产落地案例

五、生产落地案例

案例一：把单卡实验脚本改造成 8 卡可复现训练

业务背景：一个推荐模型团队，原来在单卡上跑实验，一次全量训练要跨夜。想扩到 8 卡，但第一次改造后出现了「8 卡比单卡还慢」「损失不收敛」两个问题。

怎么用：改用 DDP + DistributedSampler，去掉 DataParallel；把随机种子按 rank 派生，保证每个进程的数据增强不同但整体可复现；日志与 checkpoint 只在 rank 0 写；用 torchrun 启动并让框架自己分配 rank。

关键配置：set_epoch 每轮调用；num_workers 按 CPU 核数除以卡数分配；梯度裁剪保持在 backward 之后；日志按步数而非时间输出，避免多进程互相干扰。

效果指标：8 卡相对单卡的吞吐提升接近线性（实际会因通信与数据加载打折扣，以自家 profile 结果为准）；实验周期从跨夜缩短到当天几轮，迭代次数明显增加。

踩过的坑：一是忘了 set_epoch，每个 epoch 数据顺序完全相同，等价于训练集缩小；二是 num_workers 开太大把主机内存打满，进程被系统杀掉，日志里只看到「Killed」；三是多进程同时写同一个 checkpoint 文件导致文件损坏，改成 rank 0 写入、其余进程等待。

案例二：用 AMP 与梯度检查点把大模型塞进现有显卡

业务背景：一个团队要在有限显存下微调一个参数量不小的模型，直接跑会 OOM，降 batch 又会让训练不稳定。

怎么用：开启 AMP（前向 autocast、反向用 GradScaler 缩放损失）；对重复堆叠的块启用梯度检查点，用时间换显存；把优化器换成内存占用更小的实现；分步排查显存占用，先确认是激活值还是优化器状态占据了大头。

关键配置：autocast 的 dtype 选择要跟硬件匹配；GradScaler 与梯度裁剪配合时，要先把梯度 unscale 再做裁剪，否则裁剪的是被放大过的梯度；batch size 与梯度累积步数按「等效 batch」调平。

效果指标：在同样显存下可用 batch 变大，单步时间增加（换显存的代价），整体能跑通原本跑不起来的配置。

踩过的坑：一是先裁剪再 unscale，导致梯度尺度错误、训练发散；二是混合精度下某些层对数值敏感，需要用 autocast(enabled=False) 局部关掉；三是开了梯度检查点后忘记 model.train() 与 eval 切换，dropout 行为错乱。