训练框架

摘要

训练框架是 MaxText 的替代设计，目标是把大规模训练从"过度依赖 XLA 黑盒、层间编排不可控、控制面散落"的现状，重新组织为 XLA / Pallas / Framework 三分责任的可演进、可观测、可验证、可协作的工程体系。

它的存在填补四个空缺：跨层 overlap 调度的归属、控制面的统一 owner、跨团队契约面、训练生命周期的稳定抽象。这四个空缺在 XLA + Pallas + 手写脚本的组合下无法被填补。

训练框架同时服务训练框架团队、编译器团队和 kernel 团队。三方应围绕同一套训练任务描述、性能口径、诊断材料和验收标准协作，而不是各自维护割裂的判断依据。

在能力边界上，XLA 退回到编译 / lowering、基础图优化、device 调度、跨层 / 全局 buffer 分配；Pallas 承担层内 hot path kernel 和层内 buffer；Framework 承担层间编排（提前 / 延后发射、通信 overlap）、控制面（checkpoint、恢复、指标、调试）和跨团队契约面。

目标读者与定位

本文面向训练框架团队、编译器团队和 kernel 团队，用于对齐大规模训练场景对框架抽象、性能表达、观测诊断、正确性验证和易用性的共同要求。

本文偏向设计原则与协作契约，不定义某个具体训练框架的 API 细节，也不替代 Pallas kernel 开发规范、训练指标体系、性能评估体系、故障恢复机制或编译器工具链文档。它回答的问题是：训练框架应承担什么系统责任，以及这些责任如何转化为框架团队、kernel 实现和编译器团队之间的稳定协作边界。

背景

大规模训练任务具有长周期、高成本、有状态、强并行、强异构、强依赖编译器的固有特性。一次训练通常涉及数百至数千加速器、连续运行数天到数周，期间任意一步出错都可能毁掉整次投入。这类任务无法靠"跑一次看看"来验证，必须把可追踪、可复现、可诊断作为系统能力的起点。

团队协作上，训练任务横跨训练框架、编译器、kernel 实现、运行时与硬件支持四方。这四方视角不同、工具不同：框架看 Python 代码、编译器看 HLO/StableHLO 图、kernel 看指令、runtime 看事件流。如果四方的程序表达和诊断材料不互通，性能与稳定性问题就会在团队之间反复来回，无法收敛。

业务与硬件上下文持续扩展：稠密预训练、MoE、长序列、多模态、强化学习对模型结构、并行策略、精度、checkpoint 的要求各不相同；TPU v5/v6/v7 在内存层级、互联拓扑、计算单元上差异显著。框架如果不能在稳定抽象下表达这些差异，每次接入新模型或新硬件都会演变成对框架自身的修改。

历史背景上，当前栈以 MaxText 为参考实现。MaxText 在快速验证、社区对齐方面发挥了作用，但在大规模、长周期、跨团队协作的训练场景下暴露了结构性限制：过度依赖 XLA 黑盒做性能决策、层间编排不可控、控制面散落在脚本和约定中。本文描述的训练框架是对 MaxText 的替代设计，不是包装或扩展。

核心挑战

训练框架要解决的问题不是"做一个能跑训练的入口"，而是绕开下列具体失败模式。这些失败模式都来自大规模训练中的实际经验。

长跑异常定位困难

大规模训练中，NaN、Hang、数值漂移往往在发生后数小时才被察觉，根因可能埋在数千 step 之前的某个 kernel、某次通信、某批数据或某次编译选择。当前栈缺乏分层定位能力：异常发生时只能拿到模糊的 step 编号和错误码，无法定位到具体的 tensor、op、rank、kernel 版本，更无法快速缩小到小规模复现路径。

性能退化归因困难

step time 或 MFU 退化时，根因可能在 Pallas kernel 实现、跨层通信、XLA 编译决策、控制面开销、数据加载、硬件降频任意一层。当前栈把性能问题归到端到端 step time 一个数字上，没有结构化材料让框架、编译器、kernel、runtime 四方在同一组数据上拆解根因。

新模型 / 硬件 / 策略接入成本高

每接入一个新模型结构（MoE、长序列、多模态）、新硬件（TPU v7）、或新并行策略（专家并行、序列并行），都意味着要改框架代码、写特殊脚本、维护临时分支。演进稳定性的缺乏使得工程成本随场景数量线性增长。

跨团队上下文对不上

