RATrain：面向带宽受限异构超算的稠密大模型训练状态生命周期调度

笔记日期 / Review date: 2026-07-09 笔记作者 / Review author: Zhongzhu Zhou 论文标题 / Paper reviewed: LifeTrain: Training-State Lifecycle Scheduling for Large Language Model Training on Bandwidth-Constrained Heterogeneous Supercomputers（论文里系统的名字叫 RATrain） 作者 / Paper authors: Yao Lu, Shiqing Ma, Zhongzhi Luan, Gen Li, Jiaxing Qi, Bin Han, Hailong Yang, Depei Qian arXiv: 2606.10415 状态 / Venue: 预印本，2026 年 6 月 9 日提交（北京航空航天大学中法工程师学院）

一句话总结

今天生产环境里几乎所有大模型训练系统——Megatron-LM、DeepSpeed/ZeRO、Alpa、GSPMD——默认的前提都是「你有一块 GPU」：快速的 HBM 显存、成熟的 NCCL 式集合通信库，以及跨设备的充裕带宽。RATrain 要问的是：如果这些前提统统不成立会怎样。它的目标平台是 MT-3000，一款真实存在的国产异构超算加速器，单个计算集群的原始算力相当可观，但每集群可用训练显存只有 20GB，没有现成的 GPU 式集合通信原语，而且集群间互联带宽窄得吓人（实测约 3.7GB/s）——窄到直接搬用张量并行或 ZeRO-3 反而会拖后腿。这篇论文最核心的一步棋，是不再把梯度同步、优化器更新、参数视图物化、激活值恢复当成「step 末尾的一堆杂活」笼统处理，而是把它们每一个都当成带有明确、逐层生命周期的运行时对象，调度进标准 1F1B 流水线本来就存在的那些小空档里。在这套调度机制之上再叠加一个针对该加速器定制的 GEMM/注意力反向传播执行后端，以及一个资源可行性规划器，RATrain 就能在 20GB 显存上限下训练从 LLaMA-2-7B 到 70B 的稠密模型，在同一硬件上比它们尝试过的每一种 GPU 式基线策略快 1.04-1.37 倍，并且扩展到 1024 个计算集群时仍能维持 97% 的扩展效率——同时训练出的 loss 曲线与语义等价的基线几乎完全重合（最大相对偏差仅 0.081%）。

前置知识

这是一篇关于稠密 Transformer 训练如何在硬件上执行的系统论文，不涉及模型结构或优化理论本身。要读懂后面的设计部分，你需要四块背景知识：为什么大模型训练一开始就要拆到多设备上、具体的 1F1B 流水线调度是怎么运作的、ZeRO 式状态切分做了什么，以及「异构超算」到底和 GPU 集群有什么本质区别。下面依次搭建这四块内容，再进入 RATrain 本身。

为什么要把模型切到多设备上

一个稠密的 decoder-only 大模型（想想 LLaMA-2 或 Qwen2.5）本质上是一堆几乎相同的 Transformer block 堆叠起来。训练它意味着为每一个参数存储：(a) 参数本身，(b) 它的梯度，(c) 优化器状态（对 Adam 来说是一阶矩和二阶矩估计——通常又是参数量的 2-3 倍）。对一个 700 亿参数的模型，在混合精度下这一项就接近一个 TB 量级，还没算上激活值。没有任何单个加速器有这么大的高速显存，所以训练必须至少沿以下三个轴之一切分：

• 数据并行（DP）：每个设备都持有完整的模型副本，处理批次里不同的切片；反向传播后梯度在设备间做平均（all-reduce）。
• 张量并行（TP）：把单层内部的矩阵乘法切到多个设备上（比如切线性层的隐藏维度），这要求在每一层的前向和反向传播内部做一次 all-reduce 或 all-gather 来拼回完整结果。
• 流水线并行（PP）：不同设备各自持有连续的一段层（称为「stage」）。一个 micro-batch 的数据先流经 stage 0，再流经 stage 1，以此类推完成前向传播；反向传播则按相反顺序流回去。

这三者是可以组合的——Megatron-LM 的经典配方是节点内用 TP（互联快），节点间用 PP/DP（互联相对慢一些但还够用）。RATrain 最核心的经验性发现，也是后面「结果」部分会反复印证的一点，恰恰是：这套经典配方在 MT-3000 上是完全错的方向，因为 TP 的层内集合通信是三者中最耗带宽的原语，而集群间带宽恰恰是 MT-3000 平台最稀缺的资源。

1F1B 流水线调度，配一个具体例子

流水线并行有一个显而易见的朴素实现：先把所有 micro-batch 的前向传播都跑完，再把所有 micro-batch 的反向传播都跑完（GPipe 那种风格）。问题在于显存：stage 0 必须把每个 micro-batch 的激活值一直存着，从算出来那一刻到对应的反向传播回来为止——在 GPipe 调度下，这个时间跨度会非常长。

1F1B（“one-forward-one-backward”） 把不同 micro-batch 的前向和反向传播交错执行，让某个 stage 在任意时刻持有的在飞激活值数量不会超过必要限度。具体来说，$P$ 个流水线 stage、每个 step $M$ 个 micro-batch 的情况下，某个中间 stage 的 1F1B 调度大致是：

预热阶段（warm-up）:   F1  F2  F3  ...  F_{P-1}
稳态阶段（steady）:    F_P  B1  F_{P+1}  B2  F_{P+2}  B3  ...
收尾阶段（cool-down）: B_{M-P+1} ... B_M

预热阶段中，stage $p$ 必须先完成 $P-p$ 次前向传播，才能开始它的第一次反向传播（因为第一个反向传播要从最后一个 stage 一路传回来）。稳态阶段中，前向和反向一对一交替执行，这也是「1F1B」这个名字的来源。收尾阶段只剩反向传播，把流水线排空。论文的 Figure 6（下文重绘为 Figure 3）正好复现了这种模式：$P=4$、8 个 micro-batch 的情况下，输入侧 stage S0 还在做 F5、F6、F7 的前向传播时，输出侧 stage S3 已经在做反向传播 B1 了——这种错位正是 1F1B 相对 GPipe 的显存优势所在，因为 S0 任何时刻持有的在飞激活值数量最多是 $P$ 个，而不是 $M$ 个。

stage $p$ 需要持有激活值的在飞 micro-batch 数量大致是 $N_{act}(p) = P - p$（输入侧 stage 比输出侧 stage 持有得更多）——这一个事实就是 RATrain 三大核心机制之一（激活值恢复）的种子，下文会详细展开。

ZeRO：切分优化器状态、梯度与参数

ZeRO（Zero Redundancy Optimizer）针对的是另一种显存浪费：在朴素数据并行下，$D$ 个数据并行副本中的每一个都冗余地存储着整份优化器状态、梯度缓冲区和参数集合。ZeRO 有三个逐级更激进的阶段：

• ZeRO-1：只在 $D$ 个副本间切分优化器状态（比如 Adam 的动量/方差）。每个副本本地算出完整梯度，但只更新 $1/D$ 的参数，再把更新后的分片广播出去。
• ZeRO-2：进一步切分梯度。每个副本只需要物化自己负责更新的那部分梯度分片；其余部分在反向传播过程中可以直接 reduce 后丢弃（用 reduce-scatter 代替 all-reduce）。
• ZeRO-3：进一步切分参数本身。这时任何副本都不再在静态时持有完整模型；参数必须在每层前向/反向计算前按需 gather 回来（一次 all-gather），用完后再释放。

这个取舍是单调的：每多一级 ZeRO 都能省更多显存，但也暴露更多通信，因为 ZeRO-3 尤其需要每一层的每一次前向和反向传播都重新做一次参数分片的 all-gather——论文把这称为「参数视图重建」。在有快速 NVLink/InfiniBand 的 GPU 集群上，这通常是笔好交易。但在 MT-3000 上，正如后面会看到的，往往不是——这正是为什么 RATrain 的规划器在绝大多数生产规模配置下最终都选了 ZeRO-2 而不是 ZeRO-3 的原因。

激活值检查点，以及「恢复正好卡在关键路径上」这个问题

前向传播在每一层都会产生中间激活值，反向传播需要用它们来算梯度。把它们全部留在显存里（“full-save”）代价很高，尤其是对持有很多在飞 micro-batch 的输入侧流水线 stage 而言（回想上面的 $N_{act}(p) = P-p$）。经典解法激活值检查点（Chen et al., 2016）只保留少数「checkpoint」激活值（比如每层边界处一个），需要的时候从最近的 checkpoint 重新计算出其余部分。

问题在于：这次重计算发生在反向传播的那一刻，直接插在反向关键路径上。如果重算一层的中间激活值要花 5ms，而这恰好发生在该层反向矩阵乘法能开始之前，那这 5ms 就是纯粹加到每一个 micro-batch 反向传播上的延迟——除非能找到别的有用工作来重叠，否则这段延迟藏不掉。这正是 RATrain 的「前向侧激活值恢复」（FSR）要解决的问题：不在反向传播那一刻重算，而是提前在 1F1B 调度本来就有的空档里重算。

「异构超算」到底和 GPU 集群有什么本质区别

MT-3000 是一款真实存在、用于中国 exascale 超算平台的加速器（Lu et al., CCF THPC 2022）。从结构上看，它和 GPU 完全不是一回事：

• 基本调度单元是一个计算集群（acceleration cluster），每个集群内含 24 个 DSP（数字信号处理器，不是 CUDA core）。
• 每个 DSP 有一套显式的、软件管理的两级片上存储层次：64KB 的标量存储（SM）和 768KB 的阵列存储（AM）——没有自动缓存，程序/运行时必须显式地把数据搬进搬出。
• 同一集群内的 24 个 DSP 共享一块全局共享存储（GSM），GSM 再通过 DMA 与片外DDR交互。
• MT-3000 平台上多个集群之间，仅通过一个CPU/GP Zone 和一个内存桥连接——没有类似 NVLink 那种专用高带宽互联结构。

作者直接测得的结果是：每个集群能提供约 8.1 TFLOPS 的 FP16 算力——单看这个数字完全能和 GPU 打平——但每集群可用训练显存只有 20GB，有效 DDR 带宽约 30GB/s，集群间点对点带宽仅约 3.7GB/s。最后这个数字才是全文的关键：3.7GB/s 比 NVLink 级别的 GPU 互联带宽低了大约两个数量级。任何关键路径依赖频繁、大量集合通信的训练策略（张量并行的层内 all-reduce、ZeRO-3 的逐层参数 gather）都会被这一个数字死死卡住。

用一个简化的 Roofline 视角看这三个数字

把这三个数字——8.1 TFLOPS 算力、30GB/s DDR 带宽、3.7GB/s 集群间带宽——放到同一个「roofline」心智模型里会很有帮助，因为它们定义了三种截然不同的瓶颈区间，而 RATrain 整套设计几乎都可以读成「先搞清楚每个操作到底落在哪个区间，再针对那个区间专门优化，而不是假设所有操作都吃同一种瓶颈」。

