大模型容错训练系统设计:从故障检测、原地恢复到 TLA+ 形式化验证

本文基于一个真实落地的容错训练系统(下称 Aegis)的设计文档与源码整理,隐去了公司与内部组件信息,聚焦技术方案本身。

大模型训练的稳定性问题,是所有大规模训练团队的共同痛点。本文拆解 Aegis 的完整设计:为什么在调度平台层做容错、gRPC+NATS 双通道怎么配合、Master-Worker 状态机怎么流转、容错计数怎么算、故障节点怎么隔离,以及最有意思的一环——用 TLA+ 形式化验证状态机的安全性。

关联阅读:训练侧并行与通信见图解 Megatron 的 SP/CP/EP、分布式训练核心技术:从 FSDP 到 All-Reduce。容错是这些训练技术之下的"稳定性底座"。

一、为什么需要专门的容错系统:三大痛点

1.1 训练频繁中断是常态

大模型分布式训练几乎都是同步方式(DDP/MPI):所有进程共同构建一个通信拓扑,任何一个进程出错,其他进程都会相继退出。而训练任务动辄跨数百个 GPU 节点,数小时到数天内大概率出各种问题:

OPT 175B 训练的 LogBook:半个月内因硬件、基础设施等问题中断 40 多次;
BLOOM 176B:约 400 GPU 的集群上,平均每周遇到 1-2 个 GPU 故障;
多家大规模训练团队的统计都呈现类似特征:数天一次中断是常态,不是异常。

1.2 出错原因难排查

导致训练异常的原因五花八门:节点硬件故障、系统故障、网络问题、存储问题、训练代码问题。行业内的出错原因统计里,NET/IB、socket timeout、GPU ECC error 是常客——这些仅仅依赖重新调度是无法解决的,必须定位到具体哪个节点出错、隔离故障节点才能继续。

更麻烦的是同一个表象对应多种根因:通信超时可能是慢节点/坏节点,也可能是存储故障,甚至用户侧通信代码 bug;不同根因对应不同恢复策略。定位这些问题往往要花数小时甚至更长时间——这段时间几百张卡都在空烧。

1.3 Checkpoint 恢复的读写开销大

当前通用的恢复手段是 Checkpoint:把优化器状态、权重等以快照形式持久化。但 ckpt 大小与参数规模成正比——175B 模型带 fp32 优化器状态 + bf16+fp32 权重的 ckpt 高达 2.3TB;BLOOM 每 3 小时存一次,OPT 每 250 steps 存一次。模型和集群越大,ckpt 的读写效率和稳定性本身就成了训练效率的主要瓶颈之一。

三个痛点叠加的结论:"重调度 + ckpt 恢复"的粗粒度方案,在数百卡规模下既慢又贵。需要一套能在调度平台层做细粒度容错的系统——故障检测、原地恢复、故障节点隔离,把一次中断的代价从"整任务重启 + 2.3TB 读盘"压到"单 Pod 原地重跑"。

二、设计切入点:为什么选调度平台层

解决训练稳定性,理论上要从硬件、调度平台、训练框架甚至算法层 co-design。但很少有公司能 handle 整个链路从 0 到 1 的设计与实现。

Aegis 的选择是覆盖范围最广的一层——调度平台层:向上兼容基于 PyTorch 的各种分布式训练任务,向下兼容底层异构硬件架构(包括多种国产加速卡)。训练代码零侵入,用户只多配一个"最大容错次数"。

三、整体架构:三个组件 + 一个 CR

┌──────────────────────────────────────────────────────────────────────────┐
│ 训练平台层                                                               │
│  Job Controller                                                          │
│   ├─ 按 Job CR 创建 Worker Pod(注入 init container + tini 入口)        │
│   ├─ 创建/删除 Aegis CR(谁创建谁删除)                                  │
│   ├─ watch Pod 退出码 + CR 的 ShouldQuit/ErrorNodes                     │
│   │    ├─ Pod exit 1(节点故障)→ 调整 Affinity → 重调度到健康节点       │
│   │    └─ Retries > MaxRetry     → 关闭该任务全部 Pod(Failed)          │
└──────┬───────────────────────────────────────────┬───────────────────────┘
       │ 创建 CR                                   │ watch Pod
       ▼                                           │
┌──────────────────────────────────────────────────▼───────────────────────┐
│ Aegis Controller                                                         │
│  Master 模块(常驻)                                                      │
│   ├─ 状态机驱动:Pending→BootChecking→RunCommand⇄Troubleshooting        │
│   ├─ 容错记账:PodInfos[w].Retries,RecoverCount = max(Retries)          │
│   ├─ gRPC Server:首次通信,向 Worker 发放 NATS JWT(仅此一次)          │
│   ├─ NATS 发布 {Job}-master:bootcheck / run / troubleshoot 指令        │
│   ├─ NATS 订阅 {Job}-worker :ACK / Error / Finish / hang 上报          │
│   └─ AtlasLog 汇总 → 日志目录 → 采集管道 → 日志后端                     │
│  Prometheus:七项检测耗时 histogram + jobs counter(gpu_type×worker_num) │
└──────┬────────────────────────────────┬──────────────────────────────────┘
       │ gRPC(首次:SyncWorkerCreds    │ NATS(常驻双向,JWT 鉴权)
       │   发放 JWT + NKey)            │
       ▼                                ▼