框架写的 Python 代码、编译器看的 HLO/StableHLO 图、kernel 跑的指令、runtime 报的事件流，四方使用不同的程序表达、不同的命名约定、不同的诊断材料。一个性能或数值问题需要跨团队协同时，大量时间花在"翻译"而不是"分析"上。

XLA 黑盒限制了性能与可控性

XLA 当前对训练性能负有过多决策权，但作为黑盒编译器，它至少在三方面做不好：

• 优化决策不可控：自动 fusion、layout、sharding 决策经常覆盖用户意图，需要逆向打 flag 或 annotation 才能纠正。
• 内存规划不透明：HBM 峰值、buffer aliasing、复用决策对外不可见，内存问题只能反向推断。
• 跨层通信 / 计算 overlap 做不好：跨层通信的生命周期管理是 XLA 的弱项，无法实现可靠的"提前发射 / 延后等待"，导致训练性能低于硬件潜力。

训练框架的存在意义

训练框架不是 MaxText 之上的薄包装，而是 MaxText 的替代。它的存在是为了填补下面四个空缺，这些空缺在 XLA + Pallas + 手写脚本的组合下无法被填补。

跨层 overlap 调度的归属

TPU 训练性能高度依赖跨层通信与计算的重叠——决定哪些层提前发射、哪些 wait 延后、哪些 buffer 持有到下一步。XLA 作为编译器，在做跨层 overlap 时受限于图级静态分析，无法表达动态的层间依赖；Pallas 作为单 kernel 实现机制，无法跨层调度。这一层编排必须由 framework 承担，作为系统的一等能力而不是局部优化。

控制面的统一 owner

checkpoint、恢复、指标采集、调试、host offload、故障注入、降级路径——这些控制面逻辑必须在某处统一归属。如果散落在各个 kernel、临时脚本、启动配置里，会出现不可解释的 stall、不一致的状态、不可复现的失败。framework 作为控制面 owner，保证控制决策可记录、可审计、可回放。

跨团队契约面

framework / compiler / kernel / runtime 四方需要共用同一组材料：训练任务的程序表达、性能口径、调试产物、复现包、回归基线。framework 提供这层契约面，让四方的判断依据一致，而不是各自维护割裂的视图。

训练生命周期的稳定抽象

init / compile / run / save / restore / recover / evaluate 这套生命周期不能为每个模型、场景、硬件重写一次。framework 提供稳定的生命周期抽象，让模型结构、并行策略、精度策略、checkpoint 策略的变化在配置和扩展点层面表达，而不是改 framework 代码本身。

设计取舍与边界

与 MaxText 的关系

训练框架是 MaxText 的替代设计，不是包装、不是扩展、不是并行项目。MaxText 在快速验证、社区对齐方面有价值，但在大规模、长周期、跨团队协作的训练场景下暴露了结构性限制。本文描述的训练框架是推倒重写。

XLA / Pallas / Framework 三分

新栈的责任分配明确把 XLA 退回到它最擅长的位置，把性能关键路径交给 Pallas，把层间编排交给 framework。三方边界如下：

职责	XLA	Pallas	Framework
编译 / lowering（HLO → machine code）	✓
基础图优化（死代码消除、常量折叠、代数化简）	✓
device 调度（async launch、sync point、资源约束）	✓
跨层 + 全局 buffer 分配	✓
层内 buffer 分配		✓
层内 hot path kernel 实现		✓
自动 fusion 决策	✗（不再做）	替代实现
层间编排 / 通信 overlap / 提前 / 延后发射			✓
控制面（checkpoint / 恢复 / 指标 / 调试）			✓
跨团队契约面			✓

XLA 退回的具体含义：不再做自动 fusion 决策——fusion 边界由 Pallas kernel 边界和 framework 调度边界显式决定；不再做层内 buffer 分配——层内 HBM/VMEM 由 Pallas kernel 自治。XLA 仍然做编译、基础图优化、device 调度、跨层和全局 buffer 分配。

Pallas 的具体职责：承担层内 hot path——attention、scan、recurrent、稀疏/专家路由、特殊归约等。Pallas kernel 自治层内 buffer，并通过稳定接入契约（shape、dtype、layout、sharding、误差范围、fallback）暴露给 framework。

Framework 的具体职责：承担层间编排与控制面——跨层 wait / prefetch / overlap、Tensor 持有边界、host offload、checkpoint 触发、指标采集、故障恢复。

总体要求