算术强度指的是每搬运一字节能做多少次浮点运算。相对 DDR，一个操作要落在算力瓶颈区间（被 8.1 TFLOPS 卡住），算术强度必须超过约 $8.1e12 / 30e9 \approx 270$ FLOPs/byte；相对集群间链路，这个门槛更是高达约 $8.1e12/3.7e9 \approx 2189$ FLOPs/byte。像 Table 5 里那种大而「方」的 GEMM（比如 $8192 \times 8192$）算术强度高，只要操作数能通过 SM/AM/GSM 妥善分级搬运、避免反复回 DDR，就有机会接近算力瓶颈——这正是 Figure 3（论文 Fig.4）里那段流水线级指令交叠在争取的东西。相反，张量并行的层内 all-reduce 要跨集群间链路搬运激活张量——整块芯片上字节/FLOP 门槛最苛刻的区间——却换不来多少算力收益，这正是下文「三处结构性错配」中第一处代价最高的原因。这个 roofline 框架是我自己加上去帮助横向比较这三个数字的；论文本身是分别给出这三个数字（第 2.1 节），并没有画在同一张图上，但定性结论——集群间带宽远比另外两者更紧张——正是论文所有设计选择（避开 TP、避开 ZeRO-3、偏爱 PP）一致在优化的方向。

三处结构性错配：为什么不能直接搬 GPU 训练栈过来

论文第 2.3 节具体列出了直接复用 GPU 式训练策略在 MT-3000 上为什么会失败，值得逐条过一遍，因为每一条错配恰好对应 RATrain 后面三大核心机制中的一个。

错配一——层内集合通信。 张量并行在每一层的前向和反向传播内部插入一次 all-reduce（或 all-gather），处在关键路径上。在有 NVLink 的 GPU 节点上，这次通信足够快，能和计算重叠。在 MT-3000 上，3.7GB/s 的集群间带宽下，同样这次 all-reduce 会变成一次巨大、暴露在外的停顿，而且每层都要发生一次。

错配二——激活值驻留与恢复不均衡。 如前所述，1F1B 下输入侧流水线 stage 比输出侧持有更多在飞 micro-batch 激活值（$N_{act}(p) = P - p$）。Full-save 用高峰值显存换取这一点（20GB 上限未必能吸收得住）；经典检查点用反向时重算换取这一点（受限的 DDR 带宽让这变得很贵）。

错配三——step 末尾的收尾长尾。 如果梯度同步（GradSync）、优化器更新（UpdateShard）、参数视图准备（PrefetchW）都被推迟到梯度累积边界统一处理——大多数传统数据并行/ZeRO 实现都是这么干的，把它们当成一个整体的「step 末尾」阶段——那么这些工作就会变成附加在每个训练 step 末尾的一段串行、藏不掉的长尾。在带宽充裕的 GPU 集群上这条尾巴很短。在 MT-3000 上它可能占主导。

下图是这三处错配及 RATrain 对应的回应机制，对照的是论文原文的 Figure 2。

flowchart LR
    subgraph M1["错配一：TP-heavy 执行"]
        A1["每层插入一次\n层内 all-reduce"] --> A2["暴露在\n3.7GB/s 集群间链路上"]
    end
    subgraph M2["错配二：激活值驻留"]
        B1["输入侧 stage 持有\n更多在飞激活值"] --> B2["Full-save：显存压力\nCheckpoint：反向时重算"]
    end
    subgraph M3["错配三：step 末尾长尾"]
        C1["GradSync / UpdateShard / PrefetchW\n延迟到累积边界统一处理"] --> C2["串行、藏不掉的\n收尾长尾"]
    end
    A2 --> R1["RATrain 的回应：\nPP + DP + 轻量 ZeRO\n（避开层内集合通信）"]
    B2 --> R2["RATrain 的回应：\n前向侧激活值恢复（FSR）"]
    C2 --> R3["RATrain 的回应：\n逐层状态流水线 +\n更新-预取调度"]

Figure 1（对应论文 Fig.2，重绘）：GPU 式训练策略与 MT-3000 带宽/显存特性之间的三处结构性错配，以及 RATrain 对每一处的回应。

承上启下：训练状态生命周期是贯穿全文的统一思路

在进入架构之前，值得先点名把 RATrain 三大机制串起来的那一个概念——它会以略微不同的形式反复出现三次。稠密 Transformer 的前向传播按顺序访问层 $1, 2, \ldots, L$；反向传播按相反顺序 $L, L-1, \ldots, 1$ 访问它们。这不是无关紧要的细节——它意味着反向传播中算出来的每一个量（比方说梯度）都会以一种固定、可预测、按层索引的顺序「准备就绪」，而下一轮前向传播要消费的每一个量（比方说更新后的参数视图）也会以固定、可预测、按层索引的顺序被「需要」。因为产出顺序和消费顺序都提前已知且都绑定在层索引上，两者之间就必然存在一个窗口——从某个东西准备好到它下一次被需要之间的这段时间——而这正是 RATrain 用来调度那些本来会卡在关键路径上或堆积在 step 边界处的工作的地方。这就是论文所说的「训练状态生命周期调度」：模型的梯度、更新后的参数、可恢复的激活值，不再是不透明的 step 末尾杂物,而是带有出生时间、截止时间以及两者之间可调度窗口的运行时对象。

架构总览

RATrain 的系统主线分三个阶段，如下图所示（对应论文原文 Figure 3）。

flowchart TB
    P1["阶段一输入：模型画像\n（层数/隐藏维度/序列长度）"] --> PL["资源感知\n配置规划器"]
    P2["阶段一输入：平台画像\n（显存/带宽/拓扑）"] --> PL
    P3["阶段一输入：执行画像\n（前向/反向/更新开销）"] --> PL
    PL --> PLAN["可执行方案：\nPP/DP/ZeRO 度数、\nmicro-batch、激活策略、\n预取策略"]
    PLAN --> SL["阶段二：不变的\n非交错式 1F1B 主执行路径"]
    SL --> T1["(a) 逐层状态流水线\nGradSync 再 UpdateShard 再 PrefetchW"]
    SL --> T2["(b) 更新-预取调度\n队列管理、截止时间感知"]
    SL --> T3["(c) 前向侧激活值\n恢复（FSR）"]
    T1 --> BE["阶段三：显式存储层次\nDDR、GSM、AM、SM"]
    T2 --> BE
    T3 --> BE
    BE --> OP["FP16 GEMM 汇编流水线\n与显存驻留式注意力反向传播"]

Figure 2（对应论文 Fig.3，重绘）：RATrain 三阶段主线。规划器用离线画像选出一个资源可行的方案；stage 本地运行时围绕不变的 1F1B 主路径调度训练状态生命周期任务；后端在 MT-3000 显式存储层次上执行具体算子。

在深入细节之前，这张架构图有两点值得先强调。第一，1F1B 主执行顺序从头到尾没有变过——前向传播依然从输入侧走到输出侧，反向传播依然按相反顺序回来，梯度累积语义和优化器更新公式都没被动过。RATrain 只改变周边杂活何时、在哪里被物化执行。正是这个设计选择,让作者后面能拿一个语义等价的基线做接近逐位对比的正确性验证（见下文「正确性验证」部分），并且得到几乎为零的 loss 偏差——因为实际计算根本没变,变的只是它的调度方式。第二,三个阶段职责划分很干净：规划器回答「在跑任何东西之前，哪种配置本身是可行且快的」，stage 本地运行时回答「给定一个可行配置，运行时按什么顺序发出任务」，后端回答「具体一个 GEMM 或注意力反向 tile 在这块硬件上怎么执行才高效」。下面按论文的顺序依次深入四大核心机制（后端、状态流水线+更新-预取、FSR、规划器）。

深入之一：MT-3000 感知的执行后端

上面提到的一切——规划器的开销估计、运行时的调度决策——都依赖于拿到稳定、可预测的算子级延迟才能推理。如果一次 GEMM 的延迟每次跑都天差地别，任何静态调度决策都无法信任。所以 RATrain 首先要构建一个后端，让稠密 Transformer 训练里最贵的两个操作——GEMM 和注意力反向传播——在 MT-3000 的显式存储层次上做到又快又可预测。

为什么现成的 GEMM 库不能直接用

论文指出，现有面向多核 DSP 的 GEMM 优化「主要面向 FP32 通用 GEMM 或单算子场景」，无法直接为大模型训练关键路径提供所需的 FP16 GEMM、注意力反向传播和显式数据搬运支持。这是国产加速器系统类论文里反复出现的一个主题：通用厂商库明显落后于某个特定、高价值负载（稠密 Transformer 训练）真正需要的东西，于是系统团队最终不得不自己手写热点算子。

FP16 GEMM 数据流

论文 Figure 4（下图重绘）展示了 RATrain 如何把一次 $C=AB$ 的 GEMM 组织到整个存储层次上。设分块为 $A_g[M_g, K_g]$（从 DDR 分级到 GSM 的粗粒度块）、$A_2[M_2, K_2]$（进一步分级到 SM 的细粒度块），以及广播/累积到 AM 中的 $B_2[K_2, N_2]$、$C_2[M_2, N_2]$。

flowchart LR
    DDR["DDR（片外）"] -- "DMA 加载 Ag" --> GSM["GSM（集群共享）"]
    GSM -- "分级到 A2" --> SM["SM（每 DSP 64KB）"]
    DDR -- "广播 B2" --> AM["AM（每 DSP 768KB）"]
    DDR -- "加载 C2 作为累加器" --> AM
    SM --> VMAC["VMAC 微内核\nFP16 乘加并累积"]
    AM --> VMAC
    VMAC -- "写回 C2" --> DDR

Figure 3（对应论文 Fig.4，重绘）：FP16 GEMM 数据流。左操作数经 GSM/SM 流式搬运，右操作数和累加器住在 AM 里，VMAC 微内核负责真正的乘加运算。

性能关键的细节在 DSP 本地的汇编流水线里：RATrain 把地址生成、加载、半精度提取、广播和 FP16 乘加这几路指令流交叠起来，让下一块 tile 的操作数（$A_{next}$、$B_{next}$）准备工作和当前 tile 的乘加计算重叠。论文的 Table 1 把这展开成一个 8 路 VLIW 调度，下面重绘并附上每个功能单元的角色：

功能单元	在 GEMM 微内核里的角色
VMAC	向量乘加——真正的 FP16 乘加运算（vfmulas32）
SMAC1 / SMAC2	标量广播/地址辅助单元，负责给 VMAC 送操作数（smvaga、svbcast）
SLDST	标量加载/存储——在当前 tile 还在执行时就发出下一 tile 的半字加载（sldh A_next）
VLDST1 / VLDST2	向量加载/存储——提前从 DDR 批量取下一 tile 的 B 操作数（vldw B_next）
SIEU	标量整数/执行单元——负责排序和地址记账（seq）
SBR	标量分支——微内核的循环控制（sbr）