┌───────────────────────────┐   ┌───────────────────────────────────────┐
│ Worker Pod × N            │   │ NATS 消息总线                        │
│ ┌───────────────────────┐ │   │  {Job}-master:Master → 全员指令      │
│ │ init container        │ │   │  {Job}-worker :Worker → Master 上报  │
│ │  tini-init 装五件套:  │ │   │  Worker 的 JWT 只允许订阅本 Job 主题 │
│ │  ① tini/atlctl 二进制 │ │   └───────────────────────────────────────┘
│ │  ② 检测脚本           │ │
│ │  ③ 检测初始化器        │ │
│ │  ④ ckpt.sh+保存引擎包 │ │   ┌───────────────────────────────────────┐
│ │    (pip 装,已装跳过) │ │   │ 可观测                               │
│ │  ⑤ atlctl 控制工具    │ │   │  Prometheus(检测耗时/任务计数)      │
│ └──────────┬────────────┘ │   │  AtlasLog(结构化容错日志)           │
│            ▼ 拷入共享卷    │   │   → 采集管道 → 日志后端 → 前端       │
│ ┌───────────────────────┐ │   └───────────────────────────────────────┘
│ │ 主容器                │ │
│ │  tini(PID 1,替 bash)│ │
│ │   ├─ runner:执行/停止│ │
│ │   │   用户命令(子进程)│ │
│ │   ├─ worker:通信模块  │ │
│ │   │   订阅+上报+心跳   │ │
│ │   ├─ subservice:     │ │
│ │   │   ckpt saver 守护 │ │
│ │   └─ atlctl CLI:     │ │
│ │       hang/burn/check│ │
│ │       /taint/stop... │ │
│ │  + 七项检测脚本       │ │
│ └───────────────────────┘ │
└───────────────────────────┘

四个核心组件:

组件	职责
Job Controller	部署和管理训练作业生命周期;按需创建 Aegis CR;Pod 退出后按 Affinity 调整 annotations/labels 避免重调度到坏节点
Aegis Controller	监视 Aegis CR、控制 Worker 执行任务;内置 Master 模块负责通信与状态机
Aegis CR	跟踪训练任务实际执行状态与历史出错记录(任务数、容错次数、Worker 状态、ErrorNodes)
tini(Worker 侧)	替换容器入口的进程管理器:runner 执行/停止任务,worker 模块管通信,内置检测脚本

权限设计遵循"谁创建谁删除":Aegis CR 由 Job Controller 创建与删除、Aegis Controller 只更新;Worker Pod 由 Job Controller 创建,Aegis Controller 只 watch。两个 controller 各管一摊,不会打架。

端到端生命周期:从提交到训练

 阶段 0:提交与准备        阶段 1:建联(一次性)     阶段 2:自检
─────────────────────   ──────────────────────   ─────────────────────
用户提交 Job CR
    │
    ▼
Job Controller
  创建 N 个 Worker Pod
  + Aegis CR ─────────► Aegis Controller 起 Master
    │                        │
    ▼                        │ gRPC: SyncWorkerCreds
init container              │ (发放 JWT,仅此一次)
  tini-init 安装             ▼
  五件套到共享卷         Worker(tini)
    │                        │ 凭 JWT 连 NATS
    ▼                        │ 订阅 {Job}-master
主容器启动                  │ ACK 上报 ────► Master 确认全员订阅
  tini 成为 PID 1            │
  (连不上 Master             │ ◄── bootcheck 指令(群发)
   → RunIndependent          │ ① 环境检测:torch/CUDA/NCCL/cuDNN
   降级为普通任务)           │ ② 健康检测:all_reduce + all_gather
                             │ 结果上报 ────► Master 汇总
                             │
                             ▼
                        阶段 3:训练(RunCommand)
                       ─────────────────────
                       Master 下发 RunCommand 指令
                             │
                             ▼
                       tini runner 拉起用户命令(子进程)
                             │
                             ▼
                       训练运行中(心跳/状态持续上报)
                             │
                             ▼
                       …… 直到 Finish(成功)或 Error(进阶段 4)

阶段 4:故障与恢复的完整分支

训练进程退出 / Pod 被杀 / hang 检出
    │
    ▼ Worker 经 NATS 上报 Error(hang 场景:先只标记不杀进程)
Master:该 Pod Retries+1 → Phase=Troubleshooting → 群发 troubleshoot 指令
    │
    ▼ Worker 执行七项检测:io / gpu / runtime / rdma /
      allreduce 建联 / allreduce 性能 / hang
    │ CheckResponse 逐项上报
    ▼