大规模训练任务长周期、高成本、有状态、强并行、强异构且强依赖编译器。训练框架需要把这些复杂性以稳定抽象显式呈现，并在以下六项核心能力上提供系统能力。各项能力的具体口径见对应章节。

能力	定义	针对挑战	详见
表达训练场景复杂性	把从配置生成到结果验证的完整链路组织成可追踪、可解释、可复现的抽象	长跑异常定位、跨团队对齐	设计原则、协作契约
灵活演进	在稳定抽象下复用核心生命周期，覆盖模型/数据/优化器/并行/精度/checkpoint/profiling/调试的变化	接入成本	灵活性要求
性能闭环	把性能建模、表达、归因前置到框架设计，使性能可定义、可观察、可复现、可归因	性能归因、XLA 黑盒	性能要求
观测与诊断基线	默认产出训练/数值/系统/checkpoint/异常/配置健康度，支持分层定位	长跑异常定位、性能归因	诊断与验证 § 持续观测、§ 异常诊断
可验证性	提供从小规模正确性到大规模性能的连续验证路径	长跑异常定位、性能归因	诊断与验证 § 主动验证
规模化使用成本	标准化常见路径、配置化复杂能力、结构化失败信息	接入成本	易用性要求

设计原则

统一入口与分层抽象

训练框架应提供统一的任务入口和生命周期管理，同时保持模型、数据、优化器、并行策略、编译选项和运行时能力之间的分层边界。统一入口用于降低使用和治理成本，分层抽象用于避免框架把所有差异都固化为不可维护的特殊逻辑。

层内 Kernel 与层间编排解耦

训练框架应把层内 critical path 和层间系统编排作为两个不同层次管理。层内高性能路径由 Pallas kernel、custom VJP、primitive wrapper 或等价机制表达，框架为其提供稳定接入契约（详见 性能要求 § 层内 Pallas Kernel Critical Path）。

层间编排由训练框架负责（调度对象与策略详见 性能要求 § 层间调度、预取与延后等待）。框架应避免让临时控制逻辑破坏层内 kernel 的可分析性和可复用性，也应避免把层间调度隐式埋入单个 kernel，使端到端性能无法被解释。

控制面与数据面分离

训练框架需要区分高频数据路径和控制面逻辑。高频数据路径应尽量稳定、可编译、可 profile，并具备清晰的 kernel、子图和通信边界；控制面负责配置解析、任务状态、动态策略选择、host offload 策略、checkpoint 管理、故障恢复、指标聚合和诊断材料收集。

控制面逻辑可以灵活演进，但其决策结果必须能回写到配置快照、日志、指标和 trace 中。任何影响性能或正确性的控制面决策，例如 offload 开关、prefetch 深度、wait 位置、重试策略或故障降级路径，都应具备可追踪的上下文。

配置可组合、可校验、可追溯

训练配置应能够表达模型结构、数据、训练超参、并行策略、精度、checkpoint、profiling、调试和环境依赖。配置需要支持组合与覆盖，但最终生效配置必须可校验、可记录、可复现。

训练生命周期显式化

训练框架应显式管理初始化、编译、执行、评估、checkpoint、恢复、profiling、异常处理和资源释放。生命周期中的关键状态变化应能够被日志和指标捕获，避免隐式状态影响正确性和性能判断。

并行策略与运行时能力解耦

并行策略应作为训练任务的显式描述，而不是散落在模型代码、数据加载、编译配置和启动脚本中的隐式约定。框架需要区分“用户想要的并行语义”和“底层运行时实际采用的通信与执行方式”，并为两者之间的映射提供可检查的描述。

指标口径统一

训练框架应统一性能、稳定性和数值正确性相关指标的命名、采样、聚合和上下文标签。指标不应只服务看板展示，还应能支持回归判断、瓶颈归因、异常定位和跨团队沟通。

调试与验证能力前置

训练框架应在设计阶段就考虑调试和验证路径。关键训练任务应能够在不同规模下运行，能够切换到 synthetic data、局部模型、短步数、固定随机种子、额外数值检查或编译产物导出等模式，以支持快速定位。

故障恢复作为核心路径

checkpoint、恢复、任务重启、数据进度恢复和异常清理应是训练框架的一等能力。它们不应只在任务失败后补救，而应作为训练生命周期的一部分被持续验证。

框架与编译器共享问题上下文

