简介：这份资源围绕 PyTorch Elastic（TorchElastic）展开，面向需要以容错、弹性方式启动分布式 PyTorch 作业的算法工程师与机器学习从业者。它解决的是多节点训练中节点故障、动态扩缩容与作业编排的痛点，适合已掌握 PyTorch 基础、希望进一步实践分布式训练的中高级开发者。压缩包共 107 个文件，约 554KB，以 23 个 Python 脚本、18 个 Markdown 文档、18 个 YAML 配置、12 个 Go 源码及若干 Dockerfile、Shell 脚本为主，分别承担训练示例、说明文档、编排配置与容器化构建等职责。资源中给出了容错与弹性两种典型场景：4 节点、每节点 8 个训练器共 32 个训练器的启动方式，以及 1 至 4 节点动态伸缩、作业在单节点健康后即可开始的弹性调度思路，并涉及 etcd 作为 rendezvous 后端的配置。已有 2669 人学习，读者可据此理解 TorchElastic 的启动参数、rdzv 机制与工程化落地路径。
1. 单机多卡跑得好好的，一到多机就崩：PyTorch Elastic 到底在解决什么
你有没有遇到过这种场景：单机 8 卡训练一切正常，torchrun一跑就通，loss 曲线平滑得像教科书。结果一上多机，要么是某个节点挂了整个任务直接死掉，要么是节点恢复后训练没法从断点续上，要么是RANK和WORLD_SIZE对不上导致梯度同步卡死。这不是玄学，是分布式训练里最典型的容错问题。
PyTorch Elastic（也叫torch.distributed.elastic，简称 torchelastic）就是 PyTorch 官方给出的答案。它不是一个新框架，而是torchrun背后的运行时系统，负责进程管理、节点发现、弹性伸缩和故障恢复。简单说，它让分布式训练任务具备了「挂了能重启、节点数变了能自适应、不用手动配 MASTER_ADDR 也能跑起来」的能力。
这份资源适合两类人：一是已经在用torchrun但没搞清它内部机制的工程师，二是正准备从单机多卡迁移到多机多卡、被容错和重启问题折磨过的从业者。如果你还在用python -m torch.distributed.launch这种老写法，那更值得往下看——新旧方案在容错能力上差了一个量级。
2. 从 torchrun 到 Elastic：核心组件与启动流程拆解
2.1 为什么需要 Elastic：静态分布式训练的四个硬伤
传统的torch.distributed.launch是静态的。启动前你必须确定好节点数、每节点进程数、MASTER_ADDR 和 MASTER_PORT，一旦确定就不能变。这带来四个问题：
第一，节点故障无法恢复。任何一个 worker 进程挂掉，整个任务就卡在集合通信上，最终超时退出，之前训练的 checkpoint 如果没存好就白跑了。
第二，节点数变化需要手动干预。你想从 4 机扩到 8 机，得停任务、改配置、重新启动，没法在线伸缩。
第三，RANK 分配靠环境变量硬编码。多机场景下你得手动给每台机器设RANK，容易出错，尤其是节点重启后 IP 变了，整个映射就乱了。
第四，缺少统一的错误处理。某个进程 OOM 了，其他进程不知道，还在那等 allreduce，最后大家一起超时。
PyTorch Elastic 的设计目标就是把这四个问题一次性解决。它的核心思路是：引入一个「agent」进程作为每个节点的本地管理者，agent 之间通过共享的 rendezvous 后端（通常是 etcd 或 c10d）交换信息，动态分配 rank，监控 worker 状态，失败时按策略重启。
2.2 核心组件：agent、rendezvous、run 三层结构
Elastic 的架构可以拆成三层：
Agent 层：每个节点上跑一个LocalElasticAgent，它负责启动、监控、终止本节点的 worker 进程。agent 会记录每个 worker 的状态，当某个 worker 失败时，agent 根据配置的max_restarts决定是否重启整个 worker group。
Rendezvous 层：这是节点发现和 rank 分配的核心。所有 agent 通过 rendezvous 后端注册自己，后端维护一个「已加入节点」的列表。当达到min_nodes时，rendezvous 会生成一个 rank 分配方案，把 global rank 映射到具体节点和本地进程。常见的后端有etcd（生产推荐）、c10d（单机或小规模）、static（固定配置）。
Run 层：这是用户入口，也就是torchrun命令。它解析参数、构造 agent、启动 rendezvous，然后把控制权交给 agent。
用一句话概括流程：torchrun启动 agent → agent 加入 rendezvous → rendezvous 分配 rank → agent 启动 worker 进程 → worker 通过init_process_group建立通信 → 训练开始 → 失败时 agent 检测并重启。
2.3 最小可跑示例：单机 4 卡 torchrun 启动
先看一个能直接抄的启动命令。假设你有一个train.py，里面用DistributedDataParallel包装了模型：
# 单机 4 卡，使用 c10d 后端（默认），standalone 模式 torchrun \ --standalone \ --nnodes=1 \ --nproc_per_node=4 \ --max_restarts=3 \ --rdzv_backend=c10d \ --rdzv_endpoint=localhost:29400 \ train.py --batch_size 32 --epochs 10参数说明：
--standalone：单机模式的快捷写法，等价于--nnodes=1 --rdzv_backend=c10d --rdzv_endpoint=localhost:0，端口自动选。--nnodes：节点数，可以是固定值如2，也可以是范围如2:4表示弹性伸缩。--nproc_per_node：每节点进程数，通常等于 GPU 数。--max_restarts：worker group 最大重启次数，超过就放弃。--rdzv_backend：rendezvous 后端，单机用c10d，多机生产建议etcd。--rdzv_endpoint：rendezvous 端点地址。
在train.py里，你不需要手动设RANK，torchrun 会通过环境变量注入：
import os import torch import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP def main(): # torchrun 自动注入的环境变量 rank = int(os.environ["RANK"]) # global rank local_rank = int(os.environ["LOCAL_RANK"]) # 节点内 rank world_size = int(os.environ["WORLD_SIZE"]) # 总进程数 # 初始化进程组，使用 env:// 方式读取环境变量 dist.init_process_group(backend="nccl") # 绑定当前进程到对应 GPU torch.cuda.set_device(local_rank) device = torch.device("cuda", local_rank) model = MyModel().to(device) model = DDP(model, device_ids=[local_rank]) # 只在 rank 0 上打印和保存 if rank == 0: print(f"World size: {world_size}, starting training...") for epoch in range(10): train_one_epoch(model, device) # 保存 checkpoint 时带上 rank 信息，方便恢复 if rank == 0: torch.save({ "epoch": epoch, "model": model.state_dict(), "world_size": world_size, }, f"ckpt_epoch{epoch}.pt") dist.destroy_process_group() if __name__ == "__main__": main()这段代码的关键点：init_process_group不传init_method，默认走env://，也就是从环境变量读MASTER_ADDR、MASTER_PORT、RANK、WORLD_SIZE。这些变量由 torchrun 的 agent 在启动 worker 时设置好，你不需要手动 export。
2.4 多机启动：etcd 后端与 rank 自动分配
多机场景下，c10d后端不够用，因为它没有持久化能力，节点重启后状态丢失。生产环境推荐用etcd作为 rendezvous 后端。启动方式如下：
# 节点 0（通常作为 master） torchrun \ --nnodes=2:4 \ --nproc_per_node=8 \ --max_restarts=5 \ --rdzv_backend=etcd \ --rdzv_endpoint=10.0.0.1:2379 \ --rdzv_id=my_job_001 \ train.py --batch_size 64 # 节点 1、2、3 执行同样的命令，不需要改任何参数 torchrun \ --nnodes=2:4 \ --nproc_per_node=8 \ --max_restarts=5 \ --rdzv_backend=etcd \ --rdzv_endpoint=10.0.0.1:2379 \ --rdzv_id=my_job_001 \ train.py --batch_size 64注意几个细节：
--nnodes=2:4表示最少 2 个节点、最多 4 个节点。rendezvous 会等到至少 2 个节点加入后才开始分配 rank，后续节点可以动态加入。--rdzv_id是任务唯一标识，同一个任务的所有节点必须用同一个 id，不同任务用不同 id 避免串扰。--rdzv_endpoint指向 etcd 服务地址，所有节点必须能访问。- 所有节点的启动命令完全一致，不需要手动指定
RANK或MASTER_ADDR，这是 Elastic 相比老方案最大的便利。
etcd 的部署可以用 Docker 快速起一个：
docker run -d \ --name etcd \ -p 2379:2379 \ -p 2380:2380 \ quay.io/coreos/etcd:v3.5.0 \ /usr/local/bin/etcd \ --advertise-client-urls http://0.0.0.0:2379 \ --listen-client-urls http://0.0.0.0:2379生产环境建议用 3 节点 etcd 集群，避免单点故障。单节点 etcd 挂了，整个 rendezvous 就失效了，所有训练任务都会卡住。
3. 弹性伸缩与容错恢复：checkpoint 怎么写才不白跑
3.1 弹性伸缩的触发条件与 rank 重分配
Elastic 的「弹性」体现在两个方向：扩容和缩容。扩容是指训练过程中有新节点加入，rendezvous 检测到节点数变化后，会触发一次 re-rendezvous，重新分配 rank，然后重启 worker group。缩容是指节点掉线后，如果剩余节点数仍大于等于min_nodes，任务会继续跑，只是 world_size 变小了。
这里有个关键点：每次 re-rendezvous 都会导致 worker group 重启，也就是说训练进程会重新启动，你需要从 checkpoint 恢复。所以弹性伸缩不是「无缝」的，而是「自动重启 + 自动恢复」。如果你的代码没有正确的 checkpoint 逻辑，扩容后就是从零开始训练，那就白跑了。
触发 re-rendezvous 的条件：
- 新节点加入 rendezvous，且当前节点数小于
max_nodes。 - 某个节点失联超过
rdzv_timeout（默认 900 秒）。 - worker group 内某个 worker 失败，且
max_restarts未耗尽。
3.2 checkpoint 设计：必须保存 world_size 和 epoch
很多人的 checkpoint 只存model.state_dict()，这在弹性场景下是不够的。因为重启后 world_size 可能变了，DataLoader 的 shard 策略、学习率的 warmup 步数、optimizer 的状态都可能需要调整。一个健壮的 checkpoint 应该包含：
import torch import os def save_checkpoint(model, optimizer, epoch, global_step, save_dir, rank): """只在 rank 0 上保存，避免多进程写冲突""" if rank != 0: return os.makedirs(save_dir, exist_ok=True) ckpt = { "epoch": epoch, "global_step": global_step, "model_state": model.state_dict(), "optimizer_state": optimizer.state_dict(), "world_size": int(os.environ["WORLD_SIZE"]), "rng_state": torch.get_rng_state(), } path = os.path.join(save_dir, f"ckpt_step{global_step}.pt") torch.save(ckpt, path) # 同时写一个 latest 软链接，方便恢复时找最新 latest = os.path.join(save_dir, "latest.pt") if os.path.exists(latest): os.remove(latest) os.symlink(path, latest) print(f"[rank 0] saved checkpoint to {path}") def load_checkpoint(model, optimizer, save_dir, device): """所有 rank 都要加载，保证状态一致""" latest = os.path.join(save_dir, "latest.pt") if not os.path.exists(latest): return 0, 0 ckpt = torch.load(latest, map_location=device) model.load_state_dict(ckpt["model_state"]) optimizer.load_state_dict(ckpt["optimizer_state"]) torch.set_rng_state(ckpt["rng_state"]) return ckpt["epoch"], ckpt["global_step"]逻辑说明：world_size存下来是为了在恢复时判断是否需要调整 DataLoader 的num_replicas。global_step用于恢复学习率调度器。rng_state保证数据增强的随机性可复现。latest.pt软链接让恢复逻辑不用遍历目录找最新文件。
参数说明：save_dir建议放在共享存储上（NFS、S3、HDFS），因为重启后 rank 0 可能落在不同节点上。如果放在本地盘，节点挂了 checkpoint 就丢了。
3.3 恢复逻辑：从 latest checkpoint 续训
在train.py的main函数开头加入恢复逻辑：
def main(): rank = int(os.environ["RANK"]) local_rank = int(os.environ["LOCAL_RANK"]) world_size = int(os.environ["WORLD_SIZE"]) dist.init_process_group(backend="nccl") torch.cuda.set_device(local_rank) device = torch.device("cuda", local_rank) model = MyModel().to(device) optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4) model = DDP(model, device_ids=[local_rank]) save_dir = "/shared/checkpoints/my_job" start_epoch, start_step = load_checkpoint(model, optimizer, save_dir, device) if rank == 0: print(f"Resuming from epoch {start_epoch}, step {start_step}, world_size={world_size}") for epoch in range(start_epoch, 10): # 如果 world_size 变了，需要重建 DataLoader 的 sampler sampler = DistributedSampler( dataset, num_replicas=world_size, rank=rank, shuffle=True, ) loader = DataLoader(dataset, sampler=sampler, batch_size=32) for step, batch in enumerate(loader): if epoch == start_epoch and step < start_step: continue # 跳过已训练的 step train_step(model, optimizer, batch, device) if step % 100 == 0: save_checkpoint(model, optimizer, epoch, step, save_dir, rank) dist.destroy_process_group()这里有个容易翻车的点：DistributedSampler的num_replicas必须和当前world_size一致。如果扩容后 world_size 从 8 变成 16，但 sampler 还是按 8 来分片，数据就会重复或遗漏。所以每次 re-rendezvous 后都要重建 sampler。
3.4 用 max_restarts 和 rdzv_timeout 控制恢复行为
--max_restarts控制 worker group 最多重启几次。设太小，节点抖动一下就放弃；设太大，遇到必然失败的错误（比如代码 bug）会无限重启浪费资源。经验值：训练任务设 3 到 5 次，调试阶段设 0 或 1 次快速失败。
--rdzv_timeout控制 rendezvous 等待节点加入的超时时间，默认 900 秒。如果你的节点启动慢（比如要拉镜像、加载数据），可以调大到 1800 秒。但设太大也有问题：某个节点真挂了，你要等很久才能触发缩容。
还有一个隐藏参数--rdzv_conf，可以传一些高级配置：
torchrun \ --rdzv_backend=etcd \ --rdzv_endpoint=10.0.0.1:2379 \ --rdzv_id=my_job \ --rdzv_conf="timeout=1200,join_timeout=600,last_call_timeout=60" \ train.pytimeout：rendezvous 整体超时。join_timeout：单个节点加入的超时。last_call_timeout：最后一个节点加入后，等待其他节点确认的时间。
这些参数在节点数多、网络不稳定的环境下需要调优，默认值不一定适合你的集群。
4. 避坑与排查：那些让任务卡死的配置细节
4.1 现象：任务卡在 rendezvous 不启动，日志显示 waiting for nodes
原因：最常见的是--rdzv_id不一致，或者--rdzv_endpoint指向的 etcd 不可达。另一个常见原因是--nnodes设成了固定值如4，但实际只有 3 个节点在跑，rendezvous 会一直等第 4 个节点。
解决：先检查所有节点的启动命令是否完全一致，特别是rdzv_id和rdzv_endpoint。然后用etcdctl检查 etcd 里是否有对应的 key：
etcdctl --endpoints=10.0.0.1:2379 get --prefix /torch/elastic/如果 key 存在但节点没起来，检查--nnodes是否设成了范围2:4而不是固定值。固定值意味着必须凑齐才能启动，范围值则达到最小值就能启动。
4.2 现象：训练中途某个节点挂了，任务没有自动缩容而是直接退出
原因：--max_restarts设成了 0，或者--nnodes设的是固定值而不是范围。固定值模式下，节点数少于nnodes时 rendezvous 会失败，agent 直接退出。
解决：把--nnodes改成范围，比如2:4，并把--max_restarts设为 3 以上。同时确认 checkpoint 保存路径是共享存储，否则重启后找不到 checkpoint 还是会从零开始。
4.3 现象：扩容后 loss 突然跳变，像是重新开始训练
原因：checkpoint 里没存global_step或epoch，恢复时从 0 开始。或者 DataLoader 的 sampler 没有按新 world_size 重建，导致数据分片错乱。
解决：按 3.2 节的 checkpoint 结构保存epoch和global_step，恢复时读取并跳过已训练的 step。每次 re-rendezvous 后重建DistributedSampler，确保num_replicas等于当前world_size。
4.4 现象：NCCL 通信超时，日志报 unhandled system error
原因：多机场景下 NCCL 需要正确的网络接口配置。如果节点有多个网卡，NCCL 可能选错了网卡，导致通信走不通。另一个原因是防火墙挡住了 NCCL 使用的端口范围。
解决：设置NCCL_SOCKET_IFNAME指定网卡：
export NCCL_SOCKET_IFNAME=eth0 export NCCL_IB_DISABLE=1 # 如果没有 InfiniBand，禁用 IB 避免报错 export NCCL_DEBUG=INFO # 调试时打开，看 NCCL 选了哪个网卡如果是容器环境，还要确保--shm-size足够大，NCCL 需要共享内存做进程间通信。建议至少--shm-size=8g。
4.5 现象：torchrun 启动后 worker 进程立刻退出，没有错误日志
原因：通常是train.py里有语法错误或 import 失败，但错误被 agent 吞掉了。agent 默认只打印自己的日志，worker 的 stderr 可能没被捕获。
解决：加--log-dir参数把日志写到文件：
torchrun \ --log-dir=/tmp/torchrun_logs \ --redirects=3 \ --tee=3 \ train.py--redirects=3表示同时重定向 stdout 和 stderr 到文件，--tee=3表示同时输出到控制台和文件。这样 worker 的报错就能看到了。
5. 进阶技巧：用 torchelastic 的 API 做自定义容错策略
5.1 直接调用 elastic agent API
torchrun是命令行入口，但 Elastic 也提供了 Python API，可以在代码里直接控制 agent 行为。比如你想在 worker 失败时执行自定义清理逻辑：
from torch.distributed.elastic.agent.server.api import LocalElasticAgent from torch.distributed.elastic.rendezvous import RendezvousParameters # 构造 rendezvous 参数 rdzv_params = RendezvousParameters( backend="etcd", endpoint="10.0.0.1:2379", run_id="my_job_001", min_nodes=2, max_nodes=4, ) # 自定义 agent，重写 _invoke_run 方法插入清理逻辑 class MyAgent(LocalElasticAgent): def _invoke_run(self, role): try: super()._invoke_run(role) except Exception as e: # 失败时上传日志到对象存储 upload_logs_to_s3("/tmp/torchrun_logs") raise agent = MyAgent( name="my_trainer", start_method="spawn", rdzv_handler=rdzv_params, max_restarts=3, ) agent.run()这个用法比较底层，适合需要深度定制容错策略的场景。大部分情况下torchrun命令行就够了，不建议一上来就手写 agent。
5.2 用 elastic 的 metrics 做监控
Elastic agent 会暴露一些内部状态，可以通过torch.distributed.elastic.metrics模块打点。比如记录每次重启的次数和原因：
from torch.distributed.elastic.metrics import put_metric # 在 worker 失败时打点 put_metric("worker_restart_count", 1, "counter") put_metric("rendezvous_duration_seconds", elapsed, "timer")这些 metrics 默认是 no-op，需要配置TORCHELASTIC_ENABLE_METRICS=1环境变量才会生效。生产环境可以对接 Prometheus 做告警。
5.3 验证弹性是否真的生效：手动 kill 一个 worker
最后给一个验证方法。启动一个多机任务后，手动 kill 掉某个 worker 进程，观察 agent 是否自动重启：
# 在节点 1 上找到 worker 进程 ps aux | grep train.py | grep -v grep # kill 掉其中一个 worker（不要 kill agent 进程） kill -9 <worker_pid> # 观察节点 0 的日志，应该看到类似输出： # [agent] worker 3 failed, restarting worker group (restart 1/3) # [agent] re-rendezvous triggered, new world_size=15如果日志里出现了restarting worker group和re-rendezvous，说明弹性生效了。如果任务直接退出，检查--max_restarts和--nnodes配置。
从那以后我每次上线多机训练任务前，都会先手动 kill 一个 worker 验证容错链路，确认 checkpoint 能正确恢复、world_size 能自适应调整，才敢跑长任务。这个习惯帮我省了至少两次通宵重跑的代价。希望帮到你。
本文还有配套的精品资源，点击获取