这八路功能单元全部在同一条 VLIW 指令束里同时发出——这正是「准备下一块」和「计算当前块」能真正同时发生、而不是先后发生的关键：VMAC 还在消费当前的 $A$/$B$ tile 时，SLDST 和 VLDST1/VLDST2 已经在预取下一 tile 的数据了，所以 VMAC 算完当前 tile 时不需要等它的后继数据。没有这种重叠，每次 tile 切换都会让乘加流水线停下来等下一个操作数——这正是论文想要消灭的那种「显存访问气泡」。这在概念上等价于 CUDA kernel 里的双缓冲，只不过是显式做出来的，因为 MT-3000 没有自动预取或缓存层次来隐式完成这件事——把这段延迟藏起来的责任落在编译器/运行时身上，而不是硬件。

显存驻留式注意力反向传播：逐行拆解算法 1

注意力反向传播比一次普通 GEMM 更棘手,因为它涉及好几个相互依赖的中间张量——$P_{ij}$（注意力概率）、$P_{ij}^T$（其转置）、$GP_{ij}$（概率的梯度）、$GS_{ij}$（分数/softmax 前的梯度），以及 $GS_{ij}^T$——通过一条链路计算：$dV, dP, dS, dQ, dK$。如果把这条链上的每个中间张量都当作独立的 kernel 调用去读写 DDR，反向路径就要承担大量的片外流量。RATrain 的解法是把整条链尽量长时间地驻留在 AM/SM 内部，采用一种外层遍历 query、内层遍历 key/value的循环结构——这在思路上类似 FlashAttention 的 tiling，但重新针对的是 MT-3000 特有的瓶颈（SM 受限的左操作数分级、基于 GSM 的跨 DSP 归约），而不是 GPU 的 HBM 带宽。

下面是算法本体（论文 Algorithm 1），逐行附上解读：

算法 1：显存驻留式注意力反向传播 Tile 调度
输入：query tile Q_i, 输出梯度 GO_i, 前向保存的概率 tile {P_ij},
      key/value tile {K_j, V_j}
输出：query 梯度 GQ_i, key/value 梯度 {GV_j, GK_j}

 1: 外层驻留初始化：(Q_i, GO_i) <- LoadAM(Q_i, GO_i)
     # 把这个外层 query tile 及其输出梯度一次性加载进 AM；
     # 后面每个内层 key/value tile 都会复用它们。
 2: 为 GQ_i 分配 AM 缓冲并清零
     # GQ_i 要累加来自每个内层 j 的贡献，必须从 0 开始。
 3: for 每个 key/value tile j do
 4:     内层广播：(K_j, V_j) <- BcastAM(K_j, V_j)
         # 把这个内层 tile 的 key/value 送到每个 DSP 的 AM 里。
 5:     加载前向状态：P_ij <- LoadAM(P_ij)
         # 直接复用前向传播保存下来的 softmax 概率，
         # 而不是在反向时重新算一遍 softmax。
 6:     AM 内驻留计算：GP_ij = GO_i * V_j^T,
                        GS_ij = SoftmaxBackward(P_ij, GP_ij)
         # 标准的注意力反向链式法则，全程在 AM 内完成。
 7:     为 GV_j 做 SM 分级：P~_ij^T <- StageSM(P_ij^T)
         # 把转置后的概率 tile 搬进 SM，因为接下来的 GEMM
         # 需要它作为「左操作数」（按 RATrain 的 GEMM 约定，
         # 左操作数就住在 SM 里，见上面 Figure 3）。
 8:     GV_j^part <- P~_ij^T * GO_i
         # 这个 key/value tile 对 V 梯度的一部分贡献。
 9:     GSM 归约：GV_j <- ReduceAddGSM(GV_j, GV_j^part)
         # 不同 DSP 用不同的 query tile 算出各自的部分和，
         # 通过 GSM（而不是 DDR）在集群内归约。
10:     为 GQ_i 做 SM 分级：GS~_ij <- StageSM(GS_ij)
11:     GQ_i <- GQ_i + GS~_ij * K_j
         # 把这个内层 tile 的贡献累加到从第 1-2 步开始
         # 就驻留在 AM 里的 query 梯度上。
12:     为 GK_j 做 SM 分级：GS~_ij^T <- StageSM(GS_ij^T)
13:     GK_j^part <- GS~_ij^T * Q_i
14:     GSM 归约：GK_j <- ReduceAddGSM(GK_j, GK_j^part)
15: end for
16: 写回：WriteBack(GQ_i)
     # 每次外层迭代只写回 GQ_i；GV_j 和 GK_j 要等它们
     # 各自按 j 索引的归约循环（在别处以外层循环运行）完成后才写回。

有两个设计选择值得单独点出来：

1. 为什么外层遍历 query、内层遍历 key/value（而不是反过来）？ 因为 $Q_i$ 和 $GO_i$（query 和输出梯度）在每一个内层 $j$ 都要用到，而 $K_j/V_j$ 只是每次内层迭代里临时需要一下。把外层循环用到的张量驻留在 AM 里，避免了对 $Q_i$、$GO_i$ 做 $J$ 次重复的 DDR 加载（每个内层 tile 一次）；它只在每次外层迭代时加载一次。这和上面 Figure 3 里 GEMM 数据流「常驻复用的操作数，流式过一次性的操作数」是同一个原则。
2. 为什么用 GSM 而不是 DDR 归约？ 因为 $GV_j$ 和 $GK_j$ 各自都会收到同一集群内、处理不同 query tile $i$ 的每个 DSP 的部分贡献。朴素实现会让每个 DSP 把自己的部分和写到 DDR 上再统一求和——但 DDR 带宽是这个平台第二稀缺的资源（仅次于集群间带宽）。GSM 是片上的、24 个 DSP 共享，在这里归约完全避开了那一次片外往返。

Tile 调度背后的容量约束

Tile 的大小 $B_r$（query tile 行数）、$B_c$（key/value tile 列数）、$d$（head 维度），以及分级子块 $B_c'$、$B_r'$，都不能随便选——必须落在 AM、SM、GSM 各自固定的物理容量 $C_{AM}$、$C_{SM}$、$C_{GSM}$ 之内。设 $D_{size}$ 为一个 FP16 元素的字节数，论文的公式 1 把这写成四条同时满足的不等式：

$$D_{size}(B_r B_c + 2 B_r d + 2 B_c d) \le C_{AM}$$ $$D_{size} B_c' B_r \le C_{SM}$$ $$D_{size} B_r' B_c \le C_{SM}$$ $$D_{size} B_c d \le C_{GSM} \tag{1}$$

从左到右读：第一条约束的是 AM 的工作集：概率/分数 tile（$B_r B_c$）加上 query 及其梯度缓冲（各是 $B_r \times d$ 形状，且需要一份「伴随」副本，故系数为 2），再加上 key/value 侧对应的部分（$B_c \times d$，同样系数 2）。第二、三条约束的是转置后的概率/分数 tile 能同时经 SM 分级多少——回想前面算法逐行解读里，第 7、10、12 步都是在对应 GEMM 之前把转置 tile 搬进 SM。第四条约束的是 ReduceAddGSM 调用（第 9、14 步）期间跨 DSP 归约块要占用的 GSM 大小。

论文特别强调的一点是：这套调度不是 FlashAttention。FlashAttention 的 tiling 是为了给一台拥有大而快、但仍受带宽限制的统一显存（GPU HBM）的设备减少往返次数。RATrain 的调度针对的是完全不同的约束面——SM 受限的左操作数分级、基于 GSM 的局部（而非全局）梯度归约——因为 MT-3000 没有类似 HBM 的统一可寻址显存；SM、AM、GSM、DDR 是四块物理和语义上都独立的存储，软件必须显式在它们之间搬运数据。这是全文对「直接搬 GPU kernel 过来行不通」这一论点最具体的一个例证：不只是数值不一样，优化问题本身的形状就不一样。

深入之二：逐层状态流水线与更新-预取调度

这个机制直接针对前面动机图里的「错配三」（step 末尾收尾长尾）。

从批量式 step 末尾处理到逐层任务链

在传统的数据并行或 ZeRO 实现里，梯度同步（GradSync）、优化器参数更新（UpdateShard）、下一轮参数视图准备（PrefetchW）通常都被推迟到梯度累积边界，当作所有 micro-batch 反向传播都跑完之后的一个批量阶段统一处理。RATrain 反过来利用了前置知识部分建立的事实——反向传播按固定顺序 $L, L-1, \ldots, 1$ 访问各层，这意味着层 $l$ 的梯度会在累积窗口内某个特定、可预测的时间点完全累积完毕，远早于层 1 的反向传播真正结束。RATrain 把这个批量阶段拆解成逐层任务链：

$$\text{GradSync}(l) \to \text{UpdateShard}(l) \to \text{PrefetchW}(l) \tag{2}$$

层 $l$ 的这条任务链在层 $l$ 本地梯度累积完成的那一刻就变得可调度——不是在单个 micro-batch 的 Backward(l) 之后立刻触发，也不是推迟到 step 最末尾。论文的 Figure 5（下图重绘）展示了三个相邻层的这套调度。

flowchart LR
    subgraph L1["层 l+1"]
        BW1["Backward(l+1)"] --> GS1["GradSync(l+1)"]
        GS1 --> US1["UpdateShard(l+1)"]
        US1 --> PF1["PrefetchW(l+1)"]
        PF1 --> FN1["Forward_next(l+1)"]
    end
    subgraph L2["层 l"]
        BW2["Backward(l)"] --> GS2["GradSync(l)"]
        GS2 --> US2["UpdateShard(l)"]
        US2 --> PF2["PrefetchW(l)"]
        PF2 --> FN2["Forward_next(l)"]
    end
    subgraph L3["层 l-1"]
        BW3["Backward(l-1)"] --> GS3["GradSync(l-1)"]
        GS3 --> US3["UpdateShard(l-1)"]
        US3 --> PF3["PrefetchW(l-1)"]
        PF3 --> FN3["Forward_next(l-1)"]
    end

Figure 4（对应论文 Fig.5，重绘）：逐层状态流水线。每一层的 GradSync -> UpdateShard -> PrefetchW 任务链一旦该层的本地依赖满足就能独立调度，而不必等所有层都完成反向传播。