训练框架应为编译器团队提供足够的上下文来理解训练程序：模型结构、输入 shape、dtype、sharding、mesh、并行策略、关键配置、编译选项、运行指标和复现材料。编译器团队的诊断结果也应能够回到训练框架的指标、配置和调试流程中。

性能要求

训练框架需要把性能作为程序表达和系统协作的核心约束。性能要求不应只落在最终吞吐或利用率数字上，而应指导框架如何表达 Tensor、通信、编译单元、并行策略和算子边界，使性能优化具备结构化抓手。其中，层间通信和显存生命周期编排是训练框架的核心作用：框架需要决定哪些状态跨层保留、哪些中间结果可释放或重算、哪些通信可以前置、延后或与计算重叠。

层内 Pallas Kernel Critical Path

层内 critical path 的高性能代码应主要由 Pallas kernel 或等价 custom kernel 实现。训练框架不应把 attention、scan、recurrent block、稀疏/专家路由、特殊归约或其他高频热点退化为不可分析的 Python 组合逻辑；对必须保留在框架层表达的逻辑，也需要明确其不在性能关键路径上，或提供迁移到 kernel/primitive 的边界。

每个进入训练主路径的 kernel 都应有清晰契约：输入输出 shape、dtype、layout、sharding、tile/block 参数、accumulator 精度、边界处理、padding 语义、反向实现、可接受误差、fallback/reference 实现、benchmark 结果和 profiling 标识。框架应把这些契约与训练配置、编译产物和运行指标关联起来，使 kernel 性能问题、数值问题和编译问题能够被定位到具体实现版本。

Pallas kernel 的性能不应只通过端到端 step time 间接判断。框架应支持记录 kernel 级耗时、调用次数、输入规模、FLOPs、访存量、roofline gap、MXU/VALU/Scalar 利用、HBM/VMEM 行为和关键 tile 参数，并支持与 reference 实现或历史版本对比。

Tensor 生命周期管理

训练框架应显式关注 Tensor 从创建、布局、分片、传递、复用到释放的生命周期。模型参数、激活、梯度、优化器状态、临时 buffer、通信 buffer 和 checkpoint 相关 Tensor 都应有清晰边界，避免隐式拷贝、重复 materialization、不可解释的持有关系和难以定位的内存峰值。

Tensor 生命周期管理不仅是内存优化问题，也影响编译器的 buffer 规划、别名分析、算子调度和跨 step 稳定性。框架应尽量提供有利于编译器理解和优化的 Tensor 使用模式。

在多层模型中，Tensor 生命周期还需要跨层编排。训练框架应明确层输入、层输出、残差、KV/state、梯度、通信中间结果和可重算激活的持有边界，使显存占用不随层数简单线性膨胀，并为编译器提供可分析的释放、复用和重算机会。

通信作为一等公民

大规模训练的性能由计算和通信共同决定。训练框架需要把通信路径作为训练程序表达的一部分，而不是把通信视为并行策略的副作用。

框架应能够描述 collective、point-to-point、跨 host 通信、通信与计算重叠、通信 buffer、通信拓扑以及通信策略与并行策略之间的关系。通信应能够被观测、被建模、被调优，并能够和具体训练 step、并行切分、算子边界关联起来。

通信也需要在层间被编排。训练框架应描述每层或每类 block 在前向、反向和参数更新阶段的通信需求，并将通信触发点、通信粒度、buffer 生命周期和计算重叠关系纳入训练程序表达。

层间调度、预取与延后等待

层间编排是训练框架的核心职责。框架应能够显式表达每个 layer、block 或 stage 之间的数据依赖、通信依赖、buffer 依赖和控制依赖，并基于这些依赖安排预取、延后等待、通信/计算重叠、重算、释放和 offload。

预取和延后等待不应只作为局部优化散落在实现中，而应成为可配置、可观测、可回放的调度策略。框架需要记录 prefetch 深度、触发点、wait 位置、未隐藏等待时间、overlap ratio、队列深度和失败降级路径，使层间调度的收益和风险可以被量化。

Host Offload 与控制面逻辑

训练框架应为 host offload 提供统一策略层，覆盖可 offload 对象、触发时机、传输粒度、带宽占用、队列深度、阻塞点、恢复行为和正确性约束。offload 策略应能服务显存压力控制、checkpoint、优化器状态、低频统计、诊断材料和特定运行时能力，但不能以不可观测的方式引入隐式 stall 或破坏训练语义。