Master 按检测结果分派:
    │
    ├─ ① 节点无异常 ───► 下发重启指令 ───► tini 只清训练进程
    │                                       (原地重跑,秒级)
    ├─ ② 硬件故障 ────► Worker exit 1 退出 Pod
    │     (gpu/rdma 检测失败)   Job Controller:
    │                            调 Affinity → 重调度到健康节点
    │                            ErrorNodes 黑名单 +1
    │
    └─ ③ 慢节点 ──────► 标记(hang),暂不处理
          (通信性能低/挂起)  等 ckpt saver 完成增量保存(安全点)
                              → hangWithQuit → ShouldQuit → 踢出
                              → 新任务不再调度到该节点
    │
    ◄──── 全员重启完成,Phase 回 RunCommand,Retries 对齐 ────┘
    │
    ▼ 任一 Pod 的 Retries > MaxRetry?
       是 → Job Controller 关闭全部 Pod → Failed(交给用户/上层)
       否 → 继续训练,回到阶段 3

四、通信设计:gRPC 首联 + NATS 常驻

4.1 为什么是两段式

gRPC 只用一次:系统初始化时 Master 通过 gRPC + Protobuf 连到每个 Worker 的 init container 端口,交换初始化信息——核心目的是把 NATS 的 JWT 凭证安全地发给 Worker。
NATS 承担后续全部通信:Master 下发指令、Worker 上报状态/心跳/错误,都走 NATS 的发布/订阅。选 NATS 是因为它轻量、高性能、天然适合 K8s 跨 Pod 通信的 Master-Worker 架构。

4.2 JWT 鉴权:订阅范围收窄到单任务

NATS 的 JWT 不仅做认证,还精确声明订阅/发布的主题范围:

Master 的 JWT:授予订阅全系统主题的权限,实时监控所有任务状态;
Worker 的 JWT:只允许订阅匹配 {JobName}.> 模式的主题——每个 Worker 只能收到自己任务的消息,一个任务的 Worker 不可能听到别的任务;

流程:Master 生成 NKey 并通过 gRPC 传递给 Worker → Worker 据此生成限定订阅范围的 JWT → 连接 NATS。NKey 验证参数通过 K8s Secret 挂载维护,新任务创建时实时更新。

4.3 消息格式

每任务两类主题:{JobName}-worker(Worker 发完成确认/问题报告)、{JobName}-master(Master 下发指令)。消息体:

type Message struct {
    AegisName      string `json:"AegisName"`      // 任务名
    AegisNamespace string `json:"AegisNamespace"` // 命名空间
    PodName        string `json:"podName"`        // 'master' 表示 Master 群发
    Type           string `json:"type"`           // ACK/Order/Error/Finish
    Phase          string `json:"phase"`          // bootcheck/command/troubleshoot
    Data           string `json:"data"`           // 按 Type 不同含义
    TimeStamp      int64  `json:"timeStamp"`
}

四个 Type(ACK/Order/Error/Finish)× 三个 Phase(bootcheck/command/troubleshoot)组合出全部通信语义。

五、状态机与容错计数

5.1 Aegis CR 模型

type AegisSpec struct {
    JobName   string  // 关联的训练作业
    Namespace string
    WorkerNum int32   // Worker 数量
    MaxRetry  int32   // 用户自定义的最大容错次数
}

type AegisStatus struct {
    PodInfos     map[string]PodInfo  // 每个 Pod 的状态
    RecoverCount int32               // 任务级容错次数
    Phase        Phase               // 全局相位
    ErrorNodes   []ErrorNode         // 故障节点记录(名字+原因+时间)
}

type PodInfo struct {
    PodIP             string
    InnerPhase        Phase   // Pod 内部相位
    Finished          bool
    Retries           int32   // 该 Pod 的容错次数
    LastFailedReason  string  // 上次失败原因
    UpdatedTimestamp  int64
}

全局相位流转(代码里由 TLA+ 精确建模):

(init)                                  全员 PASS / 完成
 Pending ────► BootChecking ────► RunCommand ────► Succeed
                  ▲    ▲             │    ▲
       Pod 重启    │    │             │    │ 失败 / 被杀(Pod 重启则计数 +1)
      (不计数)     │    │             ▼    │
                  │    │      Troubleshooting
                  │    │          ▲    │    │
                  │    └──────────┘    │    ├─ 节点无异常 → 回 RunCommand(全员重启)
                  │   待重启者就绪      │    ├─ 节点异常   → exit 1 → 重调度
                  └───────────────────┘    └─ 慢节点     → 等 ckpt → 择机踢出

 出口:RunCommand 全员 Finish → Succeed
      Troubleshooting 中 Retry > MaxRetry → Failed → Job Controller 关闭全部 Pod

每条边的触发条件与计数行为(即 TLA+ Controller 进程的 8 个分支):

起点	终点	触发条件	容错计数
(init)	Pending	CR 创建	—
Pending	BootChecking	全员就绪并完成 NATS 订阅	—
BootChecking	RunCommand	bootcheck 全员 PASS	—
RunCommand	Succeed	全员 Finish	—
RunCommand	Restarting	Pod 被杀/重启	被重启者 Retries+1,recoverCount 加一(带 MaxRetry cap)
非 RunCommand	Restarting	Pod 被杀/重启	不计数(恢复期挂掉不重复罚)
Restarting	Troubleshooting	待重启者就绪	—
Troubleshooting	RunCommand	恢复完成,全员重启	retries 对齐 recoverCount
Troubleshooting	Failed	Retry > MaxRetry	触发全任务终止