GradSync(l) 被安排去和之后的反向计算（层 $l-1, l-2, \ldots$）或 stage 本地空闲时间重叠——绝不是紧跟在 Backward(l) 完成后立刻串行执行，因为那样做会让反向传播停下来等一次其实有充足时间「在后台」完成的同步，只要后面的层还在计算就行。这是个微妙但重要的点：论文并不是说同步变免费了，只是说只要运行时愿意按层粒度而不是 step 粒度去推理，同步就可以藏在反正都会发生的剩余反向工作后面。

把更新-预取建模成一个截止时间调度问题

任务链后半段 UpdateShard(l) -> PrefetchW(l) 被显式建模成一个实时系统里常见的截止时间调度问题——这是从实时系统理论向机器学习系统领域一次挺漂亮的跨界借用。设 $t_{sync}(l)$ 为 GradSync(l) 的完成时间（更新-预取链最早能开始的时刻），$t_{use}(l)$ 为下一次 Forward(l) 需要更新后的参数视图 $W_{view}(l)$ 的时刻（这条链必须完成的截止时间）。可调度窗口是：

$$t_{sync}(l) \le t < t_{use}(l) \tag{3}$$

如果 PrefetchW(l) 在这个窗口内的任何时刻完成，下一次前向传播就能拿到一份「热的」、已经物化好的参数视图，不需要额外付出任何延迟。如果它在 $t_{use}(l)$ 之后才完成，前向传播就要停下来等它——这段没被覆盖的部分正是论文所说的「下一次前向停顿」。设 $T_{upd}(l)$/$T_{pref}(l)$ 为更新/预取任务的实际延迟，$W_{upd}(l)$/$W_{pref}(l)$ 为各自能真正抢到的窗口大小（要和其他流量竞争）,那么暴露出来的延迟——藏不掉、直接加到 step 时间上的那部分——是：

$$E_{upd}(l) = \max\bigl(0,\, T_{upd}(l) - W_{upd}(l)\bigr)$$ $$E_{pref}(l) = \max\bigl(0,\, T_{pref}(l) - W_{pref}(l)\bigr) \tag{4}$$

这个 $\max(0, \cdot)$ 的形式值得多花点时间体会,因为同样的模式后面还会再出现两次（一次用在激活值恢复上,一次出现在规划器的开销模型里）：它说的是「如果任务能塞进它可用的窗口里,就完全不额外花钱；只有超出窗口的那一部分才会以可见的延迟形式出现」。这是「把工作藏在空闲时间后面」这个想法的数学形式化,也正是让规划器（深入之四）能用一套统一的记账方式去推理好几种截然不同调度机制的关键。

实际效果，在后面的消融实验里能看到确凿数字：单独关掉更新-预取调度（但保留 GradSync 的逐层状态流水线）会让暴露的收尾长尾放大 2.31 倍；把两个机制都关掉（退回批量式 step 末尾处理）会让长尾放大 4.59 倍。逐层拆解和截止时间感知的预取分别都有用,而两者组合起来比单独任何一个都更好——这一点论文的消融实验（下文详述）是用硬数字而不是空口白话证明的。

深入之三：前向侧激活值恢复（FSR）

这个机制直接针对动机图里的「错配二」（激活值驻留不均衡），也是 RATrain 三个机制里单独贡献最大的一个（消融实验显示单独关掉 FSR，step 时间会涨到 1.33 倍——三者里单一机制影响最大的一个）。

核心想法：在反向传播到达之前恢复,而不是到达时才恢复

回想前置知识部分,经典激活值检查点是在反向传播到达需要它的 stage 那一刻重新计算缺失的中间激活值——正好卡在反向关键路径上。FSR 的洞察一旦看懂了前面「按层顺序」的论证会显得几乎有点朴素：运行时其实提前就知道，某个特定 micro-batch 的反向传播接下来会在某个 stage 到达，因为 1F1B 的调度是确定性的、提前已知的（回想前面预热/稳态/收尾的结构）。所以 FSR 不是等反向传播到达再重算，而是提前在 stage 本来就有的前向侧空档或空闲气泡里重算好，把结果存进一个短生命周期的恢复缓冲区,供随后到达的反向传播直接消费。

论文 Figure 6（下图重绘）用一个 $P=4$、8 个 micro-batch 的调度具体展示了这一点。面板 (a) 里，标准调度让 stage S0 先执行 F4、F5、F6、F7,再做第一次反向传播 B1（重算发生在这里,卡在关键路径上）。面板 (b) 里，FSR 插入了小的「R」（恢复）任务——F7 之后的 R1、F8 之后的 R2 等——这样到每次反向传播（B1、B2、……）真正开始时，它需要的激活值已经躺在恢复缓冲区里了。

flowchart LR
    subgraph Std["(a) 标准非交错式 1F1B"]
        direction LR
        SF4["F4"] --> SF5["F5"] --> SF6["F6"] --> SF7["F7"] --> SB1["B1\n（重算发生在这里，\n卡在关键路径上）"] --> SF8["F8"] --> SB2["B2\n（重算发生在这里）"]
    end
    subgraph FSRp["(b) FSR 增强调度"]
        direction LR
        FF4["F4"] --> FF5["F5"] --> FF6["F6"] --> FF7["F7+R1\n（提前恢复）"] --> FB1["B1\n（激活值已就绪，\n无需等待）"] --> FF8["F8+R2\n（提前恢复）"] --> FB2["B2\n（激活值已就绪）"]
    end

Figure 5（对应论文 Fig.6，重绘）：FSR 把激活值恢复从反向关键路径（面板 a，重算恰好在反向传播需要数据的那一刻发生）挪到调度里更早出现的前向侧或空闲空档（面板 b，标注为 R1-R4 的任务），这样反向传播找到的输入已经准备好了。

关键的是,FSR 完全没有改变 1F1B 的执行顺序——F4 依然在 B1 之前，B1 依然在 F8 之前，相对位置和以前一样。它改变的只是前向/空闲空档里恢复子任务被调度的时刻，以及恢复出来的激活值存在哪里（一个短生命周期的缓冲区，而不是反向传播那一刻才现算现用）。这和逐层状态流水线是同一套哲学：改调度，不改语义。

量化这笔显存账

设 $N_{act}(p)$ 为 1F1B 下 stage $p$ 必须驻留激活值的 micro-batch 数量（大致是 $P - p$，输入侧 stage 更多），$M_{full}$ 为一个 micro-batch 完整激活值的大小，$M_{ckpt}$ 为 checkpoint 本身的大小，$M_{rec}$ 为短生命周期恢复缓冲区的大小。在 full-save 策略下,stage $p$ 的激活值显存峰值简单地是：

$$M_{act,full}(p) \approx N_{act}(p) \cdot M_{full} \tag{5}$$

这个量随着 pipeline 位置相关的重数 $N_{act}(p)$ 和（很大的）完整激活值大小 $M_{full}$ 线性增长——正是这一项在输入侧 stage 上把显存冲破 20GB 上限,这也解释了为什么论文自己的实验里 full-save 在每一个测试配置上都触发了 OOM。

在 FSR 下,长生命周期的完整激活值被便宜的 checkpoint 加一个只需短暂存在的恢复缓冲区取代：

$$M_{act,FSR}(p) \approx N_{act}(p) \cdot M_{ckpt} + M_{rec} \tag{6}$$

由于 $M_{ckpt} \ll M_{full}$（checkpoint 按设计只是一层全部中间状态里很小的一部分——比如只是层边界处的张量，而不是层内部每一个中间量），第一项急剧缩小,而恢复缓冲区项 $M_{rec}$ 根本不随 $N_{act}(p)$ 增长——它是按设计短生命周期、可复用的，不会累积。这正是 FSR 把一个 $O(N_{act}(p) \cdot M_{full})$ 的显存足迹变成 $O(N_{act}(p) \cdot M_{ckpt} + M_{rec})$ 的具体机制——它在显存这个轴上做的事，大致相当于经典检查点已经在做的事，但没有像经典检查点在延迟这个轴上那样把重算暴露在反向关键路径上。

不过 FSR 并不是一个能把恢复成本完全变没的魔术——论文对此很直白。如果前向侧恢复窗口太短,或者本地计算/显存资源恰好不可用,RATrain 会退回到普通的反向时恢复（无论哪种方式，训练语义都保持不变），没被覆盖的那部分恢复延迟仍然会出现在 step 时间估计里，遵循我们已经见过的 $\max(0, \cdot)$ 模式：

$$E_{rec}(p) = \max\bigl(0,\, T_{rec}(p) - W_{rec}(p)\bigr) \tag{7}$$

其中 $T_{rec}(p)$ 是恢复延迟，$W_{rec}(p)$ 是 stage $p$ 可用的前向侧/空闲窗口大小。当窗口足够宽松地超过恢复成本时，FSR 完全免费；当窗口不够时，它优雅地退化成接近经典检查点的水平，而不是彻底失效。这种优雅退化的性质对规划器（深入之四）很重要，因为它意味着规划器对 FSR 的开销估计是反向时检查点成本的一个严格上界，永远不会更差。

深入之四：资源感知配置规划器

后端、状态流水线和 FSR 都是运行时机制——它们让一个给定的训练配置（特定的 PP 度数、DP 度数、ZeRO 阶段、micro-batch 大小等）跑得高效且不超预算。但在组合数巨大的选择空间里，给定一个模型大小和集群数，到底该选哪种配置？这就是规划器的工作，也是把全文串成一个真正能拿去运营、而不是手调一次就再也不管的系统的关键一环。

为什么固定的经验法则不能泛化

论文说得很直接：「固定的经验法则因此很难在不同模型规模和资源约束下稳健地套用」。一条为 LLaMA-2-7B 调好的经验法则（比如「永远用 ZeRO-2、$P=2$」）没有理由在 LLaMA-2-70B 上依然最优——论文自己的结果显示，规划器需要把流水线度数一路推到 $P=48$才能塞进 20GB 每集群的上限。RATrain 没有按模型规模手调,而是构建了一套有原则的搜索加开销模型的方法。

配置空间与可行性检查

一个候选训练配置被表示为一个元组：

$$c = (P, D, Z, b, A, \pi_{act}, \pi_{pref}) \tag{8}$$

其中 $P$ 是流水线度数，$D$ 是数据并行度数，$Z$ 是 ZeRO 阶段，$b$ 是本地 micro-batch 大小，$A$ 是梯度累积步数，$\pi_{act}$ 是激活值恢复策略（full-save / checkpoint / FSR），$\pi_{pref}$ 是参数预取策略。这是一个很大但可枚举的搜索空间，规划器的第一步工作是把它裁剪到只剩显存上真正可行的配置。对候选 $c$，stage $p$ 的显存峰值分解成三个加性项：

$$M_p(c) = M^{state}_p(c) + M^{act}_p(c, \pi_{act}) + M^{buf}_p(c, \pi_{pref}, \pi_{act}) \tag{9}$$

