PyTorch分布式训练怎样容错

在 PyTorch 分布式训练（尤其是多机多卡、大规模训练）中，容错（Fault Tolerance）通常不是一个“开箱即用”的单参数开关，而是需要结合 训练框架、检查点机制、进程管理、调度系统 来共同实现。下面从原理 → 常用方案 → 具体做法系统讲一下。

一、PyTorch 分布式训练为什么容易“挂”

常见失败原因：

1. 节点 / 进程崩溃

• OOM
• 硬件故障（GPU、网卡）
• 节点被调度系统抢占

1. 网络问题

• NCCL 通信超时
• 网络抖动

1. 训练逻辑错误

• NaN / Inf
• DataLoader 死锁

1. 弹性资源变化

• 云上 Spot 实例被回收
• 扩缩容

二、核心容错思想（非常重要）

容错 ≠ 自动修复错误，而是“从最近安全状态恢复”
关键三件事：

1. 定期保存 Checkpoint
2. 进程能自动重启
3. 重启后能重新组网并继续训练

三、官方 & 主流容错方案

1️⃣ PyTorch Elastic（torch.distributed.elastic）⭐ 推荐

适用场景

• 多机多卡
• 节点可能动态变化
• 云训练 / Spot 实例

核心组件

• torch.distributed.launch ❌（已不推荐）
• torchrun ✅
• ElasticAgent
• Rendezvous（动态组网）

特点

• 支持 min / max 节点数
• 某节点挂了 → 剩余节点等待 → 新节点加入 → 重新初始化
• 不必手动写重启脚本

示例

torchrun \
--nnodes=2:4 \
--nproc_per_node=8 \
--rdzv_id=job1 \
--rdzv_backend=etcd \
--rdzv_endpoint=etcd:2379 \
train.py

注意

• 必须配合 Checkpoint
• torchrun 本身不保存模型

2️⃣ Checkpoint + 自动重启（最稳方案）

保存内容

{
"model": model.state_dict(),
"optimizer": optimizer.state_dict(),
"lr_scheduler": scheduler.state_dict(),
"epoch": epoch,
"step": step,
"rng": {
"torch": torch.get_rng_state(),
"cuda": torch.cuda.get_rng_state_all(),
"numpy": np.random.get_state(),
"python": random.getstate()
}
}

恢复

model.load_state_dict(ckpt["model"])
optimizer.load_state_dict(ckpt["optimizer"])

✅ 这是所有容错方案的基础

3️⃣ NCCL 容错参数（防止“假死”）

export NCCL_ASYNC_ERROR_HANDLING=1
export TORCH_NCCL_ASYNC_ERROR_HANDLING=1
export NCCL_TIMEOUT=1800

或在代码里：

torch.distributed.init_process_group(
backend="nccl",
timeout=datetime.timedelta(seconds=1800)
)

作用：

• 网络异常时抛错而不是卡死
• 配合 elastic 可重启

4️⃣ 高级框架级容错（强烈推荐）

✅ PyTorch Lightning

• ckpt 自动保存
• auto_restart
• 支持 DDP / FSDP

Trainer(
strategy="ddp",
plugins="torchrun",
enable_checkpointing=True
)

✅ DeepSpeed

• elastic training
• 自动 checkpoint
• ZeRO + 容错

✅ HuggingFace Accelerate

• 配合 torchrun
• 简化容错写法

四、典型容错架构（生产级）

Scheduler (K8s / Slurm)
↓
TorchElastic + etcd
↓
定期 Checkpoint → 共享存储 (NFS / S3)
↓
进程崩溃 → 自动重启 → 加载 Checkpoint

五、常见坑（非常重要）

❌ 只保存 model，不保存 optimizer
→ 学习率 & momentum 错乱
❌ Checkpoint 太稀疏
→ 一挂回到几小时前
❌ 不用 etcd / 静态 IP
→ 节点重连失败
❌ 忽略 DataLoader worker 死锁
→ 加：

persistent_workers=True

六、最小可用容错模板

torchrun \
--nnodes=1:4 \
--nproc_per_node=8 \
--rdzv_backend=etcd \
--rdzv_endpoint=etcd:2379 \
train.py

# train.py
if os.path.exists("ckpt.pt"):
load_checkpoint()
for epoch in range(start, max_epoch):
train()
save_checkpoint()

七、总结一句话

PyTorch 分布式容错 = torchrun + etcd + 周期 Checkpoint + 自动重启
如果你愿意，我可以：

• 给你一个 完整 torchrun + etcd 示例
• 或针对 Slurm / K8s / 单机多卡 给具体配置
• 或讲 FSDP / DeepSpeed 的容错差异

你现在是哪种训练环境？