SPPO:把百万 token 训练拆成可调度的序列流水线
SPPO 怎样处理 activation:卸载粒度从整层或整条序列细化到 subsequence。系统联合选择 sequence parallel、pipeline parallel、分段数和每段卸载比例,让 CPU 传输与下一段计算重叠;边界重切分则用来减少流水线 bubble。
本文目录 · 18 节
1. 长序列训练中的 activation 爆炸
长序列训练时,模型参数只随模型规模增长,activation 却随 batch、sequence length、hidden size 和层数增长;attention 中间状态还会带来更陡的开销。论文的量级例子很直接:7B 模型、4M token 序列约需 16,384GB activation;65B、4M token 的 activation footprint 可达 80TB。只靠 GPU 保存不现实,activation checkpointing 又要在 backward 前重算 forward,CPU offload 则可能把问题从显存变成 PCIe stall。
作者比较的 baseline 是 Megatron-LM 与增强后的 DeepSpeed-Ulysses。Megatron 依靠 activation recomputation,序列越长重算越重;Ulysses 按 attention head 切 sequence,受 head 数限制,而且并不能自动消除 activation 保存问题。既有 offload 通常按固定长度或固定 FLOPs 比例搬运所有 subsequence:短段计算不够长,盖不住 D2H;长段又可能搬得过多。已有 subsequence pipeline 虽能把序列当 micro-batch 推进,但当 pipeline stage 数为 p=4、subsequence 数为 N=16 时,论文指出 bubble ratio 仍可到 3/16。
2. 前置知识:一次 Transformer 训练迭代到底保存什么
训练和推理的内存结构不同。推理只做 forward,历史 token 的 K/V 需要跨 decode step 保存;训练还要做 backward,每层 forward 产生的输入、归一化结果、Q/K/V、attention 输出、MLP 中间值等都可能在求梯度时再次使用。若不保存,就必须在 backward 前重新执行相应 forward,这就是 activation recomputation。若保存到主存,就要在 forward 后做设备到主机传输,并在 backward 前做主机到设备预取。两种办法都没有改变数学梯度,但分别多付计算和数据搬运。
设 batch 为 B、总序列为 S、hidden size 为 H。仅一个形状为 [B,S,H] 的 FP16 tensor 就占 2BSH bytes;一个 Transformer layer 同时存在许多倍数不同的中间 tensor,跨几十层累加后远大于参数。FlashAttention 能避免显式保存完整 S×S attention matrix,却不能让所有线性随 S 增长的 activation 消失。因此序列从几十万增到几百万时,问题仍然会到 TB 级。
causal attention 又带来特殊依赖。把总序列切成 s0,s1,... 后,第 i 段的 query 必须看 0..i 段的 K/V;它不需要未来段。因此早期 K/V 的下一次访问非常近、复用次数很多,而本段 Q 和许多后续只给 backward 使用的 tensor,在 forward 剩余阶段长期不访问。SPPO 的 tensor 分类正是从这个 lifetime 差异出发。
3. 为什么三类常见方案都不够
只做 activation checkpointing
checkpointing 保存少量边界 activation,backward 时重跑区间内 forward。它不占 PCIe 带宽,适合计算便宜或主存带宽差的环境;但超长 attention 本身很贵,重算等于把已经昂贵的 forward 再做一次。论文给 baseline 开启 checkpointing,仍观察到 Megatron 在长序列上的吞吐和最大长度劣势,说明容量虽缓解,额外 FLOPs 已进入关键路径。
整层或整序列 offload
粗粒度 offload 的优点是实现简单,缺点是传输工作量与序列近似线性增长,而可用的遮蔽窗口不一定同步增长。把仍会在下一段 attention 使用的 K/V 搬走,会立刻产生回传;把所有 backward-only tensor 全搬走,又可能让 D2H 超过下一段计算时间。论文消融中 GPT-13B、512K 的 full offload 只有不 offload 的 0.91x,正是传输暴露的结果。
只做 sequence parallel 或固定 pipeline
Ulysses 把 sequence 维交换成 head 维,每卡负责部分 heads,要求设备数与 head partition 相容;模型 head 数有限时,扩展到更多卡受限。pipeline parallel 把层分 stage,subsequence 像 micro-batch 一样流动,能降低单卡峰值,却有 warm-up/cool-down 空洞。盲目增加 subsequence 数会减少 bubble ratio,却让单段太小,attention kernel 和通信效率下降。SPPO 因此不能只选一种并行原语,而要联合搜索。
4. SPPO 的完整数据流
- 切分候选:给定模型、总序列长度和节点数,heuristic solver 搜索
(SP, PP, N)。它避免把高通信量的 sequence parallel 随意跨节点,也避免低带宽链路上的不合适 pipeline placement,并把每层每卡 workload 限在经验上高效的 2K 到 16K token 区间。 - 生成 subsequence pipeline:因 causal mask,第
i段 attention 会用到前面各段的 K/V;各段仍按前向、反向依赖经过 pipeline stages,而不是把原始 token 顺序打乱。 - 二级 activation 管理:K/V 会被后续 subsequence 反复读取,所以留在 GPU;只在 backward 才再次需要的其他 activation 才是 offload 候选。论文估算 K/V 常驻约为
2BSHbytes,其余可部分卸载的 activation 约为36BSHbytes。 - 逐段求卸载比例:传输第
i-1段 activation 时,同时计算第i段。每段的alpha_i不相同:较早、较小的段往往多卸载,后段少卸载,使 D2H 时间尽量等于可用于遮蔽它的计算时间。 - 检查 bubble:若最优卸载后 pipeline 边缘仍空闲,SPPO 把靠近 warm-up/cool-down 边界的 subsequence 再切成 sub-subsequence,临时用更宽的 SP range 计算;中部 steady phase 仍按常规 PP 运行。
- 执行与回传:page-locked CPU buffer 和 NUMA binding 保证 D2H 带宽;forward 时按计划卸载,backward 前预取。最终参数更新语义与原训练一致,SPPO 改的是激活驻留与执行顺序,不是训练目标。
5. 关键公式与算法
给定按 FLOPs 平衡的分段 s0 <= s1 <= ... <= sN: alpha0 >= alpha1 >= ... >= alphaN, alpha_i in [0, 1] 第 i 步的 GPU activation 占用: M_i = M_(i-1) + A_i - alpha_(i-1) * A_(i-1) 选择 alpha,使每段搬运量接近同一阈值: alpha_i * A_i = M_threshold M_threshold / BW_D2H ~= T_balanced_comp 并行参数搜索目标: (SP*, PP*, N*) = arg min T_iteration(SP, PP, N)
这组式子的重点不是精确求一个全局闭式解,而是把“卸载多少”绑定到下一段能提供多少计算遮蔽。固定全卸载可能节省显存却变慢;固定不卸载可能 OOM。SPPO 的 alpha 是连接显存上限、D2H 带宽和 pipeline timing 的控制量。
6. 公式中每个变量如何对应运行时对象
A_i 不是该段所有 tensor 的总和,而是经过二级管理后可选择 offload 的 activation volume;需要持续前向复用的 K/V 已被排除。alpha_i A_i 是该段实际进入 pinned CPU buffer 的字节数,(1-alpha_i)A_i 留在 GPU 等 backward。BW_D2H 不能填硬件宣传峰值,而应使用 NUMA 绑定、page-locked buffer、并发 CUDA copy 条件下测到的有效带宽。
T_balanced_comp 是相邻 subsequence 在相关 pipeline stage 上可用于遮蔽传输的计算时间。若下一段 attention/MLP 很短,即使主存充足,M_threshold 也必须降低,否则 copy 结束不了。M_i 同时包含新段 activation 与上一段未搬完的残留,所以最危险的不一定是最大段,而是“新段已产生、旧段仍占 buffer”的重叠时刻。这个 recurrence 让 solver 检查动态峰值,而不是只加静态 tensor size。
单调关系 s0<=...<=sN 来自按计算均衡得到的非均匀分段,早期 causal attention 工作少,后段要看更长前缀、计算更多。为让各段搬运量接近统一阈值,小段需要较大 alpha、大段较小,因此 alpha0>=...>=alphaN。若模型 attention 实现或 mask 改变这一计算曲线,单调假设也要重新验证。
7. 一个完整例子
设 4 个 pipeline stages 处理 8 个 subsequences。普通排程中,stage 0 先做 s0,stage 1 尚无输入;尾部则反过来,后段空闲。对 s0 的 activation 全量 D2H,如果传输跨过 s1 甚至 s2 的计算,就同时占着新段 activation 与旧段 offload buffer,既形成显存峰值也造成 stall。
SPPO 先把 K/V 留在 GPU,只将 backward-only tensor 纳入卸载。假设 A0 较小、A1 较大,它让 alpha0 > alpha1,但保持 alpha0*A0 ~= alpha1*A1,从而每次传输都能被下一段的 balanced computation 覆盖。若得到的 8 段排程仍在左边界留出空档,论文 Table 3 的 multiplexing 会让 stage 0 左侧处理 {0,1,2},stage 1 处理 {0,1},stage 2 处理 {0},并为这些边界段分配互补的 SP ranges;中间 {3..7} 一带继续走 steady PP。这个例子展示了两层优化的顺序:先消除传输 stall,再利用仍存在的 bubble,不能把两者当成同一个技巧。
8. 从输入配置到一次迭代的伪代码
profile(model, hardware):
measure compute time for candidate subsequence lengths
measure effective D2H/H2D bandwidth and collective cost
plan(model, total_sequence, devices):
candidates = enumerate_valid_SP_PP_N()
for plan in candidates:
segments = flops_balanced_partition(plan.N)
classify activations into resident_KV and offload_candidates
solve alpha_i from overlap and peak-memory constraints
estimate iteration_time = compute + exposed_copy + bubbles
choose minimum-time feasible plan
if residual bubble is large:
create multiplexed boundary sub-subsequences
run_iteration(plan):
for each forward subsequence in pipeline order:
compute current subsequence
keep its K/V on GPU
asynchronously offload alpha_i of other activations
for each backward subsequence in reverse dependency order:
prefetch offloaded tensors before first backward use
compute gradients and release activation
perform optimizer step
实际执行还需为 copy 与 compute 使用独立 CUDA streams,并插 event 保证“copy 读完 tensor 后才能复用显存”“prefetch 完成后 backward 才能消费”。multiplexed boundary 段改变临时 SP group,必须保证 collective 顺序在所有 rank 一致,否则会 deadlock。论文的 1,155 行 multiplexing 代码说明它不是图上移动方块那么简单。
9. 实现细节为何直接影响结论
CPU buffer 必须 page-locked,否则 CUDA 异步传输可能退化为额外 staging copy;进程要绑定到离本 GPU 最近的 NUMA node,否则跨 socket 内存访问会吃掉 PCIe 带宽。节点内 NVLink/NVSwitch、节点间 InfiniBand 的拓扑决定 SP/PP 放置,论文的 heuristic 避免让高频通信跨低带宽边界。主存每节点 2TB 也很关键:系统把 GPU 容量压力转移到 CPU,不代表 CPU 容量无限。
SPPO 基于 Megatron-LM 修改约四千行,说明它复用参数、梯度和 optimizer state 的现有并行,只接管 sequence segmentation、activation hooks、copy buffer 与 pipeline schedule。它不改变 loss、optimizer 或 data order,因此数值语义应与同精度 baseline 相同;复现时可用小序列关闭 offload 的 run 对齐 loss curve,再扩大长度验证 memory/timing,避免把通信错误误当训练噪声。
10. 实验设置与准确结果
实现基于 Megatron-LM,约 4,000 行 Python/CUDA,其中 adaptive offloading 2,845 行,multiplexing sequence partition 1,155 行。物理集群有 16 台服务器,每台 8 张 80GB NVIDIA Ampere GPU、128 CPU cores、2TB 主存;节点内 NVLink/NVSwitch,节点间四路 200Gbps HDR InfiniBand,GPU-CPU 带宽 32GB/s。测试 GPT-7B、13B、65B,使用 32、64、128 GPUs,指标为 training tokens/GPU/second;baseline 均开启 activation checkpointing,DeepSpeed 额外结合 ZeRO 与 FPDT。
| 场景 | SPPO 结果 | 对照 |
|---|---|---|
| 端到端吞吐 | 跨模型和长度为 1.13x 到 3.38x | Megatron-LM / DeepSpeed-Ulysses |
| GPT-65B, 600K / 640K | 3.38x / 3.12x | 相对 Megatron-LM |
| GPT-13B 最大长度 | 1280K | Megatron 768K;Ulysses 无法训练该配置 |
| GPT-65B 最大长度 | 1024K | Ulysses 512K;Megatron 768K |
| GPT-7B 扩展 | 128 GPUs 上 4M token | Megatron-Tuned 约 2.13M,SPPO 高 88% |
11. 消融告诉我们什么
GPT-13B、512K 中,直接 full offload 相对不卸载只有 0.91x,说明“主存更多”不等于“训练更快”;加入 adaptive offload 后为 1.23x,再加入 multiplexing 为 1.44x。GPT-65B、640K 的算力足以遮蔽更多传输:full offload 为 2.1x,adaptive offload 为 3.38x,multiplexing 最终到 3.48x。相同机制在不同模型上的贡献不同,正是自适应设计存在的理由。论文也承认 GPT-7B、768K 等计算量不足的点上,SPPO 只比强 Megatron 配置略快。
12. 如何逐行解读端到端图
GPT-7B 在 512K 到 1024K 的多点加速约 1.13x 到 1.29x,说明小模型已有强 kernel 与 checkpoint baseline 时,SPPO 的优势主要是延伸 OOM 边界,而非每点都大幅提速。GPT-13B 支持到 1280K,Megatron 到 768K;这里横轴越过 baseline 的部分没有速度比,只有可运行性。GPT-65B 在 600K/640K 的 3.38x/3.12x 是最强性能点,因为重算成本高且单段计算足以遮蔽 copy。
Figure 12 的 32/64/128 GPU 序列长度扩展用 GPT-7B,SPPO 相对基准达到 1.3、2、4 倍,Megatron-Tuned 为 1、1.375、2.13 倍。128 GPU 的“高 88%”是最大可支持长度之比,不是 tokens/s。Ulysses 从 32 到 64 GPU 的异常限制与 head partition 有关,不能泛化成 DeepSpeed 在所有模型上都随卡数退化。
13. 性能模型的适用区间
当序列很短时,activation 能直接驻留 GPU,copy 与复杂 schedule 没有必要;当序列极长但下一段计算仍不足以遮蔽 D2H 时,SPPO 会保留更多 tensor、最终受 GPU 容量限制;当 CPU 主存或 PCIe 被其他任务争用时,离线 BW_D2H 会高估收益。最适区间是 activation 明显超过 GPU 余量、模型计算又足以覆盖部分传输,并且集群可提供稳定主存带宽。
模型结构同样影响区间。更多 layers 增大 pipeline parallel 空间,更多 attention heads 放宽 Ulysses 分区,GQA/MQA 会减少 K/V activation,MoE 的路由不均衡会破坏固定段 timing。SPPO 的框架仍可用,但论文现成参数不可直接移植。所谓 adaptive 是运行前基于 profile 搜配置,不是对任意抖动实时无成本重规划。
14. 方法贡献
SPPO 把长 sequence 设为 offload 与 pipeline 的共同调度单位。二级管理利用 causal attention 的 K/V 复用特性区分 tensor lifetime;sequence-aware ratio 让传输量服从可遮蔽的计算时间;multiplexing 只重切 pipeline 边界,避免所有阶段长期承担细分段的 kernel overhead。这套方法联合设计了 memory hierarchy 与 parallel schedule,范围超过单独把 activation 搬到 CPU。
15. 局限与适用边界
- 结果证明 4M token 配置可运行并有吞吐,不证明模型能从 4M 上下文中学到有效依赖或达到特定下游质量。
- 主要评测是 GPT 架构、Ampere 80GB、32GB/s CPU-GPU 通道。更慢 PCIe、更小主存或更强 GPU 会改变 compute/transfer overlap。
- 它依赖可预测的 Transformer activation lifetime;MoE、跨模态分支或不规则算子可能需要新的 cost model。
- 短序列或小模型计算不足时,adaptive offload 的收益有限,甚至全卸载会变慢。
- 论文主要报告吞吐和可支持长度,没有展示不同 offload 策略对故障恢复、CPU 内存争用或多租户干扰的影响。
16. 与合集其他论文的关系
Rethinking KV Cache Compression 处理 inference 的 KV memory,SPPO 处理 training 的 activation memory:前者提醒压缩 kernel 可能拉长端到端响应,后者显示 offload 也必须算传输与重算成本。Di-PS 在多个异构集群之间做外层参数同步,SPPO 在单个训练作业内排 sequence/pipeline,二者可上下叠加。ICEFROG 决定多个训练 job 如何获卡及冻结多少层,不替代 SPPO 的作业内执行。Hydro 用大模型 pipeline bubble 跑 HPO trial,而 SPPO 是主动重排本 job 来消掉 bubble。
17. 五个容易误解的细节
为什么 K/V 值得常驻,其他 activation 却适合搬走
一个 subsequence 的 Q 在本段 attention 完成后就失去前向用途,但早期 K/V 会被所有后续 query 段反复访问。若把 K/V 一起卸载,后续每段都要重新预取,原本一次 D2H 会变成来回搬运,且恰好落在 attention 的关键路径上。其他中间 tensor 主要等待 backward,生命周期虽长、访问频率却低,更适合利用 forward 的空档搬到主存。这不是按 tensor 大小粗暴排序,而是按复用次数和下一次访问时间分类。
为什么分段越细不一定越好
更多 subsequence 能让 pipeline 更连续,也给传输更多重叠机会,但每段太短会增加 kernel launch、边界管理与通信次数,单段计算也不足以遮住固定传输延迟。论文把每层每卡的有效工作量限制在 2K 到 16K token,就是在表达这一折中。multiplexing 只细切 bubble 邻近段,而不把 steady phase 全部切碎,原因也是避免为了填少数空档而让整个迭代承担细粒度开销。
为什么先求并行配置,再求卸载比例
SP、PP、N 决定每段在哪张卡计算、计算持续多久以及前后依赖,卸载可获得的遮蔽窗口由这些选择产生。若先固定每段卸载比例,再改 pipeline 配置,原本能被遮住的 D2H 可能暴露在关键路径上,显存峰值也会移动。SPPO 的流程先得到候选执行骨架,再按相邻段 timing 分配 alpha,最后才对剩余 bubble 做局部重切分,顺序本身就是设计的一部分。
吞吐提升和资源节省要分开看
SPPO 的 offload 允许同样 GPU 数支持更长序列,也可能在固定长度下减少 recomputation、提高吞吐。这两种收益不是一回事。某配置原来 OOM、现在能运行,证明容量扩展;某配置两边都能运行而 tokens/GPU/s 更高,才证明速度提升。GPT-13B 和 65B 的最大长度对比主要是容量证据,600K/640K 的 3.38x/3.12x 才是固定工作负载的性能证据。阅读时应避免用一个倍数同时概括两者。
落地前应该重新 profile 什么
至少要测每类 activation 的大小和访问时刻、实际 pinned-memory D2H 带宽、NUMA 绑定是否正确、每种分段长度的 attention/MLP 时间、跨节点 collective 与 P2P 拓扑,以及 backward prefetch 是否与其他通信争链路。若主存容量不足、多个 job 同时 offload 或 CPU 内存带宽共享,论文物理集群中的近零暴露传输不再成立。此时正确做法是让 cost model重新搜索,而不是照抄论文的分段数。
18. 可复现性与结论边界
从一次 forward 看 activation 如何流动
以第一个 pipeline stage 为例,归一化后的输入进入 Q/K/V projection,产生本段的三组表示;FlashAttention 消费本段 Q 与截至当前的全部 K/V,随后 attention output 经 projection、residual,再进入 MLP。forward 刚结束时,看似所有 tensor 都还能 offload,但用途不同:早期段 K/V 会被下一段 attention 立即读取,残差和 MLP 中间值则要等对应 backward。SPPO 的 hook 在 tensor 产生时记录所属层、段和类型,把 K/V 注册到 GPU resident pool,把后者放进可卸载 unit。copy stream 只处理完整 unit,避免一个 autograd 所需 tensor一半在主存、一半已被释放。
进入下一段时,计算 stream 读取 resident K/V,同时 copy stream 搬上一段候选。若 copy 提前结束,旧 GPU buffer 可回收给 allocator;若未结束,recurrence 中的 M_i 会保留这部分占用。到 backward 反向遍历时,runtime 根据预先计划在更早的段计算期间发起 H2D,event 到达后 autograd 才继续。这样 offload/prefetch 都是 schedule 的一部分,而不是运行到 OOM 后临时换页。
pipeline bubble 的来源与 multiplexing 的正确性
普通 fill-drain pipeline 中,stage 0 先启动,stage 1 必须等第一份中间 activation;最后 stage 0 又先结束,后续 stage 仍有工作。subsequence 数越多,稳定态占比越高,但序列段之间存在 causal attention 依赖,不能像独立 micro-batches 一样任意交换。multiplexing 没有让未来段提前越过依赖,而是把边界段内部的 sequence parallel 范围扩大:原本空闲的设备共同处理同一合法段的不同 token/head 分片,算完再进入原 pipeline。
Table 3 中左侧 SP range 从 stage 0 的四卡逐级缩小,右侧则反向扩大,正好利用 pipeline 两端空闲设备。steady phase 每个 stage 仍持固定层并处理连续段,所以参数不必迁移。正确性来自每个被重切分段仍执行完整 attention/MLP,collective 合并后输出与未切分相同;变化只有在哪些设备并行完成。代价是边界需要额外 layout/collective,因此 solver 只在估计 bubble 足够大时启用。
为什么 65B 比 7B 更容易得到大幅加速
在相同序列量级下,65B 每 token 的 projection、MLP 和多层计算更重,checkpointing 重算成本更高,也提供更长的 copy overlap 窗口。adaptive offload 可以搬更多 activation 而不暴露 D2H;取消重算节省的 FLOPs 也更多。7B 的单段计算较短,Megatron 的重算虽然存在,却可能仍比复杂 offload schedule 便宜,尤其序列还没有把显存逼到极限时。论文在 GPT-7B 某些点只有边际加速并非失败,而是 cost model 预期的 crossover。
模型更大也不保证永远受益。如果参数/optimizer state 已吃掉绝大部分显存,resident K/V 和 copy buffer 空间不足,必须增大 PP 或使用 ZeRO;如果 CPU buffer 装不下 TB 级 activation,offload 也会失败。SPPO 解决的是 activation 路径,仍需与参数并行共同配置。Table 4 同时列出 SP/PP/N,就是为了说明端到端可行性来自组合,而非单一 offload knob。
与 recomputation 可以怎样组合
论文对 baseline 开 checkpointing,但 SPPO 的核心路径强调近零开销 offload、避免昂贵重算。工程上二者并不互斥:对尺寸很大但重算便宜的 activation 可丢弃,对计算昂贵且下一次访问晚的 tensor 可 offload,对高频 K/V 常驻。最优决策取决于“保存字节数、D2H/H2D 时间、重算 FLOPs、可遮蔽窗口”四项。SPPO 已实现前两类 lifetime 管理,却没有给任意 tensor 做三选一全局优化,这也是可扩展方向。
若加入 recomputation,性能模型必须避免重复计算收益:被丢弃 tensor 不占 offload bytes,也不能在 backward 前等待 H2D;重算可能与其他 stage 通信重叠,但会占计算 stream,从而缩短原本遮蔽 copy 的窗口。简单在 SPPO 上再打开更激进 checkpointing,可能让 alpha 计划失效。正确做法是重新 profile 并联合搜索,而不是把两个各自最优配置相加。
训练一 billion tokens 的 wall-clock 应如何估算
论文关心的不只是单 step,而是训练固定 token 总量。若每 iteration 实际处理 global_batch × sequence_length tokens,所需 iteration 为总 tokens 除以该值,wall-clock 约为 iteration 数乘稳定态 iteration time,再加初始化、checkpoint 和故障恢复。提高 sequence length 在 global token batch 不变时未必减少优化步;若同时改变 batch,则收敛超参数也可能变化。系统论文的吞吐比较通常固定等价工作量,不能把“一个 step 含更多 token”误当同等训练目标下必然更快。
SPPO 的 tokens/GPU/s 便于归一化不同卡数,但实际作业还要乘 GPU 数得到 aggregate throughput,并核算 CPU 主存与服务器成本。用 128 张卡跑 4M 的意义是相对传统估算减少所需 GPU,不表示总成本只有短序列训练的某个固定比例。若超长样本需要特别的数据预处理、padding 或 document packing,输入 pipeline 也可能成为新的瓶颈,论文主实验未把这些生产因素全部展开。
OOM 边界为何是离散而不是平滑曲线
某长度能否运行受并行 degree、allocator block、kernel workspace 和 activation 峰值共同影响。增加少量 token 可能跨过一个 buffer 分配阈值而立即 OOM;改变 PP 又会让每 stage 层数离散变化。因此 Table 4 的最大支持长度是测试配置下的可行点,不是精确物理上限。复现时应报告最后成功点、首个失败点以及失败发生阶段,避免只给一个整数让人误以为硬件普遍上限。
SPPO 的近线性长度扩展来自把新增 activation 主要放到主存并随更多 GPU 分摊计算,但 host memory 也按节点离散增加。128 GPU 同时意味着 16 节点、约 32TB 总 CPU memory,不能说只靠 10TB GPU 显存实现。系统节省的是昂贵 GPU 容量,并利用已有主存层级;总内存数据仍然存在。
数值一致性、随机性和验证方法
offload 本身只复制 FP16/BF16 tensor,不做有损压缩,理论上回传后 bitwise 相同。但不同并行分组可能改变 collective reduction 顺序,浮点非结合性会产生微小差异;异步 stream 若缺 event 则会出现真正的数据竞争。验证可分三层:单层固定输入对比 forward/backward tensor;小模型数十步对比 loss 与 gradient norm;目标规模长跑观察吞吐、显存峰值和 loss trend。要求大规模 bitwise 全同并不现实,但偏差应与并行 baseline 的正常数值差同量级。
multiplexing 的 SP group 动态变化更需要 collective trace。每个 rank 必须以相同顺序加入相同通信组,边界段的 all-gather/reduce-scatter shape 要与未切分输出一致。任何偶发 hang 或 loss spike 都应先检查 schedule/event,而不是归因于长上下文本身。论文展示系统能训练,不等于所有移植实现自动继承正确性。
面对多租户集群时还缺什么
论文物理集群把节点资源用于一个训练 workload,因而 CPU memory bandwidth、PCIe、InfiniBand 基本可预测。多租户环境中,另一个 job 的 data loader、checkpoint 或 offload 会让带宽瞬时下降,静态 alpha 暴露传输;CPU 内存总量也需 scheduler 预留。要生产化,可把 host-memory bytes、copy bandwidth 和 NUMA locality 作为一等资源,由集群调度器做 admission,而不是只按 GPU 数放置。
运行时还可监控 copy completion 与计算 event 的时间差:连续出现正的 exposed-copy time 时降低 offload ratio、开启部分 recomputation或减少并发;CPU 有余量则反向调整。这样的闭环不在论文核心实现中,但遵循其 cost model。没有这层协调时,把多个“各自在独占测试中近零开销”的 SPPO 作业共置,可能同时失去吞吐。
这篇论文到底改变了什么抽象
传统训练框架把 sequence 当一个 tensor、activation manager 按 layer 管理、pipeline scheduler 按 micro-batch 管理,三者边界固定。SPPO 把 subsequence 变成贯穿三层的对象:它既有 causal 依赖和计算量,又有具体 activation lifetime,还能在 pipeline 中排程。这样 solver 才能用“下一段计算”决定“上一段搬多少”,再用空闲设备重切“边界段”。论文的贡献是统一控制粒度;单个 copy 优化无法产生同样的数据流。
这个抽象也解释其适用边界:只有当序列可合法分段、段间依赖可描述、tensor lifetime 可预测时,联合优化才成立。对于数据相关控制流或完全不规则稀疏模型,需要更动态的图执行器。对标准 decoder-only GPT,结构规则且迭代重复,前期 profile 成本能在长训练中摊薄,因此论文取得显著收益。
把论文的三个 challenge 与三个解法一一对齐
第一个 challenge 是固定 offload policy:按长度或 FLOPs 给所有段相同比例,会让传输与相邻计算不匹配。对应解法是 sequence-aware alpha_i。第二个 challenge 是同一段内部 tensor 的复用差异:把 K/V 与 backward-only activation 一起搬,会产生无谓回传。对应解法是 two-level activation management。第三个 challenge 是为了 kernel 效率不能无限增加 N,所以 pipeline 总有剩余边界 bubble。对应解法是 multiplexing sequence partition,只重新利用边界设备。三者分别解决“段间比例”“段内对象”“流水线边界”,不是三个同义优化。
消融也按这条链读:full offload 只证明容量可以转移到主存,可能比不搬更慢;adaptive offload 同时加入段内分类和段间比例,才稳定提高吞吐;multiplexing 在此前 schedule 上进一步填空。若实现只复刻最后的边界切分而没有前两步,GPU 可能在算更多小块的同时仍等待 D2H,得不到论文结果。
一个简化的手算容量例子
假设某 stage 可用于 activation 的 GPU 余量为 100GB,连续两段的 backward-only activation 分别为 80GB 和 120GB,常驻 K/V 共 20GB。若完全不 offload,第二段产生时占用至少 20+80+120=220GB,必然 OOM。若第一段可用的下一段计算窗口只能遮蔽 48GB D2H,令 alpha0=0.6,它仍留 32GB;第二段生成时为 20+32+120=172GB,仍不可行,所以 solver 必须同时给第二段更早分批释放、增大 PP/SP 分摊或降低段大小,不能仅调一个比例。
再假设重新分段后两段为 50GB、70GB,传输阈值 40GB,则 alpha0=0.8、alpha1≈0.57,符合小段比例更高。第一段搬完后只留 10GB,第二段出现时加常驻 K/V 为 100GB,刚好可行;每次 D2H 都是 40GB,也更容易由相邻 balanced compute 遮蔽。这个数值是解释公式的示意,不是论文测量配置,但展示 solver 为何必须联合分段、比例和峰值约束。
从研究原型走向训练平台的接口
用户层应只声明模型、序列长度、global batch 与允许使用的 CPU memory,规划器输出并行配置和预计 tokens/s;模型代码不应手工写每个 activation 的 offload。框架层需要统一 tensor lifetime metadata、pinned buffer pool、prefetch API 与 topology-aware process groups。监控层应暴露 GPU reserved/allocated memory、host buffer bytes、D2H/H2D overlap ratio、pipeline idle ratio、每段时间和重新规划原因。这样出现性能回退时能区分是模型计算变化、带宽争用还是计划错误。
checkpoint 也必须包含足以恢复 schedule 的配置,但无需保存 activation;故障后从最近 optimizer checkpoint 重启迭代并重新 profile/规划即可。若弹性缩扩 GPU,原 SP/PP/N 通常失效,不能只改变 world size。SPPO 本身不是弹性训练系统,却给出一个清楚的再规划边界:设备集合、模型结构、序列长度或有效带宽显著变化时重求计划。
最后用反事实检查论文结论
如果 GPU 显存足以保存全部 activation,SPPO 仍能运行,但最优计划应把 alpha 降到零,退化为没有 offload 的并行训练;若规划器仍大量搬运,说明 cost model 有误。如果 CPU-GPU 带宽趋近无限,offload 几乎免费,最优计划会更积极搬走 backward-only tensor,却仍保留高频 K/V,因为无谓往返会占 copy engine 和同步点。如果 pipeline stage 只有一个,PP bubble 消失,multiplexing 应自动关闭,但 sequence-aware offload 仍有用。这三个极端说明各组件并非必须同时产生收益。
反过来,若只有很慢的主存链路且重算便宜,最优系统应回到 checkpointing;若 sequence 无 causal 结构、各段独立,K/V 常驻规则与非均匀分段都要改写。论文并没有声称 SPPO 支配所有 memory technique,而是在超长 decoder 训练这一明确区域给出比固定策略更好的联合选择。判断新工作负载是否适用,应先验证 activation lifetime、计算/传输比和 bubble 三个前提,而不是先寻找相同模型名字。
还要注意,论文所称“modest GPU resources”是相对传统估算:128 张 A100 依然是大型集群,只是无需为 4M 序列堆到上千张高端卡。CPU 主存、网络和运维成本仍真实存在。系统贡献是把原本只能靠昂贵 GPU 容量承担的状态放入更便宜的内存层级,并用调度尽量不让速度为此付费,而不是把超长训练变成单机任务。这个措辞边界对评估成本尤其重要。
因此,评价复现是否成功至少要同时满足三点:目标长度不发生 OOM;loss 与未修改训练语义一致;稳定态吞吐达到按硬件重新调优后的预期。只满足容量而速度很慢,可能是带宽或 overlap 配置问题;只看到短跑吞吐而没有数值校验,也不能证明训练系统正确。三项缺一,都不足以支持论文的完整主张。
复现应先记录 GPU 型号与显存、CPU socket/内存容量、PCIe 代际、实测双向 copy 带宽、NVLink/InfiniBand 拓扑和软件 commit;再固定模型层数、hidden、heads、精度、checkpoint policy、global batch 与 sequence length。比较时必须给每个 baseline 单独调 SP/PP 等并行参数,论文 Table 4 就是在避免用同一配置偏袒 SPPO。吞吐至少预热若干 iteration,排除首次分配与编译,并报告 tokens/GPU/s 与整机 tokens/s。
正确的最终结论是:在论文硬件和 GPT 工作负载上,subsequence-aware activation placement 与边界 multiplexing 能把超长序列训练的主存传输和 pipeline 空档大幅隐藏,最高带来 3.38x 相对强 baseline 吞吐,并把 7B 扩展到 4M token。它没有证明任何 4M 训练任务都值得做,也没有消除主存、互连和模型结构对收益的约束。把这两句话同时保留,才是无需回原文也不会误读的结论。
换到新集群时,先重新测带宽、显存和段级时间,再重新求计划;直接照搬论文参数不构成有效复现。