这里 $M^{state}_p$ 涵盖该层的本地参数分片、梯度状态、优化器状态（由 $P$、$D$、$Z$ 决定——$P$ 越大，每个 stage 分到的层越少，分片越小；$Z$ 越大，跨 $D$ 个副本切分得越激进）；$M^{act}_p$ 正是上面公式 5/6 的激活值驻留项，现在显式地按所选恢复策略参数化；$M^{buf}_p$ 涵盖通信/预取/恢复/算子工作区这些短生命周期缓冲区。一个候选只有在每一个 stage 同时满足硬性每集群预算时才会被接受：

$$\max_p M_p(c) \le M_{budget} \tag{10}$$

注意这个 $\max_p$：一个配置的好坏取决于它最差的那个 stage。这正是为什么随着模型规模变大，规划器会被推向更大的 $P$（这会按比例缩小每个 stage的层数,从而缩小每个 stage 的状态量，代价是更细粒度的流水线气泡），而不是别的任何单一手段。

用暴露延迟分解来估计 step 时间

对每一个显存可行的候选,规划器还需要估计速度，不只是可行性。论文把我们已经见过两次的 $\max(0, \cdot)$「暴露延迟」模式（公式 4 和 7）推广成一个统一定义，适用于任意可调度任务 $x$：

$$E_x(c) = \max\bigl(0,\, T_x(c) - W_x(c)\bigr) \tag{11}$$

其中 $T_x(c)$ 是该任务的延迟（来自离线执行画像），$W_x(c)$ 是 1F1B 时序结构、stage 本地空闲时间或某个有界调度窗口能吸收多少这段延迟。如果一个任务完全被重叠隐藏了，$E_x(c) = 0$；只有溢出部分才会算进 step 时间。据此，总 step 时间是：

$$T_{step}(c) = T_{1F1B}(c) + E_{comm}(c) + E_{upd}(c) + E_{pref}(c) + E_{rec}(c) \tag{12}$$

其中 $T_{1F1B}(c)$ 是主路径执行时间（前向/反向的时间片，加上不可避免的流水线气泡和 stage 不均衡——也就是即使杂活开销为零你也会测到的时间），另外四个 $E$ 项分别是暴露出来的通信、更新、预取、恢复开销。这个分解之所以精妙，正是因为它把「配置本身内在的、无法避免的成本」（$T_{1F1B}$）和「调度层面产生的、原则上可以藏起来的成本」（四个 $E$ 项）区分开来——这正是论文整套设计哲学所依赖的那条分界线。

最终，规划器求解一个带约束的最小化问题：

$$c^{\star} = \arg\min_{c \in \mathcal{C}} T_{step}(c), \quad \text{s.t.} \max_p M_p(c) \le M_{budget} \tag{13}$$

算法 2（如下）把这变成对（已裁剪过的）候选空间的一次直接暴力搜索,而不是什么算法上很花哨的东西——考虑到搜索空间虽然大但是离散的、而且在花费任何 step 时间估计工作之前就能被显存可行性检查大幅预筛，这是个合理的设计选择：

算法 2：资源感知配置规划
输入：模型画像、平台画像、执行画像、搜索空间 C
输出：选中的训练方案 c*

 1: V <- 空集合
 2: for 每个候选 c in C do
 3:     按流水线度数 P 切分各层
 4:     估计每个 stage p 的显存 M_p(c)
 5:     if max_p M_p(c) > M_budget then
 6:         continue                      # 裁剪：不可行，跳过
 7:     end if
 8:     从前向/反向画像估计 T_1F1B(c)
 9:     估计暴露延迟 E_comm, E_upd, E_pref, E_rec
10:     T_step(c) <- T_1F1B(c) + E_comm(c) + E_upd(c) + E_pref(c) + E_rec(c)
11:     把 (c, T_step(c)) 插入 V
12: end for
13: return c* <- V 中 T_step 最小的候选

论文特别小心地说明的一点是：「规划器不是要取代端到端实测」。它是一个裁剪与优先级排序工具，用预先在同一平台收集的画像把一个庞大的配置空间压缩成一个短名单——不是一个可以盲目信任的先知。这一点很重要，因为规划器的估计确实存在实测误差（2.33%-2.94%，见 Table 4 和结果部分），一个自称从不需要经验验证的系统反而是过度宣称。

一个具体的例子：为什么 LLaMA-2-70B 需要 $P=48$

值得把论文自己 Table 3 里的真实数字代入公式 9-10 看看规划器为什么会落在那个点上，而不是把「70B 用 $P=48$」当作一个没有解释的结论直接接受。LLaMA-2-70B 大约有 80 层 Transformer block。如果规划器尝试用 $P=8$（对小得多的 Qwen2.5-32B 效果不错的度数），每个 stage 大约要持有 10 层的参数、梯度、优化器状态,以及输入侧 stage 若干在飞 micro-batch 的激活值——而 70B 每层的状态量本身大约是 Qwen2.5-32B 的 $70/32 \approx 2.2$ 倍，粗略推算就能看出 $P=8$ 时的 $M_p(c)$ 大概率会冲到 30GB 以上，远超公式 10 的 20GB 上限。规划器的暴力搜索（算法 2 第 3-6 行）只是不断尝试更大的 $P$——这会大致线性缩小每个 stage的层数从而缩小 $M^{state}_p(c)$——直到 $\max_p M_p(c)$ 终于压过预算线；Table 3 显示这发生在 $P=48$ 时，落在 19.46GB，刚好卡在 20GB 上限之下，几乎没有余量。这也解释了为什么 $P=48$ 配的是 $D=2$（很小的数据并行度）：96 个集群的总预算（$P \times D = 48 \times 2$）大部分已经花在满足显存约束的流水线深度上了，在这个最小可行运行点上，能留给数据并行的集群不多。

实现细节：运行时各部件如何拼在一起

在进入结果部分之前，有几个实现细节值得了解，因为不少消融和扩展性数字在这些上下文下才更容易理解。

Stage 本地运行时作为执行单元。 每个流水线 stage 被实现为一个独立的轻量级运行时，绑定到一个或多个 MT-3000 计算集群，按规划器选出的方案为自己负责的层区间执行前向、反向和状态任务。全局同步只在真正必要的点上发生——step 初始化、stage 边界通信（把激活值从 stage $p$ 交给 stage $p+1$）、以及累积边界。除此之外——逐层状态任务、激活值恢复——都由每个 stage 根据自己本地的依赖事件（层反向完成、本地梯度累积完成、参数更新完成、即将到来的截止时间）独立、局部地调度。这避开了一个全局细粒度调度器，论文指出这「符合 MT-3000 硬件的组织方式，计算集群本身就是基本执行单元」——也就是说软件架构映照的是硬件自身天然的自治单元，而不是强行套一个本身可能变成瓶颈或争用焦点的中心协调者。

按生命周期分级的显式显存管理。 RATrain 每个 stage 的显存管理器明确按对象存活时间分类：长生命周期（参数分片、优化器状态、元数据——分配一次，存活整个训练过程）、中生命周期（checkpoint、梯度桶、working-weight 缓冲区——大致存活一个 step）、短生命周期（临时激活值、通信分级缓冲区、恢复缓冲区、算子工作区——存活一个 step 里的一小段，而且被显式在恢复/预取/算子执行之间复用，而不是每次重新分配）。这种分级正是让显存预算公式（公式 9）能静态推理的关键：可以分别对每一级的贡献设界，而不需要动态跟踪每一次分配的确切生命周期。

通信争用优先级。 当通信通道或分级缓冲区发生争用时，运行时不会平等对待所有流量——它优先处理 stage 边界的传输（这是真正卡在 1F1B 关键路径上的：stage $p+1$ 收不到 stage $p$ 的激活值就没法开始前向传播），GradSync 等背景流量只有在自己的依赖满足并且资源确实空闲时才会被调度。这个优先级顺序隐含地承认了：论文里描述的「可以调度进某个窗口」的任务并非同等紧急——当多个可调度任务在同一时刻争抢同一条稀缺的集群间链路时，运行时仍然需要一条打破平局的规则。

实验与结果

整套评测完全跑在真实的 MT-3000 平台上（不是模拟），这一点值得先强调一下，是个可信度加分项——国产加速器系统类论文有时会部分依赖模拟结果，而这篇论文没有,唯一的例外是 A800 数字被明确标注为「参照性、非严格对照」（论文自己坚持要加这个限定，下文批判性分析部分我会回来讨论这个说法本身是不是还是承担了比它自称的更多论证功能）。

正确性验证：重新调度会不会改变学到的东西？

既然 RATrain 的整套前提是「改变状态操作发生的时机，而不是它们计算的内容」，第一个也是最重要的实验就应该是正确性检查，而不是速度检查。作者用 LLaMA-2-7B（序列长度 2048，全局批量 2048）跑了一次 10.28 亿 token 的训练,和一个不使用 RATrain 任何调度机制、语义等价的 Baseline-1F1B 并排对比——相同的 tokenizer、相同的初始权重、相同的数据顺序、相同的优化器/学习率调度、相同的梯度累积语义。

指标	Baseline-1F1B	RATrain
最终训练 loss	1.8312	1.8306
最终 loss 绝对差	—	0.00064
每步相对 loss 偏差最大值	—	0.081%
每步相对 loss 偏差均值	—	0.030%
每步相对 loss 偏差最终值	—	0.035%

Figure 6（对应论文 Fig.7a-b，重绘为表格）：整整十亿 token 的训练过程中，两条 loss 曲线几乎完全重合，每步相对偏差最大值不到千分之一。

这正是这套设计哲学理应给出的结果：既然 RATrain 从未改变计算图、micro-batch 顺序、梯度累积规则或优化器更新公式——只改变了状态任务的物化时机、缓冲区驻留、调度顺序——任何残余的数值差异都应该只能归因于因重排序（而非重定义范围）操作导致的浮点非结合性,而这正好就是这里观测到的量级（0.08%，不是 8% 或 80%）。如果这个结果大一个数量级，那就是一个危险信号,提示重排序过程里可能存在语义泄漏；而这个结果与「纯调度改变、无语义改变」是吻合的。

同一个实验还报告了一次参照性的吞吐对比：在同样 10.28 亿 token 的预算下，256 个 MT-3000 集群上的 RATrain 达到 29,069.73 tokens/s，而三套 8xA800 GPU 参照栈（HuggingFace+DeepSpeed、FSDP、Megatron）分别达到 24,084.54、25,702.36、20,914.00 tokens/s。论文小心地把这个对比限定为「参照性、非严格对照的跨硬件性能基线」——我认为这个限定是正确且必要的（架构、显存层次、互联组织都不一样，确实不是一次受控对比），但这个数字读起来还是很抓眼球，下文批判性分析部分我会说明读这个数字时需要一些什么样的谨慎。

与 GPU 式训练策略的端到端对比（同一硬件）