控制面逻辑需要能够灵活表达任务状态、资源状态、策略切换、异常处理和恢复动作。框架应保证控制面决策可以被指标化和审计，例如一次 step 中哪些对象被 offload、哪些 wait 被延后、哪些通信被 overlap、哪些 checkpoint 被跳过或重试。

最小化子图编译

训练框架应尽量减少非必要的完整图展开、重复编译和编译缓存失效。一次模型训练中的层结构通常是重复且可穷举的，框架不应默认把所有层都展开成一个不可复用的大图，而应识别层、block、stage 等可重复结构，将其表达为稳定、可组合、可复用的编译子图。

在组合多层模型时，训练框架应尽量复用这些子图结构，使编译器能够针对有限数量的结构单元进行分析、缓存和优化，而不是对每个重复层实例重新承担完整编译成本。可复用子图还应携带清晰的层间契约，包括输入输出、通信需求、显存持有边界和可释放/可重算状态。动态行为、调试逻辑、指标采集、checkpoint 控制和 Python 侧编排不应随意破坏这种结构复用。

基于 JAXPR / StableHLO 的 Roofline 与 Dataflow 构建

训练框架应支持从 JAXPR、StableHLO 或等价 IR 中提取训练 step 的计算、访存、通信和数据依赖结构，为 roofline 分析、dataflow 分析和跨版本性能比较提供统一输入。

这类 IR 级表达应成为训练框架和编译器团队之间讨论性能的共同语言。框架侧需要保证导出的程序表示能关联回模型结构、配置、并行策略和运行指标；编译器侧需要基于这些表示解释瓶颈、优化机会和性能回归。

在计算图 level，框架还应支持数学逻辑校验。校验对象包括 shape 与 sharding 一致性、mask 与边界语义、并行切分前后的等价性、前向/反向对应关系、优化器状态更新、checkpoint restore 后状态一致性，以及 custom kernel 与 reference 实现之间的数学语义一致性。图级校验应能够和 roofline、dataflow、profile、trace 使用同一套程序表示，避免性能分析和正确性分析割裂。

并行策略的描述与通信策略

并行策略不能只表现为启动参数或局部 sharding 注解。训练框架需要明确描述数据并行、张量并行、流水并行、专家并行、序列并行等策略如何映射到 mesh、layout、sharding 和通信模式。

框架应帮助用户理解并行策略带来的计算、通信、内存和编译代价，也应帮助编译器理解策略意图，从而在通信调度、layout 选择、重分片和融合边界上做出更合理的优化。

算子融合

训练框架应为编译器提供有利于算子融合的程序表达，减少无意义的模块边界、格式转换、中间结果固化和 host/device 往返。

框架层面的模块组织、精度策略、自定义 kernel 接入、调试逻辑和指标采集都可能影响融合机会。训练框架应避免为了局部封装便利破坏全局融合空间，同时为必须保留的边界提供清晰语义。

灵活性要求

训练框架的灵活性应体现在稳定抽象下的可组合能力，而不是无限制暴露内部实现细节。

• 模型灵活性：支持不同模型结构、层组织方式、激活重计算策略、自定义算子和特定模型族的扩展需求。
• 数据灵活性：支持真实数据、synthetic data、多数据源、数据混合、数据进度恢复和数据相关诊断。
• 优化器灵活性：支持不同优化器、学习率策略、梯度处理、混合精度和优化器状态切分方式。
• 并行灵活性：支持数据并行、张量并行、流水并行、专家并行、序列并行和混合并行的组合，并能够描述其对通信、内存和编译的影响。
• Kernel 灵活性：支持将层内 critical path 替换为不同 Pallas kernel、custom kernel 或 reference 实现，并保持统一的调用、验证、指标和回退契约。
• 控制面灵活性：支持 host offload、prefetch、wait placement、checkpoint、恢复和诊断策略的配置化切换，并保证每次策略选择可记录、可复现、可比较。
• 调试灵活性：支持按需打开数值检查、profile、IR 导出、trace、日志增强和故障注入，而不破坏主训练路径的稳定性。

协作契约

训练框架团队负责

