混合精度训练如何实现分布式训练

混合精度训练与分布式训练的整合思路

• 在分布式数据并行中，每个进程/GPU各自执行前向与反向，梯度通过 ring-all-reduce 在所有进程间求平均；混合精度则把大部分算子用 FP16/BF16 执行，关键步骤保留 FP32（如主权重、优化器状态、梯度累积与更新），并用动态损失缩放避免下溢。二者叠加后，既减少显存占用与通信量，又保持数值稳定与收敛精度。

原生 PyTorch 最小落地流程

• 初始化进程组与设备
• 使用 torch.distributed.init_process_group(backend="nccl")；设置 local_rank、调用 torch.cuda.set_device(local_rank)；模型与张量都放到对应 device。
• 数据加载
• 用 DistributedSampler 切分数据，保证各进程看到不重叠样本；DataLoader 建议开启 pin_memory 与合适的 num_workers。
• 模型包装
• 使用 torch.nn.parallel.DistributedDataParallel(DDP) 包装模型，放在 AMP 初始化之后。
• 混合精度
• 使用 torch.cuda.amp.autocast 在前向上下文指定计算精度（如 torch.float16 或 torch.bfloat16）；用 GradScaler 进行损失缩放、反向与优化步骤更新。
• 启动训练
• 通过 torch.distributed.launch 或 torch.multiprocessing.spawn 启动多进程；单机多卡常见命令：CUDA_VISIBLE_DEVICES=0,1,2,3 python -m torch.distributed.launch --nproc_per_node=4 train.py。
• 同步批归一化
• 多卡时建议开启 SyncBatchNorm（可用 Apex 的 convert_syncbn_model 或 PyTorch 原生 SyncBatchNorm），保证统计量一致。

Apex 实现路径（适合已有 Apex 代码）

• 初始化与模型
• 先 amp.initialize(model, optimizer, opt_level="O1")（或 O2），再包装 DDP；Apex 的 DDP 支持延迟梯度规约（如 delay_allreduce）以改善吞吐。
• 前向与反向
• 前向置于 amp.autocast 上下文；反向使用 with amp.scale_loss(loss, optimizer) as scaled_loss: scaled_loss.backward()，随后 optimizer.step() 与 scaler.update()。
• 同步 BN
• 使用 apex.parallel.convert_syncbn_model(model) 将普通 BN 转为同步 BN，提升多卡一致性。

大模型场景的分布式与混合精度配置

• DeepSpeed ZeRO + 混合精度
• 在配置文件中启用 fp16.enabled: true 或 bf16.enabled: true；结合 ZeRO-3 进行参数/梯度/优化器状态分片，显著降低显存占用；可配合 gradient_clipping、initial_scale_power 等参数稳定训练。
• FSDP + 混合精度
• 通过 FSDP 的 mixed_precision_policy: FP16/BF16 启用混合精度；常与 Transformer 模型的分层自动包装策略联用，兼顾显存与性能。
• Accelerate 统一入口
• 使用 Accelerate 的 Accelerator 与 accelerator.prepare(...) 一行接入多机多卡与混合精度；可通过 DeepSpeedPlugin/FSDP 插件无缝启用高级特性，代码改动最小。

关键注意事项与排错要点

• 精度与稳定性
• 优先选择硬件友好的 BF16（如 Ampere+），或在 FP16 下合理设置 初始缩放因子 与 梯度裁剪；训练中关注 loss scale 与 溢出次数，必要时增大 initial_scale_power 或调低学习率。
• 批归一化与指标
• 多卡训练使用 SyncBatchNorm；验证阶段注意跨进程 all_reduce 或只在 rank 0 做指标统计，避免重复累计。
• 启动与设备放置
• 每个进程只绑定 1 张 GPU（local_rank 与 set_device 一一对应）；确保 DistributedSampler 的 shuffle=True 与 drop_last=True 配置合理；多机训练正确设置 MASTER_ADDR/MASTER_PORT 与 RANK/WORLD_SIZE。
• 不要混用
• DP（DataParallel） 仅支持单机且不利于扩展；多卡/多机请使用 DDP 或 FSDP/DeepSpeed；Apex 与原生 AMP 二选一，避免重复包装或冲突。