科学上更干净的实验——因为它固定了硬件，只改变训练策略——是把 RATrain 和五种 GPU 式策略对比，全部在同一个 MT-3000 后端上重新实现（同样的 GEMM/注意力反向/通信实现），这样对比隔离出来的是并行组织和调度选择，而不是底层 kernel 质量的差异：

模型	方法	最优配置	峰值显存(GB)	Step 时间(s)	Tokens/s	相对慢多少
LLaMA-2-13B	RATrain	P=2,D=128,Z=2,FSR	15.84	688.09	12191.13	1.00x
LLaMA-2-13B	TP-heavy	P=2,D=64,T=2,Z=2,FSR	16.51	826.53	10149.20	1.20x
LLaMA-2-13B	ZeRO-3-heavy	P=2,D=128,Z=3,FSR	14.73	717.93	11684.48	1.04x
LLaMA-2-13B	Backward Ckpt	P=2,D=128,Z=2,Ckpt	15.73	937.04	8952.21	1.36x
LLaMA-2-13B	Full-save	—	OOM	—	—	—
LLaMA-2-13B	Tuned PP/DP/ZeRO	P=2,D=128,Z=2,Ckpt	15.73	945.84	8868.90	1.37x
Qwen2.5-32B	RATrain	P=8,D=32,Z=2,FSR	14.71	1592.51	5267.52	1.00x
Qwen2.5-32B	TP-heavy	P=8,D=16,T=2,Z=2,FSR	19.45	1922.66	4363.01	1.21x
Qwen2.5-32B	ZeRO-3-heavy	P=8,D=32,Z=3,FSR	16.54	1798.78	4663.50	1.13x
Qwen2.5-32B	Backward Ckpt	P=8,D=32,Z=2,Ckpt	14.50	2162.54	3879.06	1.36x
Qwen2.5-32B	Full-save	—	OOM	—	—	—
Qwen2.5-32B	Tuned PP/DP/ZeRO	P=8,D=32,Z=2,Ckpt	14.50	2167.81	3869.62	1.36x

Figure 7（对应论文 Table 2，重绘）：完整版 RATrain 在两个模型上都胜出；TP-heavy 和 ZeRO-3-heavy 之所以输，是因为它们比轻量的 PP+DP+ZeRO-2 引入了更多集合通信；Full-save 直接 OOM；「Tuned PP/DP/ZeRO」（可以搜索并行方式但关掉了 RATrain 三大调度机制）收敛到和纯 Backward Ckpt 差不多的水平，印证了收益来自调度而不是单纯并行搜索。

这一张表里塞了四条各自独立、可被推翻的结论，值得拆开来看，而不是笼统读成一句「RATrain 赢了」：

1. TP-heavy 输了（慢 1.20-1.21 倍），因为张量并行降低了每卡本地计算量，但引入了层内激活值集合通信，还压缩了能达到的数据并行度——直接印证了前面动机部分的错配一。这是论文对任何想不加修改就照搬 Megatron-LM 打法的人最重要的一条负面结果。
2. ZeRO-3-heavy 输了（慢 1.04-1.13 倍），因为激进的参数切分带来了参数视图物化和同步开销，而 ZeRO-2（PP+ZeRO-2 已经满足 20GB 预算）根本不需要付这笔钱。 这个结果格外有意思，因为 ZeRO-3 常被当作「只要能用就更省显存，所以用它」——这个实验展示了一种反例：当 ZeRO-2 已经够用时，用它严格更好，因为多出来的切分买来的是你不需要的显存余量，代价却是你付不起的通信。
3. Backward Ckpt 输了（两个模型都慢 1.36 倍），纯粹因为它把恢复延迟暴露在反向关键路径上——和 RATrain 完全相同的并行配置，唯一的区别只在激活策略（Ckpt vs FSR）。这是全文里对「FSR 单独就值大约三分之一的 step 时间」最干净的隔离证明。
4. Tuned PP/DP/ZeRO——允许搜索 $P, D, Z, b, A$，但关掉 RATrain 三大调度机制——最终收敛到和纯 Backward Ckpt 基本一样的水平（1.37x vs 1.36x 慢）。 这是论文对「RATrain 的收益其实只是并行度调得更好，调度机制只是个噱头」这种质疑最有力的反驳：只调并行度、不上调度机制，基本调不出什么改善。

资源受限训练能力：找到最小可行显存足迹

一个不同但互补的问题：对每个模型规模，在 20GB 每集群硬性上限下，RATrain 至少需要多少个 MT-3000 集群才能训得起来？

模型	最小集群数	配置	峰值显存(GB)	Step 时间(s)	Tokens/s
LLaMA-2-7B	8	P=2, D=4, A=128	19.57	1304.13	804.04
Baichuan2-13B	16	P=8, D=2, A=128	19.06	743.15	705.50
Qwen2.5-32B	64	P=16, D=4, A=128	18.14	873.85	1199.96
LLaMA-2-70B	96	P=48, D=2, A=16	19.46	281.32	232.96

Figure 8（对应论文 Table 3，重绘）：随着模型规模变大，规划器扩大的是流水线度数 $P$（不是张量并行，不是 ZeRO-3）来控制每 stage 显存——70B 模型上一路推到 $P=48$，每种情况下峰值显存都被死死压在 20GB 上限之下、留很少余量。

这张表里的规律,是「避开层内集合通信」这条设计原则在实证上最清楚的印证：Table 3 里每一个配置都用 $T=1$（完全不用张量并行）和 ZeRO-2（不是 ZeRO-3），仅靠扩大 $P$ 和 $D$ 来让模型从 7B 长到 70B 时仍能塞进显存预算。在 GPU 集群上，多数工程师碰到单卡显存不够的第一反应是上张量并行；这张表直接证明了在 MT-3000 的带宽特性下,这种本能反而是反效果的,而流水线并行——只需要 stage 边界处的点对点传输，而不是层内集合通信——才是该拉的那根杠杆。

规划器准确性：开销模型预测得准不准？

模型	集群数	预测 Step 时间(s)	实测 Step 时间(s)	误差
LLaMA-2-7B	256	140.92	144.28	2.33%
Baichuan2-13B	256	268.74	276.61	2.85%
Qwen2.5-32B	256	441.83	455.21	2.94%
Qwen2.5-32B	512	225.47	231.36	2.55%

Figure 9（对应论文 Table 4，重绘）：规划器的 step 时间预测（公式 12）在不同模型规模和集群数下，误差落在 2.33%-2.94% 之间，平均误差 2.67%。

这一点在实践上很重要，因为它验证了规划器确实是一个能用的搜索空间裁剪工具，而不是套上公式外壳的花架子——误差在 3% 以内意味着规划器按预测 step 时间给候选配置排的序，几乎总能和真去跑每个候选拿到的排序一致,而这正是要一个规划器的全部意义（省掉真去跑每个候选才知道哪个最快的成本）。

序列长度敏感性：RATrain 是不是只为一种输入形状调过？

任何手调过的调度系统都有一个天然的担忧：是不是只为某一个特定序列长度调过，换个长度就悄悄崩了。作者在 LLaMA-2-7B、Baichuan2-13B、Qwen2.5-32B 上测试了 512、1024、2048、3072、4096 五种序列长度，全部用 256 集群、全局批量 4096，并报告了序列长度 2048 时具有代表性的 FP16 GEMM 后端画像：

GEMM 形状	MAC 利用率	吞吐(T MAC/s)	延迟(ms)
4096x4096	64.96%	5.26	6.53
4096x11008	66.16%	5.36	17.23
11008x4096	65.13%	5.28	17.50
6656x6656	67.35%	5.46	16.63
8192x8192	68.13%	5.52	24.90

Figure 10（对应论文 Table 5，重绘）：后端在代表投影、FFN 和注意力内部矩阵乘法的一系列 GEMM 形状上维持 64.96%-68.13% 的 MAC 利用率——更大、更「方」的 GEMM 利用率略高一些，符合「固定的流水线填充开销在单位有效计算上摊得更薄」的直觉。

整体来看，Baichuan2-13B 和 Qwen2.5-32B 的训练时间从序列长度 512 到 2048 一路下降，再往后又回升；三个模型的计算利用率都从 512 涨到 2048，在 3072/4096 处略有回落。直觉的解读是：短序列没能把固定的调度/通信/状态开销摊薄到足够多的有效计算上；很长的序列则会让注意力、激活值驻留、恢复压力增长得比有效计算密度还快。序列长度 2048 对这块硬件来说是一个实用的甜点位置——但关键的是，论文把这描述成规划器资源感知搜索的涌现结果，而不是 RATrain 里硬编码的固定假设；规划器原则上会在不同序列长度下选出不同的配置，正因为它的开销模型（公式 12）已经把序列长度如何改变注意力开销、激活值驻留、恢复压力考虑进去了。

消融实验：哪个机制贡献了多少？

在我看来，这是全文信息量最大的一张表，因为它在 Qwen2.5-32B（256 集群，固定并行配置 $P=8, D=32, T=1, Z=2, b=1, A=128$——只切换被测试的那一个调度机制）上把 RATrain 的三大机制单独隔离开：

变体	归一化 Step 时间	暴露长尾放大倍数
完整 RATrain	1.00x	1.00x
去掉 FSR	1.33x	—
去掉更新-预取（U-P）	1.01x	2.31x
去掉逐层状态流水线（LSP）	1.03x	4.59x

Figure 11（对应论文 Fig.11，重绘为表格）：单独去掉 FSR 对绝对 step 时间的代价最大（1.33x）；去掉两个针对收尾长尾的机制对绝对 step 时间的代价较小（1.01x-1.03x），但会让暴露的收尾长尾具体放大 2.31x-4.59x——因为在固定的 256 集群、32 路 DP 规模下，即使长尾管理得再糟，它在整体 step 时间里也只占很小一部分；相对放大倍数很大，但对整体 step 时间的绝对贡献不大。

论文自己这次消融实验的完整 RATrain 基线是 step 时间 1790.13 秒，暴露长尾只有 14.69 秒——值得明说，因为这意味着表面上的消融比例（1.33x、1.01x、1.03x）在某种程度上被压缩了：即便是收尾长尾管理最差的情况（4.69 秒放大到 4.59 倍）也只给 1790 秒的 step 增加了大约 53 秒，也就是总 step 时间的大约 3%——一个真实、可测量、方向正确，但在这个具体规模下数值上并不算大的效应。FSR 的效应在结构上完全不同：它根本不是一个 step 边界的长尾现象，而是一个反复出现的、按 micro-batch、按 stage 计的成本（回想公式 7 的 $E_{rec}(p)$ 项，每次反向传播到达一个 stage 都会触发它，不是每个 step 一次），这正是为什么关掉它会产生比关掉任何一个收尾长尾机制都大得多的绝对 step 时间变化（33%）。读者不应该把「长尾放大倍数很大」和「绝对影响很大」混为一谈——论文自己的数字已经把这两个轴区分得很清楚了，我认为这点值得明确指出来，而不是任由那个抓眼球的 4.59x 数字单独立在那里、不带这层上下文。

