【AIaaS 全栈架构师】第 32 篇:弹性训练与容错------让训练任务永不中断
系列定位:AIaaS 全栈架构师教程,技术栈以 Go 为主。本篇聚焦分布式训练的弹性与容错------PyTorch Elastic 如何让训练在 Worker 动态增减、节点抢占、网络抖动下永不中断,以及 Checkpoint 持久化、NCCL 超时检测、Spot 实例训练的完整工程实践。
本篇你将学到
- 为什么传统的
torch.distributed.launch无法应对节点抢占,必须引入 Elastic 训练 torch.distributed.elastic的 Rendezvous 会合机制与状态机- 动态 Worker 增减的完整流程:scale up / scale down 如何不破坏训练一致性
- Checkpoint 的正确保存策略:频率、原子性、存储后端、恢复路径
- NCCL 超时、僵尸进程、死锁的诊断与治理
- Spot / 抢占式实例训练的成本与稳定性权衡,以及生产落地参数
- 用 Go 实现 Checkpoint 管理与训练恢复的辅助服务
前两篇我们用 Training Operator 和 Ray 解决了「怎么提交分布式训练任务」。但它们都隐含一个假设:Worker 数量在任务运行期间是固定的。在真实生产环境里,这个假设极其脆弱------GPU 节点会被抢占、网络会抖动、NCCL 会卡死、磁盘会写满。一个跑了几十小时的训练任务,如果因为一次抢占就从头再来,成本和周期都不可接受。这一篇我们就来彻底拆解弹性训练,让训练任务真正做到「永不中断」。
一、为什么需要弹性训练
1.1 固定 World Size 的致命缺陷
先看传统启动方式(torch.distributed.launch / torchrun 旧版)的问题。一个 8 节点 64 卡的训练任务,用 --nnodes=8 --nproc_per_node=8 启动:
节点1: rank 0-7 节点2: rank 8-15 ... 节点8: rank 56-63
所有节点通过 init_process_group 互相发现,world_size=64
这种启动方式有三个致命问题:
问题一:单点失败全盘崩溃
只要 64 个 rank 中任何一个挂掉(节点抢占、OOM、NCCL 超时),init_process_group 就会超时,整个进程组解体。PyTorch 原生的 destroy_process_group 会让所有 rank 一起退出。
问题二:无法动态伸缩
训练跑到一半,如果集群有 4 张 A100 空出来了,你没法把它「加」进正在运行的训练任务。World Size 是启动时固定的,改了就得重启。
问题三:启动时整批等待
64 个 rank 必须全部到齐才能开始训练。先起的节点干等后起的,如果有一个节点镜像拉取慢,所有节点都空转。
这三个问题在大规模、长时间训练中是致命的。Llama 3 70B 级别的训练动辄几周,期间节点抢占是家常便饭,固定 World Size 的方案根本撑不住。
1.2 Spot 实例的成本诱惑
弹性训练的另一个驱动力是成本。云厂商的 Spot / 抢占式实例通常比按需便宜 60-80%:
| 实例类型 | 按需价格 | Spot 价格 | 折扣 |
|---|---|---|---|
| p4d.24xlarge (8×A100) | $32.77/h | $9.83/h | 70% |
| p5.48xlarge (8×H100) | $98.32/h | $29.50/h | 70% |
一个 64 卡训练任务跑一周,按需成本约 35,000,Spot 只要 10,500。省下的 $24,500 足够再训一次。但 Spot 实例随时可能被回收(通常提前 2 分钟通知),没有弹性训练,这个折扣就拿不到。
1.3 弹性训练的核心承诺
弹性训练(Elastic Training)要解决的就是:在 Worker 动态增减的前提下,保证训练正确性,并尽量减少进度损失。它的核心承诺:
- 动态成员:Worker 可以中途加入或退出,不重启整个任务。
- 状态恢复:新加入或恢复的 Worker 从最近的 Checkpoint 继续,不从头训练。
- 优雅降级:Worker 减少时,剩余 Worker 继续训练,world_size 动态调整。
- 最终一致 :无论经历多少次成员变更,训练收敛性与固定 World Size 等价(数学上要求使用
join机制处理梯度缺失)。
PyTorch 从 1.10 开始提供 torch.distributed.elastic(简称 Torchrun Elastic),是这个领域的事实标准。
二、PyTorch Elastic 架构
2.1 核心概念
Torchrun Elastic 引入三个核心概念:
| 概念 | 作用 | 对应实现 |
|---|---|---|
| Rendezvous | 会合机制,让动态 Worker 互相发现并达成一致 | RendezvousHandler |
| ElasticAgent | 每个节点上的监督进程,管理本地 Worker | LocalElasticAgent |
| State | 训练状态(step / epoch / RNG),用于恢复 | 用户代码 + Checkpoint |
理解的关键:Rendezvous 不是固定的 MASTER_ADDR,而是一个动态的「注册中心」。Worker 随时可以来注册,凑够一批就组成一个新的进程组开始(或继续)训练。
2.2 整体架构
#mermaid-svg-67Oggnjdoyph2L9z{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-67Oggnjdoyph2L9z .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-67Oggnjdoyph2L9z .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-67Oggnjdoyph2L9z .error-icon{fill:#552222;}#mermaid-svg-67Oggnjdoyph2L9z .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-67Oggnjdoyph2L9z .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-67Oggnjdoyph2L9z .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-67Oggnjdoyph2L9z .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-67Oggnjdoyph2L9z .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-67Oggnjdoyph2L9z .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-67Oggnjdoyph2L9z .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-67Oggnjdoyph2L9z .marker{fill:#333333;stroke:#333333;}#mermaid-svg-67Oggnjdoyph2L9z .marker.cross{stroke:#333333;}#mermaid-svg-67Oggnjdoyph2L9z svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-67Oggnjdoyph2L9z p{margin:0;}#mermaid-svg-67Oggnjdoyph2L9z .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-67Oggnjdoyph2L9z .cluster-label text{fill:#333;}#mermaid-svg-67Oggnjdoyph2L9z .cluster-label span{color:#333;}#mermaid-svg-67Oggnjdoyph2L9z .cluster-label span p{background-color:transparent;}#mermaid-svg-67Oggnjdoyph2L9z .label text,#mermaid-svg-67Oggnjdoyph2L9z span{fill:#333;color:#333;}#mermaid-svg-67Oggnjdoyph2L9z .node rect,#mermaid-svg-67Oggnjdoyph2L9z .node circle,#mermaid-svg-67Oggnjdoyph2L9z .node ellipse,#mermaid-svg-67Oggnjdoyph2L9z .node polygon,#mermaid-svg-67Oggnjdoyph2L9z .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-67Oggnjdoyph2L9z .rough-node .label text,#mermaid-svg-67Oggnjdoyph2L9z .node .label text,#mermaid-svg-67Oggnjdoyph2L9z .image-shape .label,#mermaid-svg-67Oggnjdoyph2L9z .icon-shape .label{text-anchor:middle;}#mermaid-svg-67Oggnjdoyph2L9z .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-67Oggnjdoyph2L9z .rough-node .label,#mermaid-svg-67Oggnjdoyph2L9z .node .label,#mermaid-svg-67Oggnjdoyph2L9z .image-shape .label,#mermaid-svg-67Oggnjdoyph2L9z .icon-shape .label{text-align:center;}#mermaid-svg-67Oggnjdoyph2L9z .node.clickable{cursor:pointer;}#mermaid-svg-67Oggnjdoyph2L9z .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-67Oggnjdoyph2L9z .arrowheadPath{fill:#333333;}#mermaid-svg-67Oggnjdoyph2L9z .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-67Oggnjdoyph2L9z .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-67Oggnjdoyph2L9z .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-67Oggnjdoyph2L9z .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-67Oggnjdoyph2L9z .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-67Oggnjdoyph2L9z .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-67Oggnjdoyph2L9z .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-67Oggnjdoyph2L9z .cluster text{fill:#333;}#mermaid-svg-67Oggnjdoyph2L9z .cluster span{color:#333;}#mermaid-svg-67Oggnjdoyph2L9z div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-67Oggnjdoyph2L9z .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-67Oggnjdoyph2L9z rect.text{fill:none;stroke-width:0;}#mermaid-svg-67Oggnjdoyph2L9z .icon-shape,#mermaid-svg-67Oggnjdoyph2L9z .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-67Oggnjdoyph2L9z .icon-shape p,#mermaid-svg-67Oggnjdoyph2L9z .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-67Oggnjdoyph2L9z .icon-shape .label rect,#mermaid-svg-67Oggnjdoyph2L9z .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-67Oggnjdoyph2L9z .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-67Oggnjdoyph2L9z .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-67Oggnjdoyph2L9z :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} 节点 3 (后加入)
节点 2
节点 1
Rendezvous 后端
(C10d / etcd / 文件)
注册/心跳/会合
注册/心跳/会合
动态加入
NCCL 进程组
NCCL 进程组
重组后加入
Key-Value Store
ElasticAgent
Worker rank 0
Worker rank 1
ElasticAgent
Worker rank 2
Worker rank 3
ElasticAgent
Worker rank 4
Worker rank 5
每个节点跑一个 ElasticAgent,它负责:
- 向 Rendezvous 后端注册本节点的 Worker 数量。
- 等待凑够
min_nodes后,分配 rank,初始化进程组。 - 监督本地 Worker 进程,崩溃则重启,节点退出则通知 Rendezvous。
- 加载 Checkpoint,传递恢复状态给 Worker。
Rendezvous 后端有三种实现:
- C10d (推荐):基于一个常驻的
rendezvous server(通常跑在 Head 节点),内置在 PyTorch 中,无需额外依赖。 - etcd:用 etcd 做一致性存储,适合大规模、多任务共享。
- File:基于共享文件系统,仅用于单机调试。
2.3 Rendezvous 状态机
Rendezvous 的核心是一个分布式状态机,所有 ElasticAgent 通过它达成一致:
#mermaid-svg-OQxljJonmWzpepHu{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-OQxljJonmWzpepHu .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-OQxljJonmWzpepHu .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-OQxljJonmWzpepHu .error-icon{fill:#552222;}#mermaid-svg-OQxljJonmWzpepHu .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-OQxljJonmWzpepHu .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-OQxljJonmWzpepHu .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-OQxljJonmWzpepHu .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-OQxljJonmWzpepHu .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-OQxljJonmWzpepHu .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-OQxljJonmWzpepHu .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-OQxljJonmWzpepHu .marker{fill:#333333;stroke:#333333;}#mermaid-svg-OQxljJonmWzpepHu .marker.cross{stroke:#333333;}#mermaid-svg-OQxljJonmWzpepHu svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-OQxljJonmWzpepHu p{margin:0;}#mermaid-svg-OQxljJonmWzpepHu defs #statediagram-barbEnd{fill:#333333;stroke:#333333;}#mermaid-svg-OQxljJonmWzpepHu g.stateGroup text{fill:#9370DB;stroke:none;font-size:10px;}#mermaid-svg-OQxljJonmWzpepHu g.stateGroup text{fill:#333;stroke:none;font-size:10px;}#mermaid-svg-OQxljJonmWzpepHu g.stateGroup .state-title{font-weight:bolder;fill:#131300;}#mermaid-svg-OQxljJonmWzpepHu g.stateGroup rect{fill:#ECECFF;stroke:#9370DB;}#mermaid-svg-OQxljJonmWzpepHu g.stateGroup line{stroke:#333333;stroke-width:1;}#mermaid-svg-OQxljJonmWzpepHu .transition{stroke:#333333;stroke-width:1;fill:none;}#mermaid-svg-OQxljJonmWzpepHu .stateGroup .composit{fill:white;border-bottom:1px;}#mermaid-svg-OQxljJonmWzpepHu .stateGroup .alt-composit{fill:#e0e0e0;border-bottom:1px;}#mermaid-svg-OQxljJonmWzpepHu .state-note{stroke:#aaaa33;fill:#fff5ad;}#mermaid-svg-OQxljJonmWzpepHu .state-note text{fill:black;stroke:none;font-size:10px;}#mermaid-svg-OQxljJonmWzpepHu .stateLabel .box{stroke:none;stroke-width:0;fill:#ECECFF;opacity:0.5;}#mermaid-svg-OQxljJonmWzpepHu .edgeLabel .label rect{fill:#ECECFF;opacity:0.5;}#mermaid-svg-OQxljJonmWzpepHu .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-OQxljJonmWzpepHu .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-OQxljJonmWzpepHu .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-OQxljJonmWzpepHu .edgeLabel .label text{fill:#333;}#mermaid-svg-OQxljJonmWzpepHu .label div .edgeLabel{color:#333;}#mermaid-svg-OQxljJonmWzpepHu .stateLabel text{fill:#131300;font-size:10px;font-weight:bold;}#mermaid-svg-OQxljJonmWzpepHu .node circle.state-start{fill:#333333;stroke:#333333;}#mermaid-svg-OQxljJonmWzpepHu .node .fork-join{fill:#333333;stroke:#333333;}#mermaid-svg-OQxljJonmWzpepHu .node circle.state-end{fill:#9370DB;stroke:white;stroke-width:1.5;}#mermaid-svg-OQxljJonmWzpepHu .end-state-inner{fill:white;stroke-width:1.5;}#mermaid-svg-OQxljJonmWzpepHu .node rect{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-OQxljJonmWzpepHu .node polygon{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-OQxljJonmWzpepHu #statediagram-barbEnd{fill:#333333;}#mermaid-svg-OQxljJonmWzpepHu .statediagram-cluster rect{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-OQxljJonmWzpepHu .cluster-label,#mermaid-svg-OQxljJonmWzpepHu .nodeLabel{color:#131300;}#mermaid-svg-OQxljJonmWzpepHu .statediagram-cluster rect.outer{rx:5px;ry:5px;}#mermaid-svg-OQxljJonmWzpepHu .statediagram-state .divider{stroke:#9370DB;}#mermaid-svg-OQxljJonmWzpepHu .statediagram-state .title-state{rx:5px;ry:5px;}#mermaid-svg-OQxljJonmWzpepHu .statediagram-cluster.statediagram-cluster .inner{fill:white;}#mermaid-svg-OQxljJonmWzpepHu .statediagram-cluster.statediagram-cluster-alt .inner{fill:#f0f0f0;}#mermaid-svg-OQxljJonmWzpepHu .statediagram-cluster .inner{rx:0;ry:0;}#mermaid-svg-OQxljJonmWzpepHu .statediagram-state rect.basic{rx:5px;ry:5px;}#mermaid-svg-OQxljJonmWzpepHu .statediagram-state rect.divider{stroke-dasharray:10,10;fill:#f0f0f0;}#mermaid-svg-OQxljJonmWzpepHu .note-edge{stroke-dasharray:5;}#mermaid-svg-OQxljJonmWzpepHu .statediagram-note rect{fill:#fff5ad;stroke:#aaaa33;stroke-width:1px;rx:0;ry:0;}#mermaid-svg-OQxljJonmWzpepHu .statediagram-note rect{fill:#fff5ad;stroke:#aaaa33;stroke-width:1px;rx:0;ry:0;}#mermaid-svg-OQxljJonmWzpepHu .statediagram-note text{fill:black;}#mermaid-svg-OQxljJonmWzpepHu .statediagram-note .nodeLabel{color:black;}#mermaid-svg-OQxljJonmWzpepHu .statediagram .edgeLabel{color:red;}#mermaid-svg-OQxljJonmWzpepHu #dependencyStart,#mermaid-svg-OQxljJonmWzpepHu #dependencyEnd{fill:#333333;stroke:#333333;stroke-width:1;}#mermaid-svg-OQxljJonmWzpepHu .statediagramTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-OQxljJonmWzpepHu :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} Worker 启动
加入会合, 分配 ordinal
等待 min_workers
凑够 min_workers
且有闲超时
所有 Worker 确认
生成 world
分配 rank
训练开始
成员变更
(节点加入/退出)
超时/失败
JOIN
WAIT
CONFIRM
FINAL
RUNNABLE
min_workers = min_nodes
max_workers = max_nodes
超过 max_workers 的 Worker 等下一轮
状态含义:
- JOIN:Worker 向 Rendezvous 注册,声明自己存在。
- WAIT :等待凑够
min_nodes,同时有超时机制(timeout,默认 60s)。 - CONFIRM :凑够后,每个 Worker 确认参与。有
keep_alive_period(默认 5s)心跳,超时视为掉线。 - FINAL:所有确认的 Worker 组成最终 world。
- RUNNABLE:分配 rank(0 到 world_size-1),返回给 ElasticAgent,开始训练。
关键参数:
bash
torchrun \
--nnodes=2:8 \ # min:max,最少 2 节点最多 8 节点
--nproc_per_node=8 \ # 每节点 8 Worker
--rdzv_backend=c10d \ # 会合后端
--rdzv_endpoint=head:29500 \# 会合服务地址
--rdzv_id=job-abc123 \ # 任务唯一 ID
--max_restarts=3 \ # 最多重启 3 次
train.py
--nnodes=2:8 是弹性的灵魂:最少 2 节点就能开始训练,最多接受 8 节点。训练过程中节点数在 2-8 之间浮动,Rendezvous 自动重组。
2.4 成员变更触发重组
当某个节点崩溃或主动退出时,Rendezvous 检测到心跳超时(或收到显式退出通知),触发 re-rendezvous:
Rendezvous 节点3 (崩溃) 节点2 (存活) 节点1 (存活) Rendezvous 节点3 (崩溃) 节点2 (存活) 节点1 (存活) #mermaid-svg-w4cISKKhAAivbsYN{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-w4cISKKhAAivbsYN .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-w4cISKKhAAivbsYN .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-w4cISKKhAAivbsYN .error-icon{fill:#552222;}#mermaid-svg-w4cISKKhAAivbsYN .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-w4cISKKhAAivbsYN .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-w4cISKKhAAivbsYN .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-w4cISKKhAAivbsYN .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-w4cISKKhAAivbsYN .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-w4cISKKhAAivbsYN .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-w4cISKKhAAivbsYN .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-w4cISKKhAAivbsYN .marker{fill:#333333;stroke:#333333;}#mermaid-svg-w4cISKKhAAivbsYN .marker.cross{stroke:#333333;}#mermaid-svg-w4cISKKhAAivbsYN svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-w4cISKKhAAivbsYN p{margin:0;}#mermaid-svg-w4cISKKhAAivbsYN .actor{stroke:hsl(259.6261682243, 59.7765363128%, 87.9019607843%);fill:#ECECFF;}#mermaid-svg-w4cISKKhAAivbsYN text.actor>tspan{fill:black;stroke:none;}#mermaid-svg-w4cISKKhAAivbsYN .actor-line{stroke:hsl(259.6261682243, 59.7765363128%, 87.9019607843%);}#mermaid-svg-w4cISKKhAAivbsYN .innerArc{stroke-width:1.5;stroke-dasharray:none;}#mermaid-svg-w4cISKKhAAivbsYN .messageLine0{stroke-width:1.5;stroke-dasharray:none;stroke:#333;}#mermaid-svg-w4cISKKhAAivbsYN .messageLine1{stroke-width:1.5;stroke-dasharray:2,2;stroke:#333;}#mermaid-svg-w4cISKKhAAivbsYN #arrowhead path{fill:#333;stroke:#333;}#mermaid-svg-w4cISKKhAAivbsYN .sequenceNumber{fill:white;}#mermaid-svg-w4cISKKhAAivbsYN #sequencenumber{fill:#333;}#mermaid-svg-w4cISKKhAAivbsYN #crosshead path{fill:#333;stroke:#333;}#mermaid-svg-w4cISKKhAAivbsYN .messageText{fill:#333;stroke:none;}#mermaid-svg-w4cISKKhAAivbsYN .labelBox{stroke:hsl(259.6261682243, 59.7765363128%, 87.9019607843%);fill:#ECECFF;}#mermaid-svg-w4cISKKhAAivbsYN .labelText,#mermaid-svg-w4cISKKhAAivbsYN .labelText>tspan{fill:black;stroke:none;}#mermaid-svg-w4cISKKhAAivbsYN .loopText,#mermaid-svg-w4cISKKhAAivbsYN .loopText>tspan{fill:black;stroke:none;}#mermaid-svg-w4cISKKhAAivbsYN .loopLine{stroke-width:2px;stroke-dasharray:2,2;stroke:hsl(259.6261682243, 59.7765363128%, 87.9019607843%);fill:hsl(259.6261682243, 59.7765363128%, 87.9019607843%);}#mermaid-svg-w4cISKKhAAivbsYN .note{stroke:#aaaa33;fill:#fff5ad;}#mermaid-svg-w4cISKKhAAivbsYN .noteText,#mermaid-svg-w4cISKKhAAivbsYN .noteText>tspan{fill:black;stroke:none;}#mermaid-svg-w4cISKKhAAivbsYN .activation0{fill:#f4f4f4;stroke:#666;}#mermaid-svg-w4cISKKhAAivbsYN .activation1{fill:#f4f4f4;stroke:#666;}#mermaid-svg-w4cISKKhAAivbsYN .activation2{fill:#f4f4f4;stroke:#666;}#mermaid-svg-w4cISKKhAAivbsYN .actorPopupMenu{position:absolute;}#mermaid-svg-w4cISKKhAAivbsYN .actorPopupMenuPanel{position:absolute;fill:#ECECFF;box-shadow:0px 8px 16px 0px rgba(0,0,0,0.2);filter:drop-shadow(3px 5px 2px rgb(0 0 0 / 0.4));}#mermaid-svg-w4cISKKhAAivbsYN .actor-man line{stroke:hsl(259.6261682243, 59.7765363128%, 87.9019607843%);fill:#ECECFF;}#mermaid-svg-w4cISKKhAAivbsYN .actor-man circle,#mermaid-svg-w4cISKKhAAivbsYN line{stroke:hsl(259.6261682243, 59.7765363128%, 87.9019607843%);fill:#ECECFF;stroke-width:2px;}#mermaid-svg-w4cISKKhAAivbsYN :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} world_size=24, 训练中 重组为 world_size=16 OOM / 抢占 / 网络断 NCCL 超时 (rank 对端无响应) 请求 re-rendezvous 标记 N3 为 dead 通知: world 变更 通知: world 变更 重新加载 Checkpoint 重新加载 Checkpoint 新 world 就绪 新 world 就绪 开始新一轮训练 开始新一轮训练
注意几个细节:
- NCCL 超时是触发器:存活的 Worker 通过 NCCL 操作(AllReduce 等)感知到对端无响应,超时后主动请求重组。
- Checkpoint 是恢复依据:重组后所有 Worker 重新加载最近 Checkpoint,避免梯度不一致。
- Rank 重新分配:重组后 rank 0-15 重新映射到存活的两个节点,旧的 rank 16-23(崩溃节点的)被丢弃。
三、动态 Worker 增减
3.1 Scale Up:新 Worker 加入
新节点加入训练任务的流程:
- 新节点的 ElasticAgent 启动,向 Rendezvous 注册。
- Rendezvous 检查当前 world 是否已满(
max_nodes)。- 未满:等待当前训练 step 结束,触发重组。
- 已满:新节点排队,等下一轮有空位。
- 重组:新节点加入 world,重新分配 rank。
- 所有 Worker 加载最近 Checkpoint,继续训练。
代码层面,训练脚本需要用 elastic API 而非传统的 init_process_group:
python
import torch.distributed as dist
from torch.distributed.elastic.multiprocessing.errors import record
@record # 自动捕获异常并上报
def main():
# elastic 模式下,这些环境变量由 ElasticAgent 注入
rank = int(os.environ["RANK"])
world_size = int(os.environ["WORLD_SIZE"])
local_rank = int(os.environ["LOCAL_RANK"])
dist.init_process_group(backend="nccl")
model = MyModel().cuda(local_rank)
model = DDP(model, device_ids=[local_rank])
# 关键:加载 Checkpoint
load_checkpoint(model, optimizer, step)
# 关键:从正确的 step 继续
for step, batch in enumerate(dataloader, start=start_step):
loss = train_step(model, batch)
if step % save_every == 0:
save_checkpoint(model, optimizer, step)
3.2 Scale Down:Worker 退出
Worker 退出分两种情况:
优雅退出(主动 scale down):
python
# ElasticAgent 收到 SIGTERM
# → 通知 Rendezvous 本节点退出
# → 触发 re-rendezvous
# → 剩余 Worker 重组
非优雅退出(崩溃 / 抢占):
python
# ElasticAgent 心跳超时
# → Rendezvous 标记为 dead
# → 存活 Worker 的 NCCL 操作超时
# → 触发 re-rendezvous
两者最终都走 re-rendezvous 路径,区别在于优雅退出能保证最近 Checkpoint 已保存,非优雅退出可能丢失 save_interval 内的进度。
3.3 处理梯度缺失:Join 机制
World Size 变化时,剩余 Worker 的梯度数量与之前不同。假设原来 8 个 Worker 做 AllReduce,现在剩 6 个,那 step 的有效 batch size 从 8×B 变成 6×B。如果不处理,学习率、BatchNorm 统计都会失真。
PyTorch 提供 Join 上下文管理器处理「部分 Worker 缺席」的情况:
python
from torch.distributed.optim import ZeroRedundancyOptimizer
import contextlib
for step, batch in enumerate(dataloader):
# join_context 处理某些 rank 提前退出的情况
with contextlib.nullcontext(): # DDP 已内置 join
loss = model(batch)
loss.backward()
optimizer.step()
DDP 从 1.10 起内置 join,当某个 rank 的数据先用完(或退出),它会发送「假梯度」让 AllReduce 正常完成,剩余 rank 继续训练。这对弹性场景至关重要。
但对于 BatchNorm 等依赖 world_size 的层,仍需业务代码感知 world_size 变化并调整:
python
# 感知 world_size 变化,调整有效 batch size
effective_batch = world_size * per_gpu_batch
# 相应调整学习率(Linear Scaling Rule)
lr = base_lr * (effective_batch / base_batch)
for g in optimizer.param_groups:
g["lr"] = lr
3.4 与 Training Operator 集成
第 30 篇的 PyTorchJob 默认用的就是 elastic 启动方式:
yaml
apiVersion: kubeflow.org/v1
kind: PyTorchJob
metadata:
name: elastic-finetune
spec:
pytorchReplicaSpecs:
Worker:
replicas: 4
restartPolicy: OnFailure # 关键:失败重启而非整个 Job 失败
template:
spec:
containers:
- name: pytorch
command:
- torchrun
- --nnodes=2:8 # 弹性区间
- --nproc_per_node=8
- --rdzv_backend=c10d
- --rdzv_endpoint=elastic-finetune-worker-0:29500
- --rdzv_id=$(JOB_NAME)
- --max_restarts=10
- train.py
注意 restartPolicy: OnFailure + --max_restarts=10 的组合:单个 Worker 失败时重启,最多重试 10 次,而不是整个 Job 失败。这是弹性训练在 K8s 上的落地姿势。
四、Checkpoint 持久化与恢复
4.1 Checkpoint 的正确姿势
Checkpoint 是弹性训练的生命线。一个生产可用的 Checkpoint 保存策略必须满足:
- 原子性:写一半被中断的 Checkpoint 是「脏的」,恢复时会崩溃。
- 一致性:所有 rank 的 Checkpoint 必须对应同一个 step。
- 频率合适:太频繁拖慢训练,太稀疏丢进度。
- 存储可靠:本地磁盘不可靠,必须持久化到分布式存储。
python
import torch
import torch.distributed as dist
import os
import shutil
from pathlib import Path
CHECKPOINT_DIR = "/checkpoints"
def save_checkpoint(model, optimizer, scheduler, step, epoch, config):
"""原子性保存 Checkpoint"""
if dist.get_rank() != 0:
# 只有 rank 0 保存完整模型
dist.barrier()
return
ckpt = {
"step": step,
"epoch": epoch,
"model_state": model.module.state_dict(), # DDP 要 .module
"optimizer_state": optimizer.state_dict(),
"scheduler_state": scheduler.state_dict(),
"rng_state": torch.get_rng_state(),
"cuda_rng_state": torch.cuda.get_rng_state_all(),
"config": config,
"world_size": dist.get_world_size(), # 记录当时的 world_size
}
# 原子写:先写临时文件,再 rename
tmp_path = Path(CHECKPOINT_DIR) / f"ckpt-{step}.pt.tmp"
final_path = Path(CHECKPOINT_DIR) / f"ckpt-{step}.pt"
torch.save(ckpt, tmp_path)
tmp_path.rename(final_path) # rename 是原子的
# 清理旧 Checkpoint,只保留最近 N 个
clean_old_checkpoints(keep=3)
dist.barrier() # 所有 rank 同步,确保 rank 0 写完
def clean_old_checkpoints(keep=3):
ckpt_dir = Path(CHECKPOINT_DIR)
ckpts = sorted(ckpt_dir.glob("ckpt-*.pt"),
key=lambda p: int(p.stem.split("-")[1]))
for old in ckpts[:-keep]:
old.unlink()
原子写的核心 :rename 系统调用在同一个文件系统内是原子的。先写到 .tmp,写完再 rename 到正式文件名,这样即使写到一半进程崩溃,正式文件名指向的始终是完整的上一个 Checkpoint。
4.2 存储后端选择
| 存储 | 优点 | 缺点 | 适用场景 |
|---|---|---|---|
| 本地 NVMe | 最快 | 节点丢失即丢失 | 临时快照 |
| NFS / PVC (RWO) | 简单 | 单点,吞吐有限 | 小模型 |
| PVC (RWX, 共享) | 多 Pod 读 | 写并发需协调 | 中等规模 |
| S3 / OSS / MinIO | 可靠、廉价 | 延迟高 | 长期归档 |
| 并行文件系统 (Lustre/GPFS) | 高吞吐 | 运维复杂 | 大规模训练 |
生产推荐:训练时写本地 NVMe → 异步上传 S3,恢复时从 S3 下载。这兼顾速度和可靠性:
python
import boto3
import threading
def async_upload_to_s3(local_path, s3_key):
def _upload():
s3 = boto3.client("s3")
s3.upload_file(local_path, "my-bucket", s3_key)
threading.Thread(target=_upload, daemon=True).start()
# 保存后异步上传
save_checkpoint(model, optimizer, step)
async_upload_to_s3(f"{CHECKPOINT_DIR}/ckpt-{step}.pt", f"jobs/{job_id}/ckpt-{step}.pt")
4.3 恢复流程
恢复时,每个 Worker 从同一个 Checkpoint 加载,但只恢复属于自己的分片(对 ZeRO / TP 分片的情况):
python
def load_checkpoint(model, optimizer, scheduler):
"""加载最新的 Checkpoint"""
ckpt_path = find_latest_checkpoint(CHECKPOINT_DIR)
if ckpt_path is None:
print("No checkpoint found, starting from scratch")
return 0, 0 # step, epoch
ckpt = torch.load(ckpt_path, map_location="cpu")
# 关键:加载 RNG 状态,保证可复现
torch.set_rng_state(ckpt["rng_state"])
if torch.cuda.is_available():
torch.cuda.set_rng_state_all(ckpt["cuda_rng_state"])
# 模型加载(DDP 要先加载再 wrap)
model.load_state_dict(ckpt["model_state"])
optimizer.load_state_dict(ckpt["optimizer_state"])
scheduler.load_state_dict(ckpt["scheduler_state"])
return ckpt["step"], ckpt["epoch"]
def find_latest_checkpoint(ckpt_dir):
ckpts = sorted(Path(ckpt_dir).glob("ckpt-*.pt"),
key=lambda p: int(p.stem.split("-")[1]),
reverse=True)
return ckpts[0] if ckpts else None
注意加载 RNG 状态:弹性重组后 world_size 可能变化,DataLoader 的采样顺序会不同。恢复 RNG 保证(在 world_size 不变时)训练可复现;world_size 变化时至少保证单 rank 内的随机性一致。
4.4 Checkpoint 频率与 Save Interval
频率的权衡:
每次保存耗时 ≈ 模型大小 / 磁盘吞吐
例如 70B 模型 (140GB) 写 NVMe (3GB/s) ≈ 47 秒
如果每 100 step 保存一次,每 step 2 秒:
训练占比 = 100×2 / (100×2 + 47) = 80.9%
即 19% 时间花在保存 Checkpoint 上
实践建议:
- 频率:每 100-500 step 保存一次(根据 step 耗时调整,保证 Checkpoint 开销 < 5%)。
- 保留:最近 3-5 个 + 每隔 N 千 step 一个长期 Checkpoint。
- 异步 :用
save_async或后台线程,不阻塞训练循环。 - 校验 :保存后立即做一次
torch.load校验完整性,避免恢复时才发现损坏。
五、NCCL 超时与僵尸进程治理
5.1 NCCL 超时的本质
NCCL(NVIDIA Collective Communications Library)是 GPU 间通信的底层库,AllReduce / AllGather 等集合通信都走它。NCCL 的一个致命特性:它是阻塞式的,且对网络故障极其敏感。
Worker A 发起 AllReduce, 等待 Worker B 的梯度
Worker B 因 OOM / 抢占 / IB 网卡故障, 无响应
Worker A 阻塞在 NCCL 内核调用, 无法被 Python 层中断
超时 (默认 30 分钟) 才报错
这个 30 分钟的默认超时对生产环境太长了。一个训练任务在 NCCL 超时期间完全空转,GPU 100% 占用但 0% 有效计算。
5.2 调整超时参数
python
import datetime
dist.init_process_group(
backend="nccl",
timeout=datetime.timedelta(minutes=5), # 从 30min 降到 5min
)
或者通过环境变量:
bash
export NCCL_TIMEOUT=300 # 秒
export NCCL_BLOCKING_WAIT=1 # 阻塞等待(而非忙等,省 CPU)
5 分钟是一个折中:太短会因正常的网络抖动误判,太长则浪费 GPU。对于跨节点训练,建议从 5 分钟起步,观察误报率后微调。
5.3 NCCL 调试与诊断
NCCL 故障的诊断是分布式训练最难的部分。关键工具:
bash
# 开启 NCCL 调试日志(仅排查时开启,正常训练关闭)
export NCCL_DEBUG=INFO # 或 SUBSYS / TRACE
export NCCL_DEBUG_SUBSYS=ALL
# 常见问题日志关键词
# "transport/p2p" → NVLink / PCIe 拓扑问题
# "net" → InfiniBand / RoCE 网络问题
# "watchdog" → 看门狗超时,通常是 GPU 挂起
诊断决策树:
| 症状 | 可能原因 | 排查命令 |
|---|---|---|
| AllReduce 超时 | 节点间网络不通 | nccl-test all_reduce_perf |
| 只在某些 rank 对超时 | NVLink 链路故障 | nvidia-smi nvlink -s |
| 训练卡住但无报错 | NCCL 死锁 | py-spy dump --pid <worker_pid> |
| 重启后恢复 | 瞬时网络抖动 | 检查 IB 交换机日志 |
| 固定节点总是超时 | GPU 硬件故障 | nvidia-smi -q, dcgmi diag |
5.4 僵尸进程检测
Worker 进程崩溃但 ElasticAgent 没正确清理时,会产生僵尸进程,占用 GPU 但不工作。检测方法:
bash
# 查看哪些 GPU 被进程占用
nvidia-smi --query-compute-apps=pid,used_memory --format=csv
# 对比训练任务声明的 PID
# 如果有 PID 不属于任何 Worker,就是僵尸
ps aux | grep python | grep train.py
用 Go 写一个定期巡检的 sidecar:
go
package main
import (
"context"
"fmt"
"os/exec"
"strings"
"time"
)
type GPUPid struct {
PID string
MemUsed string
}
func getGPUComputeApps() ([]GPUPid, error) {
out, err := exec.Command("nvidia-smi",
"--query-compute-apps=pid,used_memory",
"--format=csv,noheader,nounits").Output()
if err != nil {
return nil, err
}
var pids []GPUPid
for _, line := range strings.Split(strings.TrimSpace(string(out)), "\n") {
parts := strings.Split(line, ", ")
if len(parts) == 2 {
pids = append(pids, GPUPid{PID: strings.TrimSpace(parts[0]), MemUsed: parts[1]})
}
}
return pids, nil
}
func getTrainingPIDs() (map[string]bool, error) {
out, err := exec.Command("pgrep", "-f", "train.py").Output()
if err != nil {
return nil, err
}
pids := make(map[string]bool)
for _, p := range strings.Fields(string(out)) {
pids[p] = true
}
return pids, nil
}
func main() {
ticker := time.NewTicker(30 * time.Second)
defer ticker.Stop()
for range ticker.C {
gpuPids, _ := getGPUComputeApps()
trainPIDs, _ := getTrainingPIDs()
for _, g := range gpuPids {
if !trainPIDs[g.PID] {
fmt.Printf("[ZOMBIE] GPU PID %s using %s MiB is not a training process\n",
g.PID, g.MemUsed)
// 可选:自动 kill
// exec.Command("kill", "-9", g.PID).Run()
}
}
}
_ = context.Background()
}
这个 sidecar 以 DaemonSet 形式部署到每个 GPU 节点,每 30 秒扫描一次,发现僵尸进程就告警(或自动 kill)。
5.5 NCCL 死锁的常见模式
模式一:AllReduce 与 AllGather 交错
某些 ZeRO 实现里,不同 rank 的 AllReduce 和 AllGather 顺序不一致,导致环形等待。解决:确保所有 rank 以相同顺序发起集合通信。
模式二:部分 rank 先退出
某个 rank 的数据先用完提前退出,剩余 rank 在 AllReduce 等它。解决:用 Join 机制,或确保所有 rank 的数据量一致。
模式三:网络分区
部分节点间网络中断,形成分区。NCCL 在分区内能通信,跨分区超时。解决:监控网络拓扑,及时驱逐故障节点。
六、Spot 实例训练实践
6.1 Spot 实例的两分钟通知
AWS Spot / GCP Preemptible / Azure Spot VM 都会在回收前给一个短通知:
| 云厂商 | 通知方式 | 通知时长 |
|---|---|---|
| AWS EC2 Spot | Instance Metadata termination-time |
2 分钟 |
| GCP Preemptible | Metadata preempted + ACPI 事件 |
30 秒 |
| Azure Spot VM | Scheduled Events API | 30 秒 |
训练任务必须在这 2 分钟内完成「保存 Checkpoint + 优雅退出」:
#mermaid-svg-A24lRp2z2srIVLS2{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-A24lRp2z2srIVLS2 .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-A24lRp2z2srIVLS2 .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-A24lRp2z2srIVLS2 .error-icon{fill:#552222;}#mermaid-svg-A24lRp2z2srIVLS2 .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-A24lRp2z2srIVLS2 .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-A24lRp2z2srIVLS2 .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-A24lRp2z2srIVLS2 .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-A24lRp2z2srIVLS2 .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-A24lRp2z2srIVLS2 .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-A24lRp2z2srIVLS2 .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-A24lRp2z2srIVLS2 .marker{fill:#333333;stroke:#333333;}#mermaid-svg-A24lRp2z2srIVLS2 .marker.cross{stroke:#333333;}#mermaid-svg-A24lRp2z2srIVLS2 svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-A24lRp2z2srIVLS2 p{margin:0;}#mermaid-svg-A24lRp2z2srIVLS2 .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-A24lRp2z2srIVLS2 .cluster-label text{fill:#333;}#mermaid-svg-A24lRp2z2srIVLS2 .cluster-label span{color:#333;}#mermaid-svg-A24lRp2z2srIVLS2 .cluster-label span p{background-color:transparent;}#mermaid-svg-A24lRp2z2srIVLS2 .label text,#mermaid-svg-A24lRp2z2srIVLS2 span{fill:#333;color:#333;}#mermaid-svg-A24lRp2z2srIVLS2 .node rect,#mermaid-svg-A24lRp2z2srIVLS2 .node circle,#mermaid-svg-A24lRp2z2srIVLS2 .node ellipse,#mermaid-svg-A24lRp2z2srIVLS2 .node polygon,#mermaid-svg-A24lRp2z2srIVLS2 .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-A24lRp2z2srIVLS2 .rough-node .label text,#mermaid-svg-A24lRp2z2srIVLS2 .node .label text,#mermaid-svg-A24lRp2z2srIVLS2 .image-shape .label,#mermaid-svg-A24lRp2z2srIVLS2 .icon-shape .label{text-anchor:middle;}#mermaid-svg-A24lRp2z2srIVLS2 .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-A24lRp2z2srIVLS2 .rough-node .label,#mermaid-svg-A24lRp2z2srIVLS2 .node .label,#mermaid-svg-A24lRp2z2srIVLS2 .image-shape .label,#mermaid-svg-A24lRp2z2srIVLS2 .icon-shape .label{text-align:center;}#mermaid-svg-A24lRp2z2srIVLS2 .node.clickable{cursor:pointer;}#mermaid-svg-A24lRp2z2srIVLS2 .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-A24lRp2z2srIVLS2 .arrowheadPath{fill:#333333;}#mermaid-svg-A24lRp2z2srIVLS2 .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-A24lRp2z2srIVLS2 .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-A24lRp2z2srIVLS2 .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-A24lRp2z2srIVLS2 .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-A24lRp2z2srIVLS2 .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-A24lRp2z2srIVLS2 .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-A24lRp2z2srIVLS2 .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-A24lRp2z2srIVLS2 .cluster text{fill:#333;}#mermaid-svg-A24lRp2z2srIVLS2 .cluster span{color:#333;}#mermaid-svg-A24lRp2z2srIVLS2 div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-A24lRp2z2srIVLS2 .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-A24lRp2z2srIVLS2 rect.text{fill:none;stroke-width:0;}#mermaid-svg-A24lRp2z2srIVLS2 .icon-shape,#mermaid-svg-A24lRp2z2srIVLS2 .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-A24lRp2z2srIVLS2 .icon-shape p,#mermaid-svg-A24lRp2z2srIVLS2 .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-A24lRp2z2srIVLS2 .icon-shape .label rect,#mermaid-svg-A24lRp2z2srIVLS2 .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-A24lRp2z2srIVLS2 .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-A24lRp2z2srIVLS2 .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-A24lRp2z2srIVLS2 :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} 是
否, 被强杀
Spot 回收通知
训练任务监听到?
立即保存 Checkpoint
通知 Rendezvous 退出
ElasticAgent 触发 re-rendezvous
剩余 Worker 重组
从 Checkpoint 恢复训练
NCCL 超时检测
存活 Worker 触发重组
从上一个 Checkpoint 恢复
(丢失 save_interval 内进度)
训练继续
6.2 监听中断信号
在 Pod 里监听 Spot 通知:
python
import threading
import time
import requests
def spot_watcher(model, optimizer, step):
"""监听 Spot 中断,触发优雅保存"""
while True:
try:
# AWS: 检查 Instance Metadata
resp = requests.get(
"http://169.254.169.254/latest/meta-data/spot/instance-action",
timeout=2,
)
if resp.status_code == 200:
action = resp.json()
if action["action"] == "terminate":
print(f"[SPOT] 2分钟内将被回收, 立即保存 Checkpoint")
save_checkpoint(model, optimizer, step)
# 通知 ElasticAgent 退出
os.kill(os.getpid(), signal.SIGTERM)
return
except Exception:
pass
time.sleep(5)
# 训练启动时起一个监听线程
watcher_thread = threading.Thread(
target=spot_watcher, args=(model, optimizer, step), daemon=True
)
watcher_thread.start()
在 K8s 里,更优雅的方式是利用 Pod Disruption Budget + 优雅终止:
yaml
spec:
template:
spec:
terminationGracePeriodSeconds: 120 # 给 2 分钟优雅退出
containers:
- name: pytorch
lifecycle:
preStop:
exec:
command:
- /bin/sh
- -c
- |
# 发送 SIGTERM, 训练脚本捕获后保存 Checkpoint
kill -TERM 1
# 等待保存完成
sleep 100
terminationGracePeriodSeconds: 120 与 Spot 的 2 分钟通知对齐,preStop 钩子确保训练脚本有机会保存。
6.3 Checkpoint 频率与 Spot 的权衡
Spot 实例的 MTBF(平均故障间隔)决定了 Checkpoint 频率:
假设: 64 卡 Spot 集群, 单节点每小时被抢占概率 5%
则整个集群每小时至少一个节点被抢占的概率:
P = 1 - (1-0.05)^64 ≈ 96%
即几乎每小时都有节点被抢占。
如果 save_interval = 1000 step (约 33 分钟),
每次抢占平均丢失 16.5 分钟训练进度。
因此 Spot 训练的 Checkpoint 频率要比按需实例高得多。推荐:
save_interval: 50-200 step(保证每次抢占丢失 < 5 分钟)。- 异步保存,不阻塞训练。
- 保存到 S3(Spot 节点随时消失,本地磁盘不可靠)。
6.4 Spot 集群拓扑管理
Spot 实例的可用性是波动的,某个 AZ 可能突然没有 Spot 容量。实践建议:
- 多 AZ 部署:Spread 跨 3 个 AZ,降低同时被抢占概率。
- 多实例类型:同时请求 p4d 和 p4de,容量池更大。
- 混合策略:Head 节点用按需(保证 Rendezvous 稳定),Worker 用 Spot。
- 容量预留:关键训练任务配合 Capacity Reservation / Fleet。
K8s 里的落地:
yaml
apiVersion: kubeflow.org/v1
kind: PyTorchJob
spec:
pytorchReplicaSpecs:
Worker:
replicas: 8
template:
spec:
priorityClassName: spot-preemptible # 低优先级, 易被抢占
nodeSelector:
karpenter.sh/capacity-type: spot # 用 Karpenter 调度 Spot
topologySpreadConstraints:
- maxSkew: 1
topologyKey: topology.kubernetes.io/zone
whenUnsatisfiable: ScheduleAnyway
containers:
- name: pytorch
command: [torchrun, --nnodes=2:8, ...]
lifecycle:
preStop:
exec:
command: ["/scripts/save_and_exit.sh"]
6.5 成本模型
一个实际的成本对比(训练 70B 模型 7 天):
| 方案 | GPU | 单价 | 实际利用率 | 总成本 |
|---|---|---|---|---|
| 纯按需 64×A100 | p4d.24xlarge × 8 | $32.77/h | 95% | $55,200 |
| 纯 Spot 64×A100 | p4d Spot × 8 | $9.83/h | 70%(含抢占恢复) | $12,300 |
| 混合 (Head 按需 + Worker Spot) | 1 按需 + 7 Spot | 混合 | 85% | $15,800 |
混合方案是性价比最优的:Head 节点(跑 Rendezvous)用按需保证稳定性,Worker 全部 Spot 降成本。即使 Worker 被抢占,Rendezvous 还在,恢复很快。
七、故障恢复完整流程
7.1 端到端恢复时序
S3 存储 Rendezvous ElasticAgent Checkpoint Manager 训练循环 S3 存储 Rendezvous ElasticAgent Checkpoint Manager 训练循环 #mermaid-svg-0VpTTKJcjyRO68m5{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-0VpTTKJcjyRO68m5 .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-0VpTTKJcjyRO68m5 .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-0VpTTKJcjyRO68m5 .error-icon{fill:#552222;}#mermaid-svg-0VpTTKJcjyRO68m5 .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-0VpTTKJcjyRO68m5 .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-0VpTTKJcjyRO68m5 .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-0VpTTKJcjyRO68m5 .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-0VpTTKJcjyRO68m5 .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-0VpTTKJcjyRO68m5 .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-0VpTTKJcjyRO68m5 .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-0VpTTKJcjyRO68m5 .marker{fill:#333333;stroke:#333333;}#mermaid-svg-0VpTTKJcjyRO68m5 .marker.cross{stroke:#333333;}#mermaid-svg-0VpTTKJcjyRO68m5 svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-0VpTTKJcjyRO68m5 p{margin:0;}#mermaid-svg-0VpTTKJcjyRO68m5 .actor{stroke:hsl(259.6261682243, 59.7765363128%, 87.9019607843%);fill:#ECECFF;}#mermaid-svg-0VpTTKJcjyRO68m5 text.actor>tspan{fill:black;stroke:none;}#mermaid-svg-0VpTTKJcjyRO68m5 .actor-line{stroke:hsl(259.6261682243, 59.7765363128%, 87.9019607843%);}#mermaid-svg-0VpTTKJcjyRO68m5 .innerArc{stroke-width:1.5;stroke-dasharray:none;}#mermaid-svg-0VpTTKJcjyRO68m5 .messageLine0{stroke-width:1.5;stroke-dasharray:none;stroke:#333;}#mermaid-svg-0VpTTKJcjyRO68m5 .messageLine1{stroke-width:1.5;stroke-dasharray:2,2;stroke:#333;}#mermaid-svg-0VpTTKJcjyRO68m5 #arrowhead path{fill:#333;stroke:#333;}#mermaid-svg-0VpTTKJcjyRO68m5 .sequenceNumber{fill:white;}#mermaid-svg-0VpTTKJcjyRO68m5 #sequencenumber{fill:#333;}#mermaid-svg-0VpTTKJcjyRO68m5 #crosshead path{fill:#333;stroke:#333;}#mermaid-svg-0VpTTKJcjyRO68m5 .messageText{fill:#333;stroke:none;}#mermaid-svg-0VpTTKJcjyRO68m5 .labelBox{stroke:hsl(259.6261682243, 59.7765363128%, 87.9019607843%);fill:#ECECFF;}#mermaid-svg-0VpTTKJcjyRO68m5 .labelText,#mermaid-svg-0VpTTKJcjyRO68m5 .labelText>tspan{fill:black;stroke:none;}#mermaid-svg-0VpTTKJcjyRO68m5 .loopText,#mermaid-svg-0VpTTKJcjyRO68m5 .loopText>tspan{fill:black;stroke:none;}#mermaid-svg-0VpTTKJcjyRO68m5 .loopLine{stroke-width:2px;stroke-dasharray:2,2;stroke:hsl(259.6261682243, 59.7765363128%, 87.9019607843%);fill:hsl(259.6261682243, 59.7765363128%, 87.9019607843%);}#mermaid-svg-0VpTTKJcjyRO68m5 .note{stroke:#aaaa33;fill:#fff5ad;}#mermaid-svg-0VpTTKJcjyRO68m5 .noteText,#mermaid-svg-0VpTTKJcjyRO68m5 .noteText>tspan{fill:black;stroke:none;}#mermaid-svg-0VpTTKJcjyRO68m5 .activation0{fill:#f4f4f4;stroke:#666;}#mermaid-svg-0VpTTKJcjyRO68m5 .activation1{fill:#f4f4f4;stroke:#666;}#mermaid-svg-0VpTTKJcjyRO68m5 .activation2{fill:#f4f4f4;stroke:#666;}#mermaid-svg-0VpTTKJcjyRO68m5 .actorPopupMenu{position:absolute;}#mermaid-svg-0VpTTKJcjyRO68m5 .actorPopupMenuPanel{position:absolute;fill:#ECECFF;box-shadow:0px 8px 16px 0px rgba(0,0,0,0.2);filter:drop-shadow(3px 5px 2px rgb(0 0 0 / 0.4));}#mermaid-svg-0VpTTKJcjyRO68m5 .actor-man line{stroke:hsl(259.6261682243, 59.7765363128%, 87.9019607843%);fill:#ECECFF;}#mermaid-svg-0VpTTKJcjyRO68m5 .actor-man circle,#mermaid-svg-0VpTTKJcjyRO68m5 line{stroke:hsl(259.6261682243, 59.7765363128%, 87.9019607843%);fill:#ECECFF;stroke-width:2px;}#mermaid-svg-0VpTTKJcjyRO68m5 :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} 正常训练中, step=1500 step=1530 时, 节点3 被抢占 所有 Agent 重新加载 Checkpoint 丢失进度: 1530-1500=30 step (约 1 分钟) save_checkpoint(step=1500) 异步上传 ckpt-1500.pt 上传成功 NCCL AllReduce 超时 (5min) 报告通信失败 请求 re-rendezvous (节点3 dead) 重组 world: 8→7 节点 新 rank 分配 下载最新 ckpt-1500.pt 下载完成 注入 step=1500, world_size=56 调整学习率 (Linear Scaling) 从 step 1501 继续训练
7.2 恢复正确性验证
恢复后必须验证状态一致性,避免「静默错误」:
python
def verify_recovery(model, expected_step):
"""恢复后校验模型状态"""
# 1. 检查 step 是否对齐
assert current_step == expected_step, f"Step mismatch: {current_step} vs {expected_step}"
# 2. 所有 rank 的模型参数一致(DDP 保证)
for name, param in model.named_parameters():
param_clone = param.clone()
dist.all_reduce(param_clone, op=dist.ReduceOp.AVG)
assert torch.allclose(param, param_clone, atol=1e-6), \
f"Param {name} inconsistent across ranks after recovery"
# 3. DataLoader 恢复到正确位置
assert dataloader.current_step == expected_step
print(f"[RECOVERY] Verified at step {expected_step}, world_size={dist.get_world_size()}")
这个校验在每次恢复后执行,确保不会因为 Checkpoint 损坏或 world_size 变化导致训练发散。
7.3 常见故障与应对
| 故障 | 现象 | 应对 |
|---|---|---|
| 节点 OOM | Worker 进程被 kill | 调 batch size,启用 ZeRO Offload |
| NCCL 超时 | 训练卡住无报错 | 缩短 timeout,排查网络 |
| Checkpoint 损坏 | torch.load 报错 |
回退到上一个 Checkpoint |
| Rendezvous 后端挂 | 无法重组 | 用 etcd 集群保证 HA |
| Pod 频繁驱逐 | 训练无法稳定 | 检查节点资源水位,调 priority |
| GPU 驱动崩溃 | nvidia-smi 失败 |
cordon 节点,迁移任务 |
| 网络分区 | 部分 rank 互不可达 | 驱逐分区外节点 |
| Checkpoint 写满磁盘 | PVC 满 | 异步上传 S3,本地只留最近 1 个 |
八、生产落地清单
8.1 训练脚本 Checklist
□ 使用 torchrun --nnodes=min:max 弹性启动
□ dist.init_process_group 设置 timeout=5min
□ Checkpoint 原子写 (tmp + rename)
□ Checkpoint 异步上传 S3
□ save_interval 合理 (Spot: 50-200 step)
□ 加载时恢复 RNG 状态
□ 恢复后校验参数一致性
□ 监听 SIGTERM 优雅保存退出
□ 记录 world_size 变化到日志
□ join 机制处理 rank 提前退出
8.2 平台层 Checklist
□ PyTorchJob restartPolicy: OnFailure + max_restarts
□ Pod terminationGracePeriodSeconds ≥ 120
□ Spot 实例配 preStop 钩子
□ Checkpoint PVC 用 RWX (ReadWriteMany)
□ Rendezvous 后端 (etcd/c10d) 高可用
□ GPU 僵尸进程巡检 DaemonSet
□ NCCL 超时告警 (Prometheus 监控训练 step 停滞)
□ Checkpoint 存储成本监控 (S3 生命周期策略)
□ 多 AZ / 多实例类型 Spot 策略
□ 训练任务可观测 (TensorBoard / W&B 集成)
8.3 监控指标
弹性训练的关键监控指标:
| 指标 | 含义 | 告警阈值 |
|---|---|---|
training_step_rate |
每 step 耗时 | 突增 > 2x |
rendezvous_rejoin_count |
重组次数 | > 3/小时 |
checkpoint_save_duration |
保存耗时 | > 60s |
checkpoint_recovery_duration |
恢复耗时 | > 120s |
nccl_timeout_count |
NCCL 超时次数 | > 0 |
world_size_changes |
world size 变化 | 记录但不告警 |
gpu_idle_time_pct |
GPU 空闲率 | > 20% |
spot_interruption_count |
Spot 抢占次数 | 记录 |
本篇小结
| 主题 | 核心结论 |
|---|---|
| 固定 World Size 缺陷 | 单点失败全盘崩溃、无法伸缩、启动整批等待 |
| 弹性训练承诺 | 动态成员、状态恢复、优雅降级、最终一致 |
| Rendezvous | 动态会合机制,Worker 可随时加入/退出 |
| ElasticAgent | 节点级监督进程,管理本地 Worker 生命周期 |
--nnodes=min:max |
弹性区间的声明,最少 N 节点启动,最多 M 节点 |
| 成员变更 | NCCL 超时触发 re-rendezvous,Checkpoint 恢复进度 |
| Join 机制 | 处理 rank 提前退出,保证 AllReduce 正确 |
| Checkpoint 原子性 | tmp + rename,保证一致性 |
| Checkpoint 存储 | 本地 NVMe + 异步 S3,兼顾速度与可靠 |
| NCCL 超时 | 默认 30min 太长,调到 5min |
| 僵尸进程 | DaemonSet 定期扫描 GPU PID |
| Spot 实例 | 2 分钟通知,preStop 钩子保存退出 |
| Spot 成本 | 纯 Spot 省 70%,混合方案性价比最优 |
| 恢复校验 | step 对齐 + 参数一致 + DataLoader 位置 |
核心心智模型:弹性训练 = Rendezvous(动态会合)+ Checkpoint(状态恢复)+ NCCL 监控(故障检测)。三者缺一不可,组合起来才能让训练在 Spot 实例上稳定跑满 7 天。
下篇预告
第 33 篇:KServe推理平台-标准化的K8s推理服务管理
训练搞定了,接下来就是把模型变成线上服务。K8s 上管理推理服务,如果用裸 Deployment + Service,每个模型都得手写扩缩容、灰度、健康检查,运维成本极高。有没有一个「推理服务专用」的 K8s Operator?下一篇我们讲 KServe 推理平台------它的 InferenceService CRD 如何把推理服务的生命周期、自动扩缩容、Canary/蓝绿/影子部署、多模型编排全部标准化,让推理服务的运维像 Deployment 一样简单。
如果本篇内容对你有帮助,欢迎点赞收藏!有任何疑问,欢迎在评论区交流。