• 定义稳定的训练任务抽象，包括模型、数据、优化器、并行策略、精度策略、checkpoint 和运行环境。
• 管理训练任务生命周期，包括启动、初始化、编译、执行、评估、保存、恢复、终止和资源释放。
• 产出统一的配置快照、日志、指标、profile、trace 和复现材料。
• 定义 Pallas kernel 或等价 custom kernel 的接入契约（契约内容详见 性能要求 § 层内 Pallas Kernel Critical Path）。
• 主导跨层 overlap 调度：提前 / 延后发射、跨层 wait / prefetch / 通信 overlap（详见 性能要求 § 层间调度、预取与延后等待）。
• 维护 XLA / Pallas / Framework 三方边界契约：保证三方职责划分在代码、配置、IR 层面一致。
• 将模型、数据和并行策略变化组织成可验证输入，避免将训练差异固化为不可复用脚本。
• 为性能、正确性和稳定性问题提供可定位的框架层上下文。

编译器团队负责

• 在 XLA 退回后的职责范围内工作：编译 / lowering、基础图优化（死代码、常量折叠、代数化简）、device 调度、跨层 + 全局 buffer 分配。
• 不再做自动 fusion 决策——fusion 边界由 Pallas kernel 边界和 framework 调度边界显式决定。
• 基于训练框架提供的上下文分析编译、执行、内存、通信和算子优化问题。
• 提供可解释的编译诊断、性能归因、验证工具和回归判断依据。
• 支持对 Pallas kernel、custom call、通信原语、host offload 和编译子图之间关系的分析，使 kernel 级优化和图级优化能够被统一解释。
• 将编译器优化行为暴露为可观测、可回放、可比较的证据，避免优化结果成为不可解释的黑盒。

Kernel 团队负责

• 按 Pallas kernel 接入契约交付层内 hot path kernel（attention、scan、recurrent、稀疏 / 专家路由、特殊归约等）。
• 自治层内 buffer 分配，提供稳定 shape、dtype、layout、sharding、误差范围、fallback、benchmark 和 profiling 标识。
• 配合 framework 完成层间编排所需的接口（同步点、prefetch 触发、wait 暴露）。

三方共同维护

• 训练程序的标准描述方式。
• 性能分析与正确性验证的基础口径。
• 层内 kernel 契约、层间编排契约和图级分析材料之间的对应关系。
• 最小复现材料的内容与格式。
• 新模型结构、新并行策略、新硬件环境的接入边界。
• 性能回归、数值回归和稳定性回归的判断标准。

诊断与验证

诊断与验证服务三件事：训练运行中持续观测、训练行为主动验证、异常发生时分层定位。三者共用同一套程序表达、上下文材料和指标口径，避免观测、验证、诊断成为三套割裂的体系。

持续观测

训练框架应提供统一的观测基线，使训练任务在运行中能够被理解，在异常时能够被定位，在回归时能够被比较。

• 训练健康度：loss、学习率、batch、step、跳步、NaN/Inf、梯度范数、参数规模和训练进度。
• 数值健康度：activation、gradient、optimizer state、weight update、logits、loss 分解和关键中间值分布。
• 系统性能：step time、tokens/s、MFU、内存、通信、编译时间、数据加载、checkpoint I/O 和资源利用。
• Kernel 性能：Pallas kernel 耗时、调用次数、输入规模、FLOPs、访存量、roofline gap、tile 参数、MXU/VALU/Scalar 利用和版本信息。
• 层间调度：prefetch 命中率、延后 wait 时间、overlap ratio、通信等待、buffer 持有时间、重算开销、host offload 流量、offload 队列深度和阻塞时间。
• 任务状态：任务启动、初始化、编译、运行、评估、保存、恢复、终止和异常状态。
• 恢复状态：checkpoint 成功率、checkpoint 延迟、restore 时间、回滚 step 数、恢复成功率、重试次数、丢失进度和恢复后的数值/性能对齐状态。
• 硬件诊断：设备、host、slice、rank、拓扑、心跳、collective 进度、首个 NaN/Inf 位置、hang 检测时间、异常前最后成功 op/kernel 和相关硬件/runtime 事件。
• 上下文标签：代码版本、配置快照、模型、数据、硬件、并行策略、运行环境和关键编译选项。

指标体系的详细设计由训练指标体系文档承接。训练框架在这里的责任是保证指标能够稳定产生，并且能与训练任务上下文绑定。

主动验证

训练框架需要为正确性、性能和稳定性提供连续的验证路径。大规模训练中的问题应尽量能够被缩小、隔离和重放，避免只依赖线上大规模任务暴露问题。