资源可扩展性：把集群数转化成吞吐

集群数	全局批量	Step 时间(s)	Tokens/s	相对加速	效率
256	2048	144.28	29,069.73	1.00x	100.0%
512	4096	145.75	57,558.07	1.98x	99.0%
768	6144	147.23	85,465.01	2.94x	98.0%
1024	8192	148.75	112,790.55	3.88x	97.0%

Figure 12（对应论文 Table 6，重绘）：在 LLaMA-2-7B 上做的吞吐导向扩展实验——全局批量随集群数线性增长，step 时间从 144.28s 只爬升到 148.75s，扩展到 4 倍规模时效率依然保持在 97%。

作者明确说明这是一次吞吐导向的扩展（全局批量随集群数增长），不是强扩展（固定全局批量，靠更多集群把同一个问题算得更快）——这是个重要的方法论区分，因为强扩展会随着 $D$ 增大成比例地加重跨集群通信压力（更多副本要对同样大小的梯度做 all-reduce），而这种吞吐导向的设置基本保持每个副本的工作量和通信模式不变,这是一种明显更容易维持效率的扩展方式。适度的效率下滑（100% 到 97%）被归因于梯度同步、运行时调度以及更大数据并行组规模下的系统一般性波动稍微叠加起来——按实验设计的方式，这是一个可信、不出人意料的解释。

相关工作对比

为了把 RATrain 放进上下文里，下表按论文讨论的四类前人工作重新整理了对比关系：

系统类别	代表系统	目标硬件	核心手段	与 RATrain 的关系
面向 GPU 的分布式训练	Megatron-LM, GPipe, PipeDream(-2BW), DeepSpeed/ZeRO	GPU 集群，快速互联	TP + PP + DP + 状态切分	RATrain 刻意避开把 TP 和 ZeRO-3 当作默认手段；把 PP/DP/轻量 ZeRO 加生命周期调度当作主路径
自动并行搜索	FlexFlow, Alpa, GSPMD, Whale	同构或异构 GPU 集群	搜索/编译计算图划分方案	RATrain 的规划器搜索的是一个更窄的空间（PP/DP/ZeRO/激活/预取策略），但显式建模暴露延迟和 step 末尾长尾，这些系统通常不建模
激活值显存优化	Checkpointing, Checkmate, Capuchin, bubble-filling 重计算	GPU 显存层次	用重算换显存	RATrain 的 FSR 与之互补：保留这些系统的标准 1F1B 顺序和 checkpoint 位置选择，但把重算的时机改到前向侧/空闲窗口
异构/offload 训练	ZeRO-Offload, Whale（异构 GPU）	GPU + CPU 显存分层，或异构 GPU 混部	用更慢的一层显存扩展有效容量	RATrain 不把一层更慢的存储当作 GPU 显存的延伸来用；它把参数/梯度/优化器状态/激活值/通信缓冲区都当作带有明确生命周期的一等运行时对象，跑在一块真正意义上非 GPU 的加速器上

Figure 13（综合论文第 7 节整理）：RATrain 的定位与其说是「某个已有技术的加速版」，不如说是「对调度问题本身的重新建模」——这正是它能与上述四类思路组合（而不是取代它们）的原因，比如它原则上可以作为 Alpa 式自动并行搜索底下的开销模型后端，或者用 Checkmate 式的重算位置策略、同时仍然套用 FSR 的前向侧时机调整。

局限与边界条件

明确指出这个结果不能泛化到哪里，和列出headline数字一样重要：

• 只验证了稠密 decoder-only 模型。 所有实验都用的是稠密 Transformer（LLaMA-2、Baichuan2、Qwen2.5）。MoE 模型有根本不同的通信模式（数据依赖的专家路由，而不是逐层的固定依赖），完全在本文范围之外,而训练状态生命周期这套论证（完全建立在稠密 Transformer 确定性的层序结构上）能不能迁移到 MoE 的数据依赖路由上并不显然。
• 只有一个硬件平台。 所有结果都是在 MT-3000 上测得的。论文的设计原则（集群间带宽紧张时避开层内集合通信；围绕层序调度状态生命周期；在通用调度层之下保留一个加速器特定的执行后端）可以合理地类比迁移到其他带宽受限的异构加速器上，但这是一个类比论证，不是一个已验证的结果——没有第二块硬件平台被测试过。
• A800 对比明确是参照性、非严格对照。 论文自己对此很小心,我想保留这个限定，而不是过度解读一种「RATrain 打赢了 8xA800」的粗糙印象。
• 10.28 亿 token 是正确性验证的规模，不是完整预训练规模。 接近零的 loss 偏差是「在这个规模下调度改变不会改变训练语义」的有力证据；它本身并不能证明一次完整的数万亿 token 预训练也会表现出同样可以忽略的漂移（不过考虑到这些改变只关乎时机/物化，不关乎计算内容,并没有明显的机制会让漂移灾难性地累积）。
• 规划器需要先在同一平台上做画像才能规划。 2.33%-2.94% 的预测误差是在作者已经在真实 MT-3000 平台上收集执行画像之后测得的。把 RATrain 的规划器部署到一个没做过画像的新平台（甚至是一个明显不同的 MT-3000 配置）上，都需要先重做这一步画像工作——规划器裁剪的是一个有好画像支撑的搜索空间；它不能消除获得画像本身的成本。
• 没有讨论容错或弹性扩缩容。 在 1024 个集群的规模下，硬件或链路故障对任何超算规模的训练任务来说都是现实的运维担忧,论文没有讨论 checkpoint/重启行为、集群数变化时的弹性重规划,或故障恢复——这是一个被完全留在讨论范围之外、但确实存在的系统性关切。
• 数据并行梯度同步的具体拓扑没有细说。 GradSync 的开销模型（公式 4）把通信当作一个暴露延迟标量处理，但没有说明底层的 all-reduce 在跨集群、经 CPU/GP Zone 中转的结构上具体用的是环形、树形还是分层拓扑——这个细节会实质性地影响 1024 集群下观测到的 97% 扩展效率再往上推会怎样变化。

批判性分析：不足与可改进之处

「参照性」的 A800 对比承担的论证功能，比论文自己那句限定语允许的更多。 摘要和引言都用了一种轻推读者往「RATrain 在 MT-3000 上能和 8xA800 GPU 栈打平」方向去理解的措辞，那个数字（29,069.73 对比 20,914-25,702 tokens/s）第一眼看确实很抓人。但这次对比控制了 token 预算和序列长度，却完全没控制架构、显存层次和互联——恰恰是决定训练吞吐的三个最大变量。一个更站得住脚的版本，要么 (a) 干脆去掉这次跨硬件对比，让同硬件的 Table 2 对比（那才是真正受控的）单独承担整个性能论证；要么 (b) 如果保留这次跨硬件对比作为背景信息，就应该把成本归一化后的吞吐（每美元、每瓦的 tokens/s）和原始 tokens/s 一起报告出来——因为「同样的 tokens/s，但只花一小部分 A800 集群的成本/功耗」很可能才是作者真正想让读者带走的结论,而目前这个实验，在没有功耗/成本数字的情况下，是支撑不起这个结论的。

没有跨模型规模的消融实验。 整个消融实验（Table/Figure 11）只在一个模型 Qwen2.5-32B、一个集群数 256 上跑。既然 Table 3 显示主导性的杠杆会随模型规模变化（流水线度数从 7B 时的 $P=2$ 涨到 70B 时的 $P=48$），那么完全有可能 FSR 相对两个收尾长尾机制的相对贡献也会随模型规模变化——比如在 $P=48$ 时，流水线 stage 多得多、每个 stage 计算量相应小得多，step 末尾收尾长尾可能会占到比 $P=8$ 时观测到的 3% 大得多的比例。论文在这一点上完全没有给出证据，单一规模的消融，无论做得多仔细，都排除不了这种可能。

正确性验证的基线和性能对比的基线不完全是同一个「Baseline-1F1B」。 第 6.2 节的正确性实验对比的是 RATrain 和「一个语义等价的 Baseline-1F1B」，而第 6.3 节的性能实验对比的是五种不同的 GPU 式策略（TP-heavy、ZeRO-3-heavy、Backward Ckpt、Full-save、Tuned PP/DP/ZeRO）。如果能至少对其中一个性能基线（理想情况是 Backward Ckpt，因为它和 RATrain 用的是完全相同的并行配置，唯一区别只在激活策略）也报告一下正确性结果（loss 曲线、相对偏差），会大大增强论文的说服力——这样才能确认同样接近零的语义漂移，在被拿来比速度的那些具体配置上也依然成立，而不只是在一个精心挑出来的参照基线上成立。

规划器的开销模型只在它自己选中的配置上被验证过，没在它拒绝的配置上验证。 Table 4 报告的是规划器在每个模型/集群数组合下选中的配置的预测准确度。一个真正更有说服力的验证，还应该报告规划器拒绝掉的若干配置（比方说一个显存可行、但被规划器放弃转而选了 ZeRO-2 的 ZeRO-3 候选）的预测-实测对比——这样才能直接证明开销模型对不同候选之间的相对排序，而不只是对赢家的点估计准确度，是符合实际的。现在论文展示的是规划器选出的东西接近它自己的预测,但没展示它是不是正确地预测了为什么其他候选更差。

如果我给这篇论文做系统类会议评审，我会补充： (1) 一次针对 20GB 显存预算本身的敏感性研究——如果预算是 15GB 或 30GB，规划器选出的配置以及相对 GPU 式基线的加速比会怎么变？这能澄清 RATrain 的优势是这套方法的一般性质，还是恰好针对 MT-3000 特定 20GB 约束调出来的。(2) 一个第二块异构加速器平台（即便是模拟的，只要明确标注），来实证支撑目前只靠类比论证的「这些原则能超越 MT-3000 迁移」这句话。(3) 一次显式的、成本/功耗归一化后的 A800 参照栈对比,因为这很可能是「用 MT-3000+RATrain 替代商用 GPU」这件事最有力的现实论据,而论文恰好把做出这个论证所需的证据留在了桌面上没有采集。

把整套流程串起来：一个训练 step 的伪代码

前面四段「深入」很容易让人只见树木不见森林。下面把整个 RATrain 训练 step 从头到尾写成一份伪代码，展示四大机制分别插在不变的 1F1B 主路径的什么位置。这是我自己整合论文算法 1-2 和 Figure 4-6 写出来的一份控制流总览；论文原文没有恰好以这种形式出现，但每一行都能追溯回上文某个具体机制。