5.2 容错计数:一个 3-Worker 的完整演算

计数规则:任务容错次数 = max(各 Pod 的 Retries),上限 MaxRetry。关键细节:只有全局相位处于 RunCommand 时发生的失败才计容错——troubleshooting 阶段挂掉的 Pod 不重复计数。

设 3 个 worker,容错上限 2:

三 Pod 执行任务,worker 0 被 kill → Controller 监控到,记录 index 0 容错 +1,相位置 Troubleshooting。此时计数 (1:0:0),任务 RecoverCount=1;
troubleshooting 期间 index 1/2 任意 Pod 挂掉 → 不计容错(相位非 RunCommand)。计数仍 (1:0:0);
worker 0 拉起,1/2 也跟着重启任务(同步训练,一损俱损) → 相位回 RunCommand,计数 (1:1:1),RecoverCount=1;
worker 1 内部任务出错,主动上报 → 计数 (1:2:1),RecoverCount=2,相位置 Troubleshooting;
worker 1 恢复,Controller 重新下发任务 → 计数 (2:2:2),RecoverCount=2;
又有 worker 失败 → 计数 (2:2:3)。超过上限 2 → Job Controller 监控到 CR 变化,关闭该任务所有 Pod,交给用户/上层处理。

这里"同步训练一损俱损"的特性反而被利用了:恢复时全员重启,计数天然对齐——max 取值即任务级容错次数。

5.3 Fall Back:容错系统自己失败怎么办

容错系统不能成为新的单点。两条回退路径:

Worker 侧:tini 启动时若连不上 Master/NATS,直接置 RunIndependent=true——降级为普通训练任务原样执行,退出码即任务退出码;
Controller 侧:Worker Pod 始终无法与 Aegis Controller 建联时,Job Controller 依据 CR 中的控制权标记判断——正常流程中的只需拉起 Pod;流程外的则清理全部 Pod、记一次容错,回退到普通任务语义。

5.4 各种 Case 的完整决策表

阶段	事件	结果
Bootcheck	掉卡	exit 1,上报 ErrorNode
Bootcheck	执行错误命令	常规检测查不出,正常展示结果
RunCommand	掉卡	上报失败,等待 troubleshooting
RunCommand	pod 被强删	记录容错,拉起后重跑
Troubleshooting	掉卡	exit 1,上报 ErrorNode
Troubleshooting	执行错误命令(节点无问题)	继续执行命令
Troubleshooting	断网	正常执行但无法上报,结果在 Pod 内展示;掉卡 exit 1,否则继续 RunCommand
未开始 Bootcheck	任一 Pod 断网	所有 Pod fall back 为普通训练任务

六、Worker 侧实现:tini 替换容器入口

这是整个方案里最"釜底抽薪"的一步:把 Pod 的启动命令从 bash -c 改为 /opt/Aegis/bin/tini -c。tini 成为 PID 1,用户命令成为它的子进程——于是"清理训练任务进程、原地重跑"变成了可能。

containers:
- name: train
  command:
  - /opt/Aegis/bin/tini
  - -c
  args:
  - |
    python3 train.py   # 用户原始命令,原样保留

tini 内部三个模块:

runner:具体任务的执行与停止。失败时只清理训练进程、不杀 Pod——这是"原地恢复"的实现基础;
worker:通信模块,统一管理 NATS 订阅、上报、接收指令;
工具拷贝:init 阶段把内置检测脚本(bootcheck/troubleshoot)拷进用户容器。

代码里还有一个务实的细节:tini 启动时先尝试连 Master,失败就 DrainNatsConn 并进入独立运行模式——容错链路任何一环不通,训练照常能跑,只是少了容错保护。

bootcheck:训练前的双重自检

所有 Worker 就绪后、正式训练前,Master 下发 bootcheck 指令:

环境检测——验证软件栈一致性:

import torch
print(torch.__version__)                  # PyTorch 版本
print(torch.cuda.device_count())          # 可用 CUDA 设备数
print(torch.cuda.get_device_name())       # 设备名称
print(torch.version.cuda)                 # CUDA 版本
print(torch.cuda.nccl.version())          # NCCL 版本
print(torch.backends.cudnn.version())     # cuDNN 版本

健康检测——用一次真实的集合通信验证拓扑:

dist.init_process_group(backend='nccl')
local_tensor = torch.tensor([rank]).float().cuda()
dist.all_reduce(local_tensor, op=dist.ReduceOp.SUM)   # all-reduce
gather_list = [torch.zeros_like(local_tensor) for _ in range(world_size)]
dist.all_gather(gather_list, local_tensor)             # all-gather 校验

跑一个 all_reduce + all_gather,校验每个 rank 拿到的值是否等于期望——通信拓扑在训练前就被验证过,而不是等到训练几小时后因为 NCCL 挂掉才发现。多 Worker 输出不一致(如 CUDA 版本不同)时,汇总日志直接标"多版本"。