• 小规模可验证：支持小模型、小 batch、短步数、synthetic data、固定 seed 和局部并行策略，用于快速验证训练逻辑。
• 数学逻辑可验证：支持在计算图 level 校验模型结构、mask、路由、状态更新、前向/反向关系、并行切分语义和 custom kernel/reference 等价性。
• 数值可验证：支持对 forward、backward、loss、梯度、参数更新、优化器状态和 checkpoint restore 的对齐与漂移检查。
• 数值范围可验证：支持对数值误差累计进行定量范围分析，包括 dtype、accumulation order、归约树、重计算、混合精度、kernel 实现和并行策略带来的误差预算、容忍区间和随 step 漂移趋势。
• 分布式可验证：支持验证不同并行策略下的 shape、sharding、通信语义、数据一致性和恢复一致性。
• 性能可验证：支持基于相同配置、相同输入和相同硬件环境比较性能变化，并能保留足够上下文用于归因。
• 编译可验证：支持导出和保留足够的程序表示、编译选项和运行信息，使编译器问题能够被复现和分析。

异常诊断

异常诊断需要把硬件问题、运行时问题和模型数值问题分层定位，并把硬件或 runtime 事件关联回训练上下文（模型层、kernel、并行策略、通信组、输入 batch、checkpoint 状态、配置变更），使训练框架团队、编译器团队和硬件/运行时支持人员能够基于同一组材料定位问题，而不是只看到底层错误码。

故障恢复预案。 框架应为常见故障类型提供完整预案，包括进程失败、节点失败、设备失败、网络异常、存储异常、数据异常、编译失败、checkpoint 损坏、数值异常和 hang。每类故障都应定义检测信号、保护动作、checkpoint 策略、恢复入口、回滚范围、重试策略、降级路径、人工介入条件和恢复后验证方式。

恢复路径统计与演练。 框架需要持续记录 MTBF、MTTR、RTO、RPO、恢复成功率、重试次数、丢失 step、checkpoint 写入耗时、checkpoint restore 耗时、恢复后首个有效 step、恢复后的 loss/数值指标对齐和恢复导致的吞吐变化。关键任务应通过故障注入或恢复演练验证这些路径，而不是等真实事故发生后再补齐。

NaN/Inf 定位。 框架应捕获首个异常 step、首个异常 tensor、op/kernel 名称、rank/device/slice、输入输出分布、dtype、sharding、最近的 checkpoint、最近的编译产物和相关 Pallas kernel 版本，并支持缩小到小规模或 reference 路径复现。

Hang 定位。 框架应记录心跳、step 进度、collective 进度、barrier 状态、device/host 状态、最后成功 launch、最后成功 collective、等待中的通信对象、超时阈值、runtime 日志、trace/profile 和拓扑信息。诊断材料应支持区分模型代码死锁、数据输入阻塞、通信不匹配、runtime hang、host offload 阻塞、设备异常和硬件互联问题。

易用性要求

训练框架的易用性应服务规模化使用和跨团队协作，而不是只服务单次任务启动。

• 入口清晰：训练任务应有标准启动方式、标准配置加载方式和标准输出目录。
• 默认可靠：默认配置应覆盖常见训练路径，并尽量提供安全的 checkpoint、日志、指标和异常处理行为。
• 错误可操作：错误信息应指向可能的配置、数据、并行、编译、运行时或环境问题，而不是只暴露底层异常。
• 扩展可控：新增模型、优化器、并行策略、指标、profile 或调试机制时，应接入明确扩展点。
• 调试低成本：用户应能以低成本切换到更小规模、更强日志、更强检查或更强 profile 的运行模式。

非目标

为了明确边界，下面这些事情训练框架不做。不做不代表不重要，而是由其他系统或团队承担，框架只提供接口或协作面。

• 不负责模型结构设计本身：模型架构由模型团队和研究人员定义，框架只提供模型组装、并行切分、生命周期管理的稳定抽象。
• 不负责 XLA 编译器内部：XLA 保留 lowering、基础图优化、device 调度、跨层 / 全局 buffer 分配等职责；框架不干涉编译器内部决策，只通过稳定的程序表达和编译选项影响它。
• 不负责 Pallas kernel 实现：kernel 由 kernel 团队按 pallas-kernel 规范开发，框架只定义接入契约（shape、dtype、layout、sharding、误差范围、fallback）和调度边界。
• 不负责硬件 fleet 管理 / 物理 TPU 分配：由 xpk / SkyPilot / GKE 承接，框架消费分配好的 mesh，不管物理切片。
• 不负责训练数据采集与标注：由数据 pipeline 承接，框架只提供数据接入、进度恢复、混合采样的接口。
• 不负责任务调度排队 / 优先级 / 多租户：由任务调度与管理体系（见 Infra 路线图）承接，框架只暴露标准启动、状态、终止接口。
• 不替代 JAX 作为前端：framework 基于 JAX 生态（JAX / Pallas / StableHLO / XLA），不重写前端 autodiff、不重写 numpy-style API。
• 不优化单机 demo 易用性：框架目标是大规模训练稳定运行，单机 demo 不是首要设计目标。