一个 RATrain 训练 step（针对流水线 stage p，给定算法 2 选出的方案 c*）：

 1: # ---- 预热/稳态/收尾主路径（1F1B 本身不变）----
 2: for 1F1B 顺序中每个已调度的 Forward(l) 或 Backward(l) 任务 do
 3:     if 任务是 Forward(l) 且该层使用激活值恢复（pi_act = FSR）then
 4:         if 在对应 Backward 预计到达之前存在可用的空闲或前向侧窗口 then
 5:             把 Recover(l) 调度进这个窗口                     # 深入之三（FSR）
 6:         end if
 7:         经 MT-3000 感知的 GEMM/注意力后端执行 Forward(l)      # 深入之一
 8:     else if 任务是 Backward(l) then
 9:         if Recover(l) 已经在更早的窗口里完成 then
10:             直接消费恢复缓冲区（E_rec(l) = 0）
11:         else
12:             退回到反向时恢复（E_rec(l) = T_rec(l) - W_rec(l) > 0）
13:         end if
14:         经 MT-3000 感知的注意力反向后端执行 Backward(l)        # 深入之一
15:         if 该层本地梯度累积现在已完成 then
16:             标记 GradSync(l) 为可调度                          # 深入之二
17:         end if
18:     end if
19: end for
20: # ---- Stage 本地生命周期调度，与上面并发运行 ----
21: for 每个 GradSync(l) 已可调度的层 l do
22:     调度 GradSync(l) 与后续 Backward(l-1), Backward(l-2), ...
        或 stage 本地空闲时间重叠                                # 深入之二，公式 2
23:     一旦 GradSync(l) 在 t_sync(l) 完成：
24:         在 [t_sync(l), t_use(l)) 内调度 UpdateShard(l)
            然后 PrefetchW(l)                                    # 深入之二，公式 3
25:         if PrefetchW(l) 在 t_use(l) 之前完成 then
26:             下一次 Forward(l) 使用热的、已预取的 W_view(l)     （E_upd=E_pref=0）
27:         else
28:             下一次 Forward(l) 停顿等待未覆盖的部分             （E_upd, E_pref > 0）
29:         end if
30: end for
31: # ---- Step 边界 ----
32: 累积边界处，只剩真正的全局同步
     （是为下一个 step 做初始化，不是为这一个 step 做逐层杂活）

一句话总结这份伪代码：第 2、7、8、14 行（真正的前向/反向计算）相对朴素的 1F1B 基线什么都没变；RATrain 加的所有东西都活在第 3-6、9-13、15-30 行——关于何时运行恢复/同步/更新/预取任务的调度决策，而基线要么会急着立刻做这些事，要么会把它们推迟到 step 边界统一处理。

公式索引：本文出现的每一个公式,交叉引用

公式	计算的是什么	用在哪里
1	注意力反向传播的 AM/SM/GSM tile 容量约束	约束算法 1 里合法的 tile 形状 $B_r, B_c, d$
2	逐层状态任务链 GradSync -> UpdateShard -> PrefetchW	定义逐层状态流水线（深入之二）
3	可调度窗口 $[t_{sync}(l), t_{use}(l))$	更新-预取调度的截止时间约束
4	暴露的更新/预取延迟 $E_{upd}, E_{pref}$	$\max(0,\cdot)$ 暴露延迟模式首次出现
5	Full-save 激活值显存峰值 $M_{act,full}(p)$	FSR 用来改进的基线显存成本
6	FSR 激活值显存峰值 $M_{act,FSR}(p)$	展示 FSR 显存足迹为何小得多
7	暴露的激活值恢复延迟 $E_{rec}(p)$	同一 $\max(0,\cdot)$ 模式用在 FSR 的退化情形上
8	候选配置元组 $c = (P,D,Z,b,A,\pi_{act},\pi_{pref})$	规划器的搜索空间表示
9	每 stage 显存峰值分解 $M_p(c)$	可行性检查的输入
10	硬性可行性约束 $\max_p M_p(c) \le M_{budget}$	在算法 2 里裁剪候选空间
11	通用暴露延迟定义 $E_x(c)$	把公式 4 和 7 统一进同一套记账体系
12	总估计 step 时间 $T_{step}(c)$	规划器的目标函数
13	带约束的最小化 $c^\star = \arg\min T_{step}(c)$	算法 2 真正求解的问题

留意这张表里的结构性模式：公式 4、7、11 都是同一个 $\max(0, T_x - W_x)$ 想法的具体实例，先用在更新/预取上，再用在激活值恢复上，最后推广成「任意可调度任务」。这是刻意的统一，不是巧合——正是这一点让一个规划器（公式 12）能把三种看起来毫无关系的机制的开销，加总进同一个 step 时间估计里。

一览设计取舍

决策点	RATrain 的选择	显而易见的 GPU 式替代方案会怎么做	在哪里被测试
大模型的主要扩展轴	流水线并行（70B 时 $P$ 一路到 48）	张量并行（切单层内部）	Table 3（最小可行配置）、第 6.3 节
ZeRO 阶段	ZeRO-2（切梯度+优化器状态）	只要显存允许就上 ZeRO-3（连参数也切）	Table 2（ZeRO-3-heavy 慢 1.04-1.13 倍）
激活策略	FSR（提前在空闲/前向侧窗口恢复）	经典的反向时检查点	Table 2（Backward Ckpt 慢 1.36 倍）；Figure 11（消融）
状态任务时机	逐层、事件触发（依赖满足就立刻做 GradSync/UpdateShard/PrefetchW）	所有反向传播跑完后统一批量处理	Figure 11（去掉 LSP 后长尾放大 4.59 倍）
注意力反向 tiling	外层遍历 query、内层遍历 K/V，驻留 AM/SM，经 GSM 归约	针对 GPU HBM 优化的 FlashAttention 式 tiling	Figure 10（相对 DDR 分级基线快 1.24-1.54 倍）
配置选择	离线画像驱动、对离散空间做搜索（算法 2）	按模型规模手调经验法则	Table 4（预测误差 2.33%-2.94%）

常见问题解答

RATrain 需要改模型结构或优化器吗？ 不需要。所有机制都严格作用在调度/运行时这一层。计算图、梯度累积规则、优化器更新公式都明确没有被动过（论文第 5.3 节，也被上文正确性研究实证验证过，见 Figure 6）。

RATrain 是 Megatron-LM 或 DeepSpeed 的替代品吗？ 不完全是——它是专门为 MT-3000 特定的显存层次和带宽特性打造的训练运行时。概念上它扮演的角色和 Megatron-LM/DeepSpeed 在 GPU 集群上扮演的角色类似，但具体机制（GEMM/注意力反向后端、逐层状态流水线、FSR）都是后端特定的，换一块加速器就需要大改，尽管调度哲学本身（把状态当作生命周期管理的对象，带宽紧张时避开层内集合通信）原则上是可迁移的。

为什么论文几乎完全避开张量并行？ 因为 TP 的定义性特征——每一层前向和反向传播内部的一次 all-reduce 或 all-gather——恰好是 MT-3000 实测 3.7GB/s 集群间带宽吃不消的那种高频、延迟敏感的集合通信。相反，流水线并行只需要 stage 边界处的点对点激活值传输，发生频率低得多（每个 stage 边界每个 micro-batch 一次，不是每层每个 micro-batch 一次）。

如果前向侧恢复窗口的估计错了，FSR 会不会悄悄出错、污染训练？ 不会——这是整套系统设计得比较细致的一个地方。如果前向侧窗口实际不够用，RATrain 会退回到普通的反向时恢复（公式 7 的 $\max(0,\cdot)$ 形式正是让这种退回成立的原因：「暴露」出来的成本就是没能藏起来的那部分，从零平滑地一直涨到反向时恢复的全部成本作为严格上限）。训练语义——算了什么，不是什么时候算——无论哪种情况都保持不变。

RATrain 的想法对 GPU 集群有帮助吗？ 部分有。GPU 集群通常带宽够用，错配一（层内集合通信）和错配三（step 末尾长尾）没那么要命，但错配二（激活值驻留/恢复）在 GPU 上也是真实成本——FSR 的核心想法（在空闲/前向侧窗口恢复激活值，而不是在反向关键路径上）本身是架构无关的，原则上可以移植进一个面向 GPU 的流水线并行运行时,不过论文没有测试这一点。

如果规划器的显存估计有一点点偏差，选出来的配置实际跑起来 OOM 会怎样？ 论文没有描述这种情况下的显式运行时兜底方案。考虑到规划器自己报告的 2.33%-2.94% step 时间预测误差，以及 Table 3 里若干选中配置离 20GB 上限只差零点几个 GB（比如 LLaMA-2-70B 是 19.46GB），一点没被开销模型捕捉到的额外显存压力就有可能把一个「可行」配置推向真正的 OOM——这是一个值得指出、而不是默认不存在的缺口。

复现性说明

论文全程依赖真实硬件评测,并给出了足够的配置细节（流水线/数据并行/ZeRO 度数、micro-batch 和累积步数、显存预算），足以精确重建每个实验的训练配置——Table 2、3、6 尤其给出了精确的 $(P, D, Z, b, A, \pi)$ 元组。但完整复现被硬件访问权限死死卡住：MT-3000 是一款专用超算加速器,不是商用可购买的 GPU，所以在能接触到这块具体平台（或等价平台）的机构之外，复现在今天基本不可行。这里抓取到的论文正文没有提到代码开源；想在此基础上做后续工作的读者需要根据论文描述自己重新实现 GEMM/注意力反向后端（算法 1、Table 1 的 VLIW 调度）和规划器（算法 2）——论文描述得足够详细，可以尝试，但从零实现出来的数字需要拿论文报告的延迟画像（Table 5）单独核实过才能信任。

结语

RATrain 的贡献与其说是一个单一的新 kernel 或一个聪明的调度小技巧，不如说是一次重新建模：稠密大模型训练里的梯度同步、优化器更新、参数准备、激活值管理，不再是「step 之间发生的杂活」，而变成带有按层索引的出生时间、截止时间、以及两者之间可调度窗口的运行时对象——这直接来自每一个稠密 Transformer 本来就自带的、确定性的前向/反向反向序遍历。这次重新建模，配上一个加速器特定的执行后端和一个资源可行性规划器，让一个真正受限的平台（每集群 20GB、集群间带宽 3.7GB/s）能以 97% 的扩展效率把 7B 到 70B 的稠密模型训到 1024 个集群规模，在同一硬件上比它们尝试过的每一种 GPU 式基线策略都快 1.04-1.37 倍，而且做到这一切的同时,对实际学习轨迹的扰动不超过千分之零点八。给任何在非 GPU 加速器上搭建训练基础设施的人最直接的一条经验：别问「我怎么把 Megatron-LM/DeepSpeed 搬过来」，要问「这块芯片具体的显存层次和带宽特性让什么东西变贵了，我能不能围着它重新调度,而不是硬扛这个代价」。