七、异常检测与分级恢复:资源监控、三级策略与慢节点处理

这是容错系统的"感觉神经 + 决策中枢"。设计上分为三层:常态资源监控(持续采集,发现异常)、异常处理机制(检测到异常后的响应链路与分级恢复策略)、慢节点处理(介于正常与故障之间的灰色地带,需要单独策略)。

7.1 资源监控体系:常态采集 + 专项检测

监控分两个层次,解决两类问题。

第一层:常态资源监控(发现"不对劲")

部署专门的异常检测服务,与 Master 通过 protobuf + gRPC 通信,对训练过程做实时在线监测:

指标周期采集:持续收集各训练 Worker 的 CPU、GPU、内存、网络带宽使用数据,刻画每个节点的资源画像——异常的最早信号往往是"指标形态变了"(GPU 利用率掉零、带宽突降、内存只涨不跌),而不是显式的报错;
日志实时监控:对 Worker 的训练日志做正则表达式匹配,捕捉异常模式(NCCL timeout、CUDA error、OOM 关键字等)——很多故障在日志里先于指标暴露;
异常信息结构化记录:检测到异常时,记录四元组——异常 Pod ID、所在节点、异常原因代码、异常发生时间。这份数据既喂给恢复决策(7.2),也沉淀进 ErrorNodes(供后续分析哪些机器反复出问题)。

第二层:专项检测(定位"哪里坏了")

常态监控发现异常后,进入 troubleshoot 阶段的七个专项检测,每项有独立的超时预算——这些预算本身就是故障特征工程:正常情况下 io 检测几秒就该完成,拖到 30 秒说明存储有问题;allreduce 建联/性能检测给足 122s/182s,因为通信类检测天然波动大:

检测项	超时预算	检测什么
io check	30s	存储读写是否正常
gpu check	66s	GPU 健康(`nvidia-smi` 等,ECC/掉卡/温度)
runtime check	毫秒级出结果	运行时环境(CUDA/驱动可用性)
rdma check	25s	RDMA 网卡状态
allreduce connection	122s	集合通信能否建联(拓扑完整性)
allreduce performance	182s/轮	集合通信实际性能(慢节点检测的主力)
hang	—	训练进程是否挂起

七项检测结果以 CheckResult{CheckType, Pass, Result, TimeStamp} 结构经 NATS 上报,Master 汇总进 CheckResponse 决定走哪条恢复路径。每项检测的耗时都是 Prometheus histogram(维度:gpu_type × worker_num × round)——检测本身也被监控,哪种卡型、多大规格的任务、第几轮检测耗时多少,全部可查。

7.2 异常处理机制:从检测到恢复的完整链路

第一步:异常上报。监控服务侦测到异常,立即通过 gRPC 通知 Master(常态监控走 gRPC 直连,Worker 的 troubleshoot 结果走 NATS)。

第二步:全局停训。Master 负责中断该任务下所有 Pod 的训练任务——同步训练下一个 rank 卡住全员都得等,与其空等,不如立刻全体停下进入诊断。

第三步:按相位分派(Master 的 handleError 逻辑):

RunCommand 相位收到 Error → 记录容错日志 + 对应 Pod 的 Retries+1、InnerPhase 置 Troubleshooting → 下发 troubleshoot 指令做专项检测;
BootChecking / Troubleshooting 相位收到 Error → 标记该 Pod 进"需要退出"集合——诊断阶段还挂,说明问题严重,直接走退出重调度;
若消息相位与 CR 当前相位不一致(网络延迟导致的旧消息)→ 按 Finish 处理,防止过期消息污染状态。

第四步:分级恢复策略——先看七项检测的结果如何映射到异常分类,再按分类从最小代价开始逐级升级:

 七项检测结果                      异常分类               恢复路径
──────────────────────          ─────────────         ─────────────────────
 io/gpu/runtime/rdma 全 PASS     进程/瞬时异常    ──►  Level 1 重启进程
 + 训练进程确实退出                                     (tini 只清训练进程,
                                                        原地重跑,秒级)
                                                        │ 失败则降级 ▼
 runtime 报错/环境损坏            容器异常         ──►  Level 2 重启容器
                                                        (进程级救不回来,
                                                        环境可能坏了)
                                                        │ 检测出硬件问题
                                                        │ 则升级 ▼
 gpu check 失败(掉卡/ECC)      硬件错误         ──►  Level 3 隔离节点
 rdma check 失败(网卡坏)                              Worker exit 1 退出 Pod
                                                       → Job Controller 调
                                                         Affinity(annotations/
                                                         labels)
                                                       → Pod 重调度到健康节点
                                                       → ErrorNodes 黑名单,
                                                         后续任务不再上该机
─────────────────────────────────────────────────────────────────────────
 allreduce perf 显著偏低          慢节点(灰色地带) ──►  单独策略(见 7.3):
 hang 检出(挂起)                                      标记不杀 → 等 ckpt
                                                       安全点 → 择机踢出