优先级

能力按"对 MaxText 替代目标的贡献度"分三档。优先级指导资源分配，不代表 P2 不重要。

优先级	能力	理由
P0	XLA 边界重新划分（自动 fusion 退出、Pallas hot path 接入）	决定新栈结构性差异，不做就不是 MaxText 替代
P0	跨层 overlap 调度（提前 / 延后发射、跨层 wait / prefetch / overlap）	framework 的核心不可替代职责，XLA / Pallas 都做不了
P1	控制面统一（checkpoint / 恢复 / 指标 / 调试 / offload / 降级）	框架可用性的最低门槛，散落即不可维护
P1	生命周期稳定抽象（init / compile / run / save / restore / recover）	让框架能被使用而不只是设计
P2	观测与诊断基线（NaN / Hang 定位、分层归因）	长跑稳定性，依赖 P1 控制面
P2	跨团队契约面（共用程序表达、复现包、回归基线）	协作效率，依赖 P0 / P1 稳定
P2	性能归因工具链（roofline / dataflow / IR 分析）	性能调优的杠杆，属于精细化阶段

阶段化路线图（具体阶段、交付物、时间、负责人）由 Infra 月度路线图 承接，本文不展开。

支撑机制

上述要求通过下列产物落地，具体实现可随框架演进。每项机制的内容口径见对应章节。

机制	内容要点	服务对象	详见
配置快照	生效配置、覆盖来源、环境变量、关键运行参数	配置可组合、可校验、可追溯	设计原则
任务上下文	任务 ID、代码版本、模型、数据、硬件、并行策略、输出目录	协作契约双方共用同一组材料	协作契约
Kernel 注册与契约	版本、shape/dtype/layout/sharding、fallback/reference、benchmark、误差范围、profiling 标识	层内 critical path 接入	性能要求 § 层内 Pallas Kernel
层间调度策略	prefetch、wait placement、通信重叠、重算、释放、host offload 决策	层间编排可追踪	性能要求 § 层间调度
checkpoint 与恢复	训练状态、优化器状态、数据进度、框架元数据	故障恢复作为核心路径	诊断与验证 § 异常诊断
指标与日志	结构化日志、训练/系统/数值指标、异常事件	观测与诊断基线	诊断与验证 § 持续观测
profile 与 trace	step、通信、内存、编译、I/O、运行时调度	性能归因、Hang 定位	诊断与验证 § 持续观测、§ 异常诊断
IR 与编译产物	JAXPR、StableHLO、HLO、编译选项、调试材料	编译可验证、性能归因	性能要求、诊断与验证 § 主动验证
图级分析材料	roofline、dataflow、数学逻辑校验、误差预算、跨版本对比	性能与正确性统一输入	性能要求
最小复现材料	性能/数值/编译回归的可重放输入	回归分析	诊断与验证 § 主动验证
硬件诊断包	rank/device/slice、拓扑、心跳、collective 进度、runtime 日志、trace/profile、最后成功 op/kernel	NaN/Hang 定位	诊断与验证 § 异常诊断
验证模式	synthetic data、小规模、固定 seed、数值检查、数学逻辑校验、误差范围分析、故障注入、恢复演练	主动验证与故障演练	诊断与验证 § 主动验证、§ 异常诊断

验收口径

关键训练任务应能够通过统一训练框架完成端到端启动、运行、恢复、观测和验证。框架应支持在不同模型规模、并行策略和硬件环境下复用同一套任务抽象，并能为训练指标体系、正确性机制、性能评估体系、Pallas kernel 性能分析和编译器诊断工具链提供稳定的数据与控制接口。

当训练任务出现性能退化、数值异常、编译失败、通信瓶颈、内存问题、host offload stall、NaN、Hang、硬件异常或恢复失败时，训练框架应能够提供足够的上下文，使训练框架团队、kernel 实现、编译器团队和运行时/硬件支持人员基于同一组材料进行分析、复现和验收。