训练框架与 Profiling

训练框架 = 并行策略 + 显存优化 + 可观测性的组合。PyTorch 原生两条路: DDP(数据并行,每卡持全量模型副本,反向后对梯度做一次 All-Reduce,等价 ZeRO-0——只解决算力扩展、无显存收益);FSDP(Fully Sharded Data Parallel,官方实现的 ZeRO-3 机制): 每个权重张量按 rank 分片,前向时 All-Gather 恢复本层完整参数、用后即弃,反向时对梯度 Reduce-Scatter 得到分片,优化器状态随分片天然分布——每卡显存 $16\Psi/N_d$,通信量约 DDP 的 1.5 倍(3Φ vs 2Φ);按层分片 + 参数预取把通信藏进计算,近线性扩展。DeepSpeed 提供 ZeRO-1/2/3 + CPU/NVMe Offload,机制与 FSDP 等价: 单机多卡、少写代码选 FSDP;跨机、超大规模、需要 Offload 或既有生态选 DeepSpeed。Profiling 瓶颈定位(torch.profiler + Chrome trace): 把单步时间拆成四块——① 计算(CUDA kernel 耗时与 SM 占用率,小 kernel 过多 → 算子融合/减少同步);② 通信(NCCL 时间占比,梯度 All-Reduce 是否与反向重叠,未重叠 → 调 bucket 大小/开异步);③ IO(DataLoader worker 是否打满、磁盘/网络是否拖后腿,调 worker 数与 prefetch_factor);④ 空闲(GPU 空转,同步点过多或单卡串行)。排查顺序: 先看 CUDA time 占比,再按 forward/backward/optimizer 拆,最后定位到具体 kernel 或通信原语。

面试高频“DDP/FSDP/DeepSpeed 区别与选型”“FSDP 分片原理”“训练变慢、GPU 利用率低怎么排查”。

裸 DDP 在单卡放不下 $16\Psi$ 时无能为力;FSDP/ZeRO 把每卡显存降到 $16\Psi/N_d$,让 N 卡能训练单卡 N 倍大的模型,通信量仅多约 1.5 倍、扩展性近线性;profiler 把“训练慢”从经验判断变成可量化的时间账(compute/comm/IO/idle 四类占比),瓶颈定位从小时级缩短到分钟级。

对比三者显存账: DDP 每卡 $16\Psi$(全量副本,ZeRO-0),FSDP/ZeRO-3 每卡 $16\Psi/N_d$,DeepSpeed 额外支持 CPU/NVMe Offload;分别说明解决的是算力扩展还是显存扩展。

讲清 FSDP 分片原理: 权重张量按 rank 分片 → 前向 All-Gather 恢复本层参数 → 用后即弃 → 反向 Reduce-Scatter 梯度分片;为什么通信量约 DDP 的 1.5 倍(3Φ vs 2Φ)、为什么仍近线性扩展。

选型: 单机多卡、想少写代码选什么(FSDP);跨机/需要 Offload/超大规模/既有生态选什么(DeepSpeed);FSDP 的 FULL_SHARD 与 SHARD_GRAD_OP 有何区别。

用 torch.profiler 定位瓶颈: 把单步时间拆成 compute/comm/IO/idle 四类;通信未与反向重叠怎么调(bucket 大小、异步梯度),IO 打不满怎么调(worker 数、prefetch_factor、数据格式),给出排查顺序。

手算场景: 70B 在 8×80 GB 节点上为什么 DDP 装不下($16\Psi = 1120$ GB > 640 GB);FSDP/ZeRO-3 每卡 $1120/8 = 140$ GB 仍超 80 GB,须 16 卡或再开 Offload(GPU 仅剩 $2\Psi/N_d = 17.5$ GB);完整说出两层账。