三级策略的核心思想:能用最小代价解决的绝不用大代价。绝大多数瞬时故障(进程 crash、NCCL 瞬断)在第一级就消化掉了,用户无感;只有真正的硬件故障才会走到节点隔离。而"节点隔离 + Affinity 调整"把故障机器从集群里逐步"淘汰"出去,这正是 1.2 节"必须隔离故障节点才能继续训练"的解法。慢节点则因为"没坏但拖后腿、立即踢会丢进度"的特殊性,单独走 7.3 的择机路径。

7.3 慢节点处理:不立刻杀,选 checkpoint 边界踢

慢节点是最微妙的一类问题,值得单独说。

慢节点的特殊性:它没"坏"——硬件检测项可能全过,但性能低下。同步训练下,每个 step 所有 rank 都要等最慢的那个,一个慢节点能拖垮整个集群的吞吐。但如果发现就立刻踢掉,当前所有训练进度全部作废——为了赶走一个"只是慢"的节点,损失几小时训练进度,得不偿失。

处理设计分四步:

检测:两个来源——allreduce performance check(182s/轮实测集合通信性能,性能显著低于同侪的节点被识别)和 hang 检测(识别训练挂起);
标记:给慢节点打标,记录在案,但不立即行动;
择机踢出:等下一个 checkpoint 完成后,自动结束标记为慢的 Pod——此刻训练进度已安全落盘,踢掉慢节点的损失降到最小。这是整个慢节点策略的灵魂:处理的时机不是"何时发现",而是"何时损失最小";
调度隔离:确保踢出后的新任务/新 Pod 不会被调度到该节点(与硬件故障的隔离机制复用)。

实际链路:标记、ckpt 感知与踢出怎么落地

这四步在工程上由三方配合实现,不是简单的"打标丢给调度器":

① hang 的两级处理(hang ≠ 立即杀)。系统定义了两个相位:hang(只标记)和 hangWithQuit(升级到踢出)。Worker 侧的 tini 发现训练可能挂起(长时间无进度),主动上报 hang 消息;Master 收到后做三件事:把 hang 指令群发给该任务所有 Worker → 每个 Worker 收到后只置一个"挂起中"标记(SetIsHanging(true))——训练进程原样保留,一个不杀;同时把 CheckResponse{CheckType: hang, Pass: false} 记入容错日志,该 Pod 的 Retries+1 进入 Troubleshooting。只有当判定升级为 hangWithQuit 时,Master 才置 ShouldQuit=true——这是"踢出"的决定,但执行要等安全点。

② checkpoint 感知:托付给专用保存引擎,分四层落地。"这一轮 ckpt 是否安全落盘"Aegis 自己不监听文件系统,而是托付给一个专用的 checkpoint 保存引擎,代码里能看到完整的集成骨架:

安装层(已落地):init container 阶段,安装器依次装二进制、各类检测器、checkpoint 组件、控制工具。checkpoint 组件分两步:ckpt.sh 脚本就位;保存引擎的 Python 包经 pip 装进用户容器——安装时先 pip show 探测,用户镜像已带该包则跳过,因为它要和训练框架在同一 Python 进程空间里工作(脚本里有"把 Megatron-LM 路径加进 PYTHONPATH"的注释,印证它面向 Megatron 这类框架做保存代理);
专用 run 标识(已落地):ckpt.sh 以带 fault 前缀的专用 run_id 启动 saver——即这是容错场景专用的保存轮次,与正常训练的周期性 checkpoint 区分。这也解释了这套 checkpoint 引擎的双重使命:增量保存降低全量 ckpt 的 TB 级写盘开销(痛点 1.3),同时为慢节点踢出提供安全点;
进程守护层(逻辑已写好):tini 里有一个 supervisor 模式的守护——循环拉起 saver,saver 意外退出后 3 秒自动重启,stdout/stderr 直通 Pod 日志。保证"保存引擎活着"这件事不依赖人工;
事件对接层(在引擎内部,本仓库不可见):"保存完成 → 放行慢节点退出"的精确事件对接在保存引擎内部实现——开源代码里守护逻辑尚未在主流程接线,这层属于设计意图与内部集成。诚实地讲:能确认的是"容错专用保存 + 进程守护"的集成骨架,不能确认的是完成事件的精确触发机制。

这样分层的意义:即使"择机踢出"的最后一环依赖引擎内部实现,保证 checkpoint 安全区存在这件事本身(saver 常驻、容错专用 run、增量保存)已经是系统设计的一部分,而不是假设用户自己管好 ckpt。

③ 重调度与节点黑名单。慢节点 Pod 退出后,平台侧的 Job Controller watch 到 Pod 退出 + CR 中的 ShouldQuit/ErrorNodes,按 Affinity 把新 Pod 调度到健康节点;而"新任务不再调度到该节点"靠 ErrorNodes 黑名单:Aegis 把故障/慢节点合并进 CR.Status.ErrorNodes(taint 动作),平台调度器消费这份黑名单;运维也能通过控制通道手动 taint/untaint 节点(支持批量)。

总结分工:检测与决定(Aegis Master 状态机)→ 时机控制(checkpoint engine saver)→ 执行(Pod 退出)→ 调度隔离(平台 Job Controller 读 ErrorNodes)。

