第32篇-弹性训练与容错-让训练任务永不中断

【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 动态增减的前提下,保证训练正确性,并尽量减少进度损失。它的核心承诺:

  1. 动态成员:Worker 可以中途加入或退出,不重启整个任务。
  2. 状态恢复:新加入或恢复的 Worker 从最近的 Checkpoint 继续,不从头训练。
  3. 优雅降级:Worker 减少时,剩余 Worker 继续训练,world_size 动态调整。
  4. 最终一致 :无论经历多少次成员变更,训练收敛性与固定 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,它负责:

  1. 向 Rendezvous 后端注册本节点的 Worker 数量。
  2. 等待凑够 min_nodes 后,分配 rank,初始化进程组。
  3. 监督本地 Worker 进程,崩溃则重启,节点退出则通知 Rendezvous。
  4. 加载 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 就绪 开始新一轮训练 开始新一轮训练

注意几个细节:

  1. NCCL 超时是触发器:存活的 Worker 通过 NCCL 操作(AllReduce 等)感知到对端无响应,超时后主动请求重组。
  2. Checkpoint 是恢复依据:重组后所有 Worker 重新加载最近 Checkpoint,避免梯度不一致。
  3. Rank 重新分配:重组后 rank 0-15 重新映射到存活的两个节点,旧的 rank 16-23(崩溃节点的)被丢弃。

三、动态 Worker 增减

3.1 Scale Up:新 Worker 加入

新节点加入训练任务的流程:

  1. 新节点的 ElasticAgent 启动,向 Rendezvous 注册。
  2. Rendezvous 检查当前 world 是否已满(max_nodes)。
    • 未满:等待当前训练 step 结束,触发重组。
    • 已满:新节点排队,等下一轮有空位。
  3. 重组:新节点加入 world,重新分配 rank。
  4. 所有 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 保存策略必须满足:

  1. 原子性:写一半被中断的 Checkpoint 是「脏的」,恢复时会崩溃。
  2. 一致性:所有 rank 的 Checkpoint 必须对应同一个 step。
  3. 频率合适:太频繁拖慢训练,太稀疏丢进度。
  4. 存储可靠:本地磁盘不可靠,必须持久化到分布式存储。
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 容量。实践建议:

  1. 多 AZ 部署:Spread 跨 3 个 AZ,降低同时被抢占概率。
  2. 多实例类型:同时请求 p4d 和 p4de,容量池更大。
  3. 混合策略:Head 节点用按需(保证 Rendezvous 稳定),Worker 用 Spot。
  4. 容量预留:关键训练任务配合 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 一样简单。


如果本篇内容对你有帮助,欢迎点赞收藏!有任何疑问,欢迎在评论区交流。

相关推荐
移动云开发者联盟1 小时前
密态计算落地!MobileClaw解锁安全AI新范式
大数据·人工智能·安全
一次旅行1 小时前
2026‑08‑19 AI产业深度解读|Mojo正式开源、Agent记忆剂量研究、OpenAI收紧模型安全管控
人工智能·开源·mojo
HackTwoHub1 小时前
开箱即用 AI 渗透测试系统,整合多类AI工具、Web安全、内网攻防、二进制逆向、API测试、流量分析等多个安全领域
人工智能·安全·web安全·网络安全·系统安全·网络攻击模型·安全架构
TAN-90°-1 小时前
Deep Learning for Computer Vision——Training CNNs and CNN Architectures
人工智能·深度学习·神经网络·算法·机器学习·计算机视觉·cnn
0x3F(小茶)1 小时前
Tokenization(分词算法):一切大语言模型的地基
人工智能·算法·语言模型
MartinYeung51 小时前
[论文学习]ChainWatch:面向MCP-Based AI智能体系统中多步攻击的杀伤链对齐序贯检测框架
人工智能·学习
集萃智造机器人2 小时前
集萃智造采摘机器人:以具身智能破解果蔬采收难题,助力智慧农业升级
人工智能·机器人
fīɡЙtīиɡ ℡2 小时前
AI 应用评测体系
人工智能·学习