简介这份资源围绕 PyTorch ElasticTorchElastic展开面向需要以容错、弹性方式启动分布式 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 \ --nnodes1 \ --nproc_per_node4 \ --max_restarts3 \ --rdzv_backendc10d \ --rdzv_endpointlocalhost:29400 \ train.py --batch_size 32 --epochs 10参数说明--standalone单机模式的快捷写法等价于--nnodes1 --rdzv_backendc10d --rdzv_endpointlocalhost:0端口自动选。--nnodes节点数可以是固定值如2也可以是范围如2:4表示弹性伸缩。--nproc_per_node每节点进程数通常等于 GPU 数。--max_restartsworker group 最大重启次数超过就放弃。--rdzv_backendrendezvous 后端单机用c10d多机生产建议etcd。--rdzv_endpointrendezvous 端点地址。在train.py里你不需要手动设RANKtorchrun 会通过环境变量注入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(backendnccl) # 绑定当前进程到对应 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(fWorld 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, }, fckpt_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 \ --nnodes2:4 \ --nproc_per_node8 \ --max_restarts5 \ --rdzv_backendetcd \ --rdzv_endpoint10.0.0.1:2379 \ --rdzv_idmy_job_001 \ train.py --batch_size 64 # 节点 1、2、3 执行同样的命令不需要改任何参数 torchrun \ --nnodes2:4 \ --nproc_per_node8 \ --max_restarts5 \ --rdzv_backendetcd \ --rdzv_endpoint10.0.0.1:2379 \ --rdzv_idmy_job_001 \ train.py --batch_size 64注意几个细节--nnodes2: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_okTrue) 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, fckpt_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_locationdevice) 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(backendnccl) torch.cuda.set_device(local_rank) device torch.device(cuda, local_rank) model MyModel().to(device) optimizer torch.optim.AdamW(model.parameters(), lr1e-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(fResuming 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_replicasworld_size, rankrank, shuffleTrue, ) loader DataLoader(dataset, samplersampler, batch_size32) 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_backendetcd \ --rdzv_endpoint10.0.0.1:2379 \ --rdzv_idmy_job \ --rdzv_conftimeout1200,join_timeout600,last_call_timeout60 \ train.pytimeoutrendezvous 整体超时。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 里是否有对应的 keyetcdctl --endpoints10.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_IFNAMEeth0 export NCCL_IB_DISABLE1 # 如果没有 InfiniBand禁用 IB 避免报错 export NCCL_DEBUGINFO # 调试时打开看 NCCL 选了哪个网卡如果是容器环境还要确保--shm-size足够大NCCL 需要共享内存做进程间通信。建议至少--shm-size8g。4.5 现象torchrun 启动后 worker 进程立刻退出没有错误日志原因通常是train.py里有语法错误或 import 失败但错误被 agent 吞掉了。agent 默认只打印自己的日志worker 的 stderr 可能没被捕获。解决加--log-dir参数把日志写到文件torchrun \ --log-dir/tmp/torchrun_logs \ --redirects3 \ --tee3 \ train.py--redirects3表示同时重定向 stdout 和 stderr 到文件--tee3表示同时输出到控制台和文件。这样 worker 的报错就能看到了。5. 进阶技巧用 torchelastic 的 API 做自定义容错策略5.1 直接调用 elastic agent APItorchrun是命令行入口但 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( backendetcd, endpoint10.0.0.1:2379, run_idmy_job_001, min_nodes2, max_nodes4, ) # 自定义 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( namemy_trainer, start_methodspawn, rdzv_handlerrdzv_params, max_restarts3, ) 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_METRICS1环境变量才会生效。生产环境可以对接 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_size15如果日志里出现了restarting worker group和re-rendezvous说明弹性生效了。如果任务直接退出检查--max_restarts和--nnodes配置。从那以后我每次上线多机训练任务前都会先手动 kill 一个 worker 验证容错链路确认 checkpoint 能正确恢复、world_size 能自适应调整才敢跑长任务。这个习惯帮我省了至少两次通宵重跑的代价。希望帮到你。本文还有配套的精品资源点击获取