对比一下三种节点的处置差异,能看出策略的精细度:

节点状态	检测手段	处置时机	处置动作
瞬时故障(进程 crash)	运行时报错	立即	原地重启进程,秒级,不计数损失
硬件故障(掉卡/ECC)	gpu/rdma 检测失败	立即	Pod 退出重调度 + 节点隔离
慢节点(能跑但拖后腿)	allreduce 性能/hang 检测	checkpoint 边界	踢出 Pod + 节点不再接收新任务

7.4 检测与恢复的可观测

所有检测、恢复动作都被指标化(Prometheus):

耗时 histogram(gpu_type × worker_num × round 三维):七项专项检测各自耗时、bootcheck 总耗时、troubleshoot 总耗时;
任务 counter:created / success / failed,按 gpu_type × worker_num 统计;
结构化容错日志(AtlasLog):Pod/Node 信息 + 相位 + 消息,阶段开始/结束逐条记录。

运维因此能回答这类问题:“A800×64 的任务,io 检测 P99 是多少?最近失败率是否上升?失败都发生在哪类检测上?”——检测系统本身就是被监控的对象,故障特征随时间漂移(固件升级、集群扩容)也能被感知到。

八、TLA+ 形式化验证:状态机的数学证明

这是代码库里最有工程含金量的部分。Aegis 的核心是一个 Master-Worker 相位状态机:

相位转换路径多:正常路径、Pod 重启路径、重试耗尽路径、主动退出路径交织;
并发事件:多个 Worker 同时 Finish、Pod 随时重启、消息可能丢失/重发;
边界条件多:Pod 重启时哪些 Worker 需要重置?Retry 超限怎么处理?

单元测试和人工 review 无法覆盖所有状态组合。项目用 TLA+ 做形式化验证,分四步:

8.1 建模:PlusCal 写两个并发进程

用 PlusCal 算法语法编写规约(pcal.trans 翻译成 TLA+,TLC 模型检查器执行),把系统抽象为两个并发进程:

Controller 进程(复刻 Go 代码的状态机)——一个循环里 8 个 either 分支,TLC 非确定性地探索每个分支,等价于穷举所有可能的动作交错:

#	分支	动作
①	初始化	置 Pending,全员 innerPhase=Pending
②	正常推进	全员 Finish 且无 Pod 重启 → `NextPhase(phase)`,retries 对齐 recoverCount
③	RunCommand 中 Pod 重启	置 Restarting,被重启者 retries+1,recoverCount 加一(带 cap)
④	非 RunCommand 的 Pod 重启	只重置被重启者,其他人保持原相位
⑤	Restarting 完成	→ Troubleshooting
⑥	Retry 耗尽	Troubleshooting 中全员 Finish 且超限 → Failed
⑦	Failed 收尾	recoverCount=MaxRetry+1,置 shouldQuit
⑧	retries 对齐	BootChecking/RunCommand 时 retries 同步 recoverCount

Environment 进程(模拟环境扰动)——3 个分支:任意 worker 突然 Finish、任意 Pod 突然重启、外部置 shouldQuit。环境进程的存在,让"任何故障时序"都成为被探索的路径——包括你没想到的"两个 worker 同时 Finish + 第三个 Pod 恰好重启 + 消息重发"组合。

状态变量就是 CR 的镜像:phase/shouldQuit/recoverCount 全局三份,innerPhase[w]/finished[w]/retries[w]/podRestarted[w] 每个 worker 一份。

8.2 定义安全性质:P0 五条 + P1 四条

P0(状态机层)五条:

Invariant	含义
`ValidPhase`	全局 Phase 只能是已知值集合中的成员
`ValidInnerPhase`	每个 Worker 的 InnerPhase 只能是已知值
`RecoverCountBound`	recoverCount 不会无限增长(≤ MaxRetry + 1)
`RetriesBound`	每个 Worker 的 retries 不会无限增长
`PhaseConsistency`	未完成 Worker 的 InnerPhase 必须与全局 Phase 对齐(带豁免:phase=Restarting 时允许 innerPhase 是前驱相位)

P1(消息去重层)新增四条——网络重发会导致同一 Finish 被处理两次,方案用 SequenceNum + 去重防御:

Invariant	含义
`NoFinishDoubleProcessed`	同一 Worker 在同一 seqNum 轮次的 Finish 不会被重复处理
`MasterSeqNumMonotonic`	Master 的 seqNum 单调不减
`WorkerSeqNumBound`	Worker 的 seqNum 不会超过 Master 的
`StaleFinishRejected`	过期 seqNum 的 Finish 被拒绝

8.3 跑 TLC:穷举所有可达状态

AtlasPhase.cfg 配置 Workers={w1,w2,w3}、MaxRetry=2,TLC 从初始状态出发穷举所有可达状态:

P0(状态机):检查 26,778 个 distinct states,5 条 invariant 全部通过;
P1(加消息层):在 P0 上加入 seqNum + dedup 机制,用 CONSTRAINT SeqNumBound(masterSeqNum ≤ InitSeqNum+4)限制状态空间保证有限,检查 44,826 个 distinct states,9 条 invariant 全部通过。

所谓"数学证明"的实际含义就是:四万多个状态空间里的每一条执行路径,9 条安全性质全部成立——不是测试了 N 个用例,而是覆盖了全部可能的用例空间。

8.4 验证过程中抓出的 bug(最有价值的部分)

验证不是一次通过的,spec/PLAN.md 完整记录了迭代发现:

spec 自己的 bug:ValidInnerPhase 忘了包含 Succeed/Failed——shouldQuit 触发时所有 worker 的 innerPhase 被置为 Failed,invariant 立刻被打脸。修正定义后通过;
PhaseConsistency 过严:非 RunCommand 的 Pod 重启(分支④)只重置被重启者,其他 worker 保持原相位,原 invariant 不允许这种组合 → 放宽为允许前驱相位;
反向修复了 Go 生产代码:TLC 发现 RunCommand 中多次 Pod 重启会让 recoverCount 无限增长(分支③没有上限),于是给 Go 代码加了 recoverCount = Min(recoverCount+1, MaxRetry) 的 cap(对应 workflow.go:87-89)——形式化验证驱动了生产代码的修复,这是它超越测试的核心价值;
spec 与 Go 行为不一致:PENDING exception + 已 dedup 的过期 Finish 在 spec 里被错误拒绝,Go 代码实际允许通过(handler 允许 PENDING 例外,dedup 检查跳过重复)→ spec 增加 Case 4 分支对齐;
工程细节:模块名错误导致 TLC 找不到文件、cfg 漏列 invariant、PlusCal 缺翻译块——形式化工具链本身也有工程坑。

8.5 小结:什么时候值得上 TLA+

这套验证的成本并不高(核心规约数百行),但收益明确:状态机 + 并发消息 + 重试的组合空间,恰好是单元测试最薄弱的区域——测试只能覆盖你想到的路径,TLC 穷举你没想到的。对一个"计数错了就误杀/漏杀训练任务"的系统,这是值得的保险。

为什么消息层的验证尤其重要:容错计数(第五节那个 (1:0:0)→(2:2:3) 演算)是全系统的决策依据——如果消息重发导致一次失败被记两次,任务会提前被判死;如果 seqNum 乱序导致旧 Finish 覆盖新状态,计数会失真。TLA+ 从数学上证明了在任何消息时序、任何 Pod 重启组合下,计数和相位都不会进入非法状态。这是"测试覆盖不到的角落,用穷举证明补上"。

九、容错日志

Master 负责汇总容错日志(阶段开始/结束、各 Worker 上报),结构化记录:

type AegisLog struct {
    PodName, PodIP, NodeName, NodeIP string
    TimeStamp int64
    Phase     string
    Message   string
}

日志收集链路在方案评审时对比过五种路径(经 event/watch/grpc 转发 vs 直接落库 vs 日志采集管道),最终选择采集管道方案:controller 把容错日志输出到日志目录,由采集器(loggie)收集到后端(openobserve),前端查询直接转发——避免了为日志在控制面组件间增加转发链路。用户侧展示形如:“您的最大重启次数为 N 次,目前已重启 M 次;当发现硬件问题,系统将自动帮您重调度,目前已发生 K 次”。

十、设计哲学小结

粗粒度恢复 → 细粒度恢复:从"整任务重调度 + TB 级 ckpt 读盘"压缩到"单 Pod 原地清进程重跑",恢复代价降了一个数量级;只有硬件故障才升级到重调度 + 节点隔离。
同步训练的弱点变设计资产:"一损俱损"意味着恢复必然全员重启,计数天然对齐——max(Pod.Retries) 就是任务容错次数,不需要复杂的分布式对账。
容错系统自身必须可降级:连不上 Master 就退化为普通任务(tini 的 RunIndependent),任何容错链路故障都不会阻断训练本身。
通信按生命周期分层:一次性凭证分发用 gRPC(点对点、可靠),常驻指令/上报用 NATS(pub/sub、轻量),JWT 把每个 Worker 的订阅范围钉死在自己任务里。
状态机用形式化方法兜底:并发 + 重试 + 消息重发的组合空间靠测试覆盖不完,TLA+ 穷举可达状态证明 9 条安全性质——分布式容错这种"错了就丢训练进度"的系统,值得上这种保险。
谁创建谁删除的权限切分:两个 controller 各管一摊,CR 的创建权与更新权分离,避免双写冲突。

一句话:Aegis 把"训练挂了怎么办"从事后人工排查(数小时、烧几百张卡)变成了平台内置的自动流水线——检测(bootcheck/异常检测)→ 定位(troubleshooting)→ 分级恢复(进程/容器/硬件)→ 记账(CR 容错计数)→ 兜底(上限关闭),并用 TLA+ 证明了这套状态机在任意并发与消息时序下不会出错。

参考

OPT-175B LogBook、BLOOM-176B 训练报告(故障统计为公开数据)
TLA+ Spec 语言 / TLC Model Checker
NATS JWT 鉴权与 NKeys
关联:图解 Megatron 的 SP/CP/EP、分布式训练核心技术:从 FSDP 到 All-Reduce