AI Training Parallelism

导言

AI 训练时，有些分布式训练的常见并行概念需要了解。
例如，TP, VP, SP, VPP

数据并行 DP¶

(Data Parallelism)

思想：把同一个模型放在多个GPU上，batch数据平均分布到各个GPU上，并行计算。
难点：注意参数同步和信息过期问题。
优点：
- 加速比线性。
- 部署简单工作量小，每个节点内的计算效率高。
- 由于部署简单，是最先采用的并行方式。
缺点：
- 需要在每个节点复制所有模型参数，显存重复度高，利用率低，并不适合大模型的部署。只适用于训练样本较多而模型较小的情况。
- 数据分布在不同机器，需要allreduce同步权重 FS（fully shared）DP

特点+具体操作：

同构模型，不同数据：每个节点都包含完整的模型，以及模型的参数（weight，parameter），输入数据则根据模型的并行度进行拆分，分别被每个节点读取；
- 假设我们有8张GPU卡或者昇腾的NPU卡来训练图片分类的模型，训练的批量为160，那么每张卡上面分到的批量数据（min-batch）为20，每张卡基于样本数据完成训练。
独立运行：每个节点读取相应的输入数据后，分别独自处理模型的前向和反向传播，并得到Gradients，归并所有的梯度并更新梯度；
梯度聚合：因为各张卡上处理的数据样本不同，所以获得的梯度会有些差别。因此，需要对梯度进行聚合（求和、均值）等计算来保持和单卡训练相同的结果，最后再更新参数。
统一通讯：所有节点之间的通信，主要包括前向传播的Loss归并以及反向传播的gradient归并以及更新，这些通信则是通过相应的通信原语（gather/reduce/broadcast）操作。
参数更新：梯度聚合会让各卡的模型以相同的梯度值同时进入参数更新阶段，然后针对新的数据进行下一轮训练。

¹

模型切分的共同背景

模型并行把模型或计算图拆到多个设备，适用于单卡放不下完整模型的情况。它需要人为设计切分边界并支付通信成本；张量并行、流水线并行与专家并行只是不同切法，不需要再用一个空泛的“模型并行”标题把它们强行包在一起。

张量并行 TP¶

Tensor Parallelism 是模型并行的一种

思想：通过在多个计算设备上分片模型参数来实现并行。
举例：一个操作中进行并行计算，主要是矩阵-矩阵乘法。张量并行训练是将一个张量沿特定维度分成 N 块，每个设备只持有整个张量的 1/N，同时不影响计算图的正确性。这需要额外的通信来确保结果的正确性。
底层逻辑是矩阵乘法的拆分计算, 矩阵乘法中列并行与行并行这两种张量并行的方式，以及它们在前向传播和反向传播中的区别。¹⁴
难点：额外的通信
优点：每个设备只持有整个张量的 1/N

张量并行要具体场景具体设计

Megatron-LM模型在训练中采用更为手工的方式，将每个transformer块都进行了分割，实现了高性能的计算。²

LLM¹¹

序列并行 SP¶

Sequence Parallelism 不是一种单独算法

SP 只说明输入或激活沿序列维切分，没有说明 Attention 如何取得其他分片的上下文。真正决定显存、通信、头数限制和迁移成本的是 Attention 内部选择 AllGather、All-to-All、P2P Ring，还是它们的二维组合。

两条技术路线¶

同一个名称下实际有两条不同路线：

TP 配套的 Megatron SP：只把 LayerNorm、Dropout 等逐 token 激活沿序列维切开，Attention 和 MLP 仍由 TP 负责。它的目标是消除 TP 区域之间的冗余激活，不能单独解决长上下文 Attention。⁸
长上下文 SP/CP：网络输入和各层激活始终保持 N/P 的序列分片；Attention 再通过 AllGather、All-to-All 或 Ring 取得全局上下文。Ulysses、Ring Attention、Megatron Context Parallelism（CP）和 USP 都属于这条路线。

三个容易混淆的名称

Megatron SP 是 TP 配套的激活优化。
Megatron CP 是贯穿所有层的长上下文切分，Attention 默认以 Ring 交换 KV。
USP 在原论文中指 Unified Sequence Parallelism，即 Ulysses × Ring 的二维组合；部分框架也把 Ulysses Sequence Parallel 简写成 USP，阅读配置时需要回到具体通信原语。

长上下文序列并行方案选择示意图 — 自绘示意图：AllGather、Ulysses、Ring 与 USP 的核心区别分别是整段复制、按头换轴、KV 分块环传和二维组合。

演进脉络¶

方案	出现时间	当时的设计初衷	核心数据移动	论文或官方效果
Ring Self-Attention（早期 SP）	2021-05	单卡不再持有完整序列，突破自注意力的长序列显存上限	固定本地 Q，K/V 分块沿 Ring 传递	论文在 64 张 P100 上相对 TP 支持 `13.7×` 最大 batch、`3.0×` 最大序列；稀疏 Attention 超过 114K token。⁵
Megatron SP	2022-05	减少 TP 区域之间 LayerNorm、Dropout 激活的重复存储，少做全量重计算	TP 的 AllReduce 拆成 ReduceScatter + AllGather	与选择性重计算合用时，论文报告激活显存减少 `5×`；530B GPT 训练达到 54.2% MFU，比全量重计算的 42.1% 快 29%。这些数字不是 SP 单项消融。⁸
DeepSpeed-Ulysses	2023-09	用高带宽集合通信替代低效 P2P，把长序列 Attention 包装成可复用模块	Attention 前后各一次 All-to-All，在序列维与头维之间换轴	官方实验报告比既有基线支持 `4×` 更长序列、通信量降低超过 `10×`、吞吐最高 `2.5×`，并训练百万 token 序列。¹⁰²³
Ring Attention	2023-10	用 blockwise Attention 把 KV 通信隐藏在计算后面，并让可处理长度随设备数扩展	每轮只收发一个 KV block，在线合并 softmax 统计量	论文支持百万 token，并给出最多随设备数线性扩展的上下文长度；“无额外暴露通信”依赖 block 计算足以覆盖 P2P。⁹
Megatron CP	2023 Q3；MCore 0.5.0+	把 Ring Attention 工程化到 Megatron/Transformer Engine，兼顾因果负载均衡与现有 Flash/Fused Attention	默认 P2P Ring；当前也支持 AllGather、All-to-All 和分层 A2A+P2P	官方文档给出每卡激活显存约按 CP 度缩小，并在 175B GPT 示例中优于全量重计算；页面未给出可脱离图表复述的统一加速数字。²⁰²⁴
USP（Unified SP）	2024-05	不在 Ulysses 与 Ring 间二选一，而是适配头数和节点内/节点间拓扑	二维进程网格：一维 All-to-All，另一维 P2P Ring	论文在 2×8 A800 上训练 LLaMA3-8B：208K 序列达到 147.26 TFLOPS/GPU、47% MFU；64K/80K 时最佳二维配置比纯 Ring 高 13%/12%。¹⁸
Ulysses-Offload / FPDT	2024-08 论文；2026-07 教程	Ulysses 已切分 rank 间序列，但单个 rank 内一次性 Q/K/V、Attention 反向状态和非 Attention 算子仍限制超长上下文；目标是把这些峰值变成可流水的小工作集	Attention 按全局 `C` 分块；A2A 后把暂时不用的 Q/K/V 与反向状态卸载到 pinned CPU，并用双缓冲覆盖传输；FFN 与 loss 也分块	FPDT 论文报告同硬件最长序列最高提升 `16×`；8B 模型可在 4 卡训练 2M token，MFU 超过 55%。这些是特定 A100/网络与调参结果。³⁷³⁸
AllGather CP	2024-10 进入 MCore `cp_comm_type`	为 CP 提供最直接、最容易复用标准 Attention kernel 的通信后端	前向 AllGather 完整 K/V；反向 ReduceScatter K/V 梯度	它不是独立论文方案，MCore 明确标注当前 AllGather 不能异步、不能与 Attention 计算重叠；没有独立的统一性能数字。²¹²²
ALST（Arctic Long Sequence Training）	2025-06 论文；2026-07 教程	发现只有 Ulysses 仍会在 logits、MLP、checkpoint activation 等位置 OOM，于是把跨卡 SP、逐 token tiling、CPU offload 和运行时优化组合起来	Attention 仍是 Ulysses A2A；MLP/loss 沿序列分 Tile；checkpoint activation 可卸载到 CPU	Llama-8B 在论文的 H100 + ZeRO-3 配置中从 32K 扩展到单卡 500K、8 卡 3.7M、32 卡 15M；追求的是最长可训练序列，不是最短 step time。³¹³²
MindSpeed-MM Chunk Layer	2026-07-25 分支实现	在 FSDP2 单个 rank 内把完整 Decoder Layer 改成序列分块执行，以降低当前 Chunk 的激活峰值，并让 FA/GDN 跨块状态可重计算、可卸载	没有 SP 通信；FA 串行携带因果前缀 K/V，GDN 串行携带卷积与 delta state；第二版把所选缓存的 D2H 从重复累计前缀改为每个值只搬一次	最新提交没有测试或性能基准；源码只证明 BSND/TND、FA/GDN 路径和 `O(S)` cache D2H 设计，不能据此宣称无限长度或统一加速。⁴⁸⁴⁹

效果数字的边界

上表来自不同年份、模型、GPU、网络、精度和 baseline，只能证明各方案在原实验中的可行性，不能横向排名。同一集群上的最终选择应固定模型、序列长度、CP 度和重计算策略，比较峰值显存、step time、MFU 与通信暴露时间。

下面统一用这些符号：B 是 batch size，S 是全局序列长度，P 是序列或上下文并行度，Hq/Hkv 是 Q/KV 头数，D 是单头维度，M 是隐藏维度。为便于看懂，伪代码只保留决定数据布局和通信语义的步骤，但不会用省略号隐藏关键的数据移动。

技术图怎么读

既有方案使用自绘四联图：设计逻辑、前向流程、多 Rank 时序、张量数据流。MindSpeed-MM Chunk Layer 不是跨 Rank 算法，因此改用五视图：再增加一幅物理形态前后对照，并把 Rank 时序替换成 NPU、staging 与 pinned CPU 的组件生命周期。蓝色表示本地张量流，橙色表示跨 Rank 通信，紫色表示循环、归并或控制；论文原图紧随相应论文方案之后，效果图只说明原实验，不能跨论文直接排名。

反向、重计算与卸载¶

前向只要沿箭头走一遍，确实比较直观；反向最容易产生的误解是：“PyTorch 有 autograd，所以 Ring 梯度通信和 FPDT 的分块调度也是系统自动推导出来的。” 实际上，autograd 更像一个按依赖倒序叫号的调度员，它不会凭空发明分布式算法：

普通 PyTorch 算子已经登记了梯度公式，模型作者调用 loss.backward() 后，autograd 自动串起这些公式。
自定义分布式算子必须由框架作者手写 backward()：包括梯度张量如何切分、该做 All-to-All 还是 Ring、累加发生在哪个 rank，以及返回给上游的布局。autograd 只负责在恰当时机调用它。⁴¹
activation checkpoint 必须显式圈定要重算的函数或层。反向到达这个边界后，checkpoint 包装器恢复输入和随机数状态、重新执行一次前向，再让 autograd 对新图求导；系统不会自行猜测最合适的重算边界。⁴²
activation offload 必须显式决定保存什么、搬到哪里、何时预取。通用 save_on_cpu 或 saved-tensor hook 可以自动执行登记好的 pack/unpack，但 FPDT 还有自己的 pinned buffer、CUDA stream 和预取次序。⁴³

层次	谁设计反向或内存策略	训练时看起来是否自动	实际需要显式实现的部分
Linear、LayerNorm 等普通算子	PyTorch/kernel 作者	是	算子的局部 VJP（vector-Jacobian product）已注册
Ulysses A2A	框架作者	是，autograd 自动调用	反向再做一次 A2A，并交换前向的 scatter/gather 维度
Ring / USP Attention	Ring kernel 作者	是，autograd 自动调用	FlashAttention block backward、`dQ` 本地累加、`dK/dV` 环传累加
activation checkpoint	模型/训练框架配置者	启用后自动	checkpoint 边界、输入保存、RNG 恢复、重跑哪些通信
activation offload	框架或策略配置者	登记策略后自动搬运	保存对象、CPU buffer、预取距离、stream 与生命周期
FPDT Attention/FFN/loss	DeepSpeed FPDT 作者	是，autograd 自动调用	分块循环、A2A 反向、FlashAttention backward、权重梯度累加、卸载与预取

序列并行反向、重计算与激活卸载的责任和数据流 — 自绘四联图。左上区分“autograd 自动调度”与“框架作者手写 VJP”；右上展开 USP 的逆序反向；左下展示 checkpoint 与 offload 的组合时序；右下展示 FPDT 的 CPU/GPU 梯度流水。

先分清三种内存策略¶

假设一段前向为 Y=f(X)，反向需要 X、中间激活 A 和上游梯度 dY。三种方案解决的是同一显存问题，但付出的代价不同：

普通保存：前向把 A 留在 GPU，反向直接读取。速度最快，HBM 占用最高。
重计算：前向只保留边界 X；反向先重新执行 f(X) 得到 A，再计算梯度。省 HBM，增加一次前向计算；若 f 内含 A2A/Ring，通信也会重放。
卸载：前向把 A 搬到 CPU，反向再预取回来。它不重新计算 A，但增加 D2H/H2D、pinned RAM 和同步成本。
重计算 + 卸载：checkpoint 的边界输入可以在 CPU；反向先 H2D 取回 X，再重跑 f(X)，最后执行自定义 backward。FPDT 还可能在这次重跑的前向里，再次建立自己的 Q/K/V/O/LSE CPU 状态。

卸载不是重计算

从 CPU 取回旧值叫 reload，从输入重新算出新值才叫 recompute。 两者可以同时存在。只看显存曲线时它们都像“激活消失后又回来”，但性能瓶颈分别是 PCIe 带宽与额外 FLOPs/通信，调优方式完全不同。

下面是一个不绑定具体框架的执行骨架。训练脚本没有手改 .grad；checkpoint 与自定义算子的 backward 返回梯度后，autograd 继续沿图向前一层传播：

def train_step(batch):
    hidden = checkpoint(transformer_block, batch.hidden)
    loss = language_model_loss(hidden, batch.labels)
    loss.backward()                                # 应用层通常只调用这一句

def checkpoint_backward(boundary_x_cpu, grad_y, rng_snapshot):
    boundary_x = copy_to_gpu(boundary_x_cpu)
    restore_rng_state(rng_snapshot)
    with enable_grad():
        recomputed_y = transformer_block(boundary_x)
        # 若 block 内有 A2A/Ring/FPDT，这里的前向通信会再执行一遍
    autograd_backward(recomputed_y, grad_y)
    return boundary_x.grad

def custom_distributed_backward(grad_y, saved_or_recomputed_state):
    grad_x, grad_weight = hand_written_vjp(
        grad_y, saved_or_recomputed_state
    )
    return grad_x, grad_weight

DeepSpeed-Ulysses 的 A2A 是最小例子：前向记录 scatter_idx 与 gather_idx，反向再次调用同一个 A2A Function，但把两个维度对调。也就是说，梯度调用是自动的，逆布局变换是源码作者手写的。把源码中的 batch 维、stream/handle 和不可导参数返回占位折叠后，关键逻辑可以归一化为：²⁶

@staticmethod
def backward(ctx, grad_output):
    return _SeqAllToAll.apply(
        ctx.group, grad_output, ctx.gather_idx, ctx.scatter_idx
    )

重计算还有两个分布式约束。第一，所有 rank 必须以相同顺序、相同次数重新进入 collective，否则不是变慢，而是直接死锁；第二，Dropout 等随机算子的 RNG 状态必须与原前向一致，否则重算得到的不是同一个函数。当前 PyTorch checkpoint 默认保存相关 RNG 状态，并推荐显式使用非 reentrant 实现，但第三方框架仍要正确管理 TP/SP 专属 RNG tracker。⁴²⁴⁵

Ring Self-Attention¶

方案一：先环传 K 算分数，再环传 V 算输出。

起因：2021 年的方案面对的是最直接的长序列问题：若每张卡都保存完整 Q/K/V 和 S×S 分数矩阵，显存很快耗尽；只做 TP 又没有切掉序列长度。
小白直觉：每张卡把自己的 Q 留在原地，让其他卡的 K 依次“排队来做题”；所有题目的分数归一化后，再让 V 按相同顺序走一圈，累加出本地 Q 的输出。
关键对象：本地 Q_i/K_i/V_i 为 [B,S/P,H,D]；每轮分数块 score_ij 为 [B,H,S/P,S/P]；O_i 为 [B,S/P,H,D]。与现代 Ring Attention 不同，早期方案要先得到并保存分数块，再开始 V 阶段。

Ring Self-Attention 逻辑、流程、时序与数据流 — 自绘四联图。重点看左下时序：K 和 V 分成两个完整的 Ring 阶段，因此通信轮数多，分数块生命周期也更长。

# 教学伪代码：每个 rank i 独立执行，ring_next 表示把块发给下一个 rank
q_i, k_block, v_block = project_qkv(x_i)       # [B,S/P,H,D]
score_blocks = []

for step in range(P):
    owner = (i - step) % P
    score_ij = matmul(q_i, transpose(k_block)) / sqrt(D)
    score_ij = score_ij + causal_mask(i, owner, S, P)
    score_blocks.append((owner, score_ij))
    k_block = ring_send_recv(k_block)

scores = concat_by_owner(score_blocks, dim="key_sequence")
probs = softmax(scores, dim="key_sequence")
prob_blocks = split_by_owner(probs, parts=P)
o_i = zeros_like(q_i)

for step in range(P):
    owner = (i - step) % P
    o_i = o_i + matmul(prob_blocks[owner], v_block)
    v_block = ring_send_recv(v_block)

return output_projection(o_i)

Ring Self-Attention 论文方法原图 — 论文 Figure 2 原图：上半部分传 K 计算 Attention score，下半部分再传 V 计算输出；这正是它与 2023 年 online-softmax Ring Attention 的核心差别。[^5]

Ring Self-Attention 最大序列扩展效果 — 论文 Figure 5 原图：左图展示 SP 相比 TP 的可训练长度扩展，右图展示稀疏 Attention 下接近理想的长度扩展。数值只对论文的 BERT/P100 实验成立。[^5]

效果与局限：它证明了“Q 留在本地、K/V 环传”可以突破单卡序列上限，但两阶段通信和分数块存储并不理想。今天更常见的 Ring Attention 用 blockwise/FlashAttention 与 online softmax 边算边归并，通常不再照搬这份实现。

Megatron SP¶

方案二：只切 TP 外围的逐 token 激活。

起因：Megatron 的 TP 已经切分了 Attention/MLP 权重，但 LayerNorm、Dropout 和残差仍会在每个 TP rank 保存同一份完整序列激活，浪费显存。
小白直觉：像多人合写一本书。真正需要多人共同计算的 Attention/MLP 前，把各自页码临时拼成整本；算完马上把结果按页码散回去。Norm、Dropout、残差只看单个 token，所以各人处理自己的页即可。
关键边界：本地 X_i/N_i/Y_i 是 [S/P,B,M]；进入 TP 区域前 AllGather 成 [S,B,M]；TP 线性层生成的部分和再由 ReduceScatter 归并为 [S/P,B,M]。它减少激活复制，没有让 Attention 获得超长上下文的能力。

Megatron SP 逻辑、流程、时序与数据流 — 自绘四联图。左上清楚标出完整序列只存在于 TP Attention/MLP 区域；右下展示 `AllGather → TP Linear → ReduceScatter` 的共轭关系。

# 一个 TP 子层的教学伪代码；真实 Transformer 会在 Attention 和 MLP 两处复用
residual_i = x_i                              # [S/P,B,M]
norm_i = layer_norm(x_i)                      # 逐 token，本地完成
norm_full = all_gather(norm_i, dim="sequence", group=tp_group)
                                                  # [S,B,M]
z_partial = tensor_parallel_sublayer(norm_full, weight_shard_i)
                                                  # [S,B,M] 的部分和
y_i = reduce_scatter(z_partial, dim="sequence", group=tp_group)
                                                  # [S/P,B,M]
y_i = residual_i + dropout_with_sp_rng(y_i)
return y_i

Megatron SP 论文方法原图 — 论文 Figure 5 原图：`g` 与 `ḡ` 分别是一对前后向共轭的 AllGather/ReduceScatter；虚线框说明哪些区域属于 SP、哪些仍属于 TP。[^8]

Megatron SP 激活显存效果 — 论文 Figure 7 原图：蓝柱是 SP，黄柱是 SP 与选择性重计算合用。约 `5×` 显存缩减属于**组合结果**，不能全部归因于 SP。[^8]

实现与迁移：Megatron 当前在 TP Linear 的映射函数中实现这对通信，固定代码版本可见 layers.py 的相关路径。²⁵ 在 MCore 内接入通常只是配置；迁移到新框架时，需要同时处理 TP Linear 边界、Dropout RNG 和残差布局，否则代码看似能跑，数值语义却可能改变。

DeepSpeed-Ulysses¶

方案三：用 All-to-All 在序列维和头维之间换轴。

起因：Ring 的多轮 P2P 代码复杂，且小块通信可能难以吃满高带宽网络。Ulysses 希望只在 Attention 前后各做一次集合通信，中间仍调用普通的本地 Attention kernel。
小白直觉：开始时每人拿“所有题型的一小段题目”；交换后，每人改成拿“少数题型的全部题目”。这样每个人都能独立完成自己那些 Attention heads，再把答案换回原来的序列分片。
布局变化：输入 Q_i 是 [B,S/P,Hq,D]；All-to-All 同时在序列维 gather、头维 scatter，得到 [B,S,Hq/P,D]。K/V 同理变成 [B,S,Hkv/P,D]；输出再逆变换回 [B,S/P,Hq,D]。

DeepSpeed-Ulysses 逻辑、流程、时序与数据流 — 自绘四联图。右上红字是最关键的硬约束：`P` 必须能整除参与切分的 Attention heads；对 GQA，通常还要受更少的 KV heads 约束。

q_i, k_i, v_i = project_qkv(x_i)              # [B,S/P,Hq|Hkv,D]

q_head = all_to_all(q_i,
                    scatter_dim="head", gather_dim="sequence",
                    group=ulysses_group)       # [B,S,Hq/P,D]
k_head = all_to_all(k_i,
                    scatter_dim="head", gather_dim="sequence",
                    group=ulysses_group)       # [B,S,Hkv/P,D]
v_head = all_to_all(v_i,
                    scatter_dim="head", gather_dim="sequence",
                    group=ulysses_group)       # [B,S,Hkv/P,D]

o_head = local_attention(q_head, k_head, v_head, global_mask)
                                                  # [B,S,Hq/P,D]
o_i = all_to_all(o_head,
                 scatter_dim="sequence", gather_dim="head",
                 group=ulysses_group)          # [B,S/P,Hq,D]
return output_projection(o_i)

DeepSpeed-Ulysses 论文方法原图 — 论文 Figure 2 原图：两根红色箭头分别是 Attention 前后的 All-to-All；蓝色区域内每个 rank 已拥有完整序列，但只拥有一部分 heads。[^10]

DeepSpeed-Ulysses 7B 模型吞吐效果 — 论文 Figure 4 原图：指定 7B GPT 实验中，Ulysses 从 8K 扩展到 256K，而对照组在 256K OOM。该图不是对所有硬件和模型的普遍保证。[^10]

实现与局限：DeepSpeed 把输入校验、两次 All-to-All 和本地 Attention 封装在 DistributedAttention 中。²⁶ 纯 Ulysses 的并行度不能超过可切分头数；MHA 通常看 Q heads，GQA 的 K/V 布局常把约束收紧到 KV heads，MQA 只有一个 KV head 时尤其棘手。packed sequence、cross-attention 或 TP 同时切头时，还要重新核对维度映射。

Ring Attention¶

方案四：KV 边走，Attention 边算，softmax 在线合并。

起因：早期 Ring Self-Attention 要先存分数、再传 V；Ulysses 又受头数整除限制。Ring Attention 希望每张卡始终只保存本地 Q 和一个 KV block，并把通信隐藏在块计算之后。
小白直觉：本地 Q 是考生，KV block 是依次经过考场的题本。每来一本就立即算这一部分答案，并保存“当前最大分数、归一化总和、加权输出”三个统计量；题本走完一圈，局部答案就等价于看过完整上下文。
关键对象：Q_i 为 [B,S/P,Hq,D]，当前 KV_j 为 [B,S/P,Hkv,D]；块 Attention 返回 O_ij 与 LSE_ij，online merge 把它们稳定地合并为 O_i/LSE_i，无需物化完整 S×S 分数矩阵。

Ring Attention 逻辑、流程、时序与数据流 — 自绘四联图。左下展示异步收发下一块 KV 与当前 FlashAttention 并行；右下展示真正保证全局 softmax 正确性的 `O/LSE` 在线归并。

q_i, k_block, v_block = project_qkv(x_i)
o_i = zeros_like(q_i)
lse_i = full(shape=[B, Hq, S/P], value=-inf)

for step in range(P):
    owner = (i - step) % P
    recv_request = async_ring_send_recv(k_block, v_block)
    block_mask = causal_mask(i, owner, S, P)
    o_block, lse_block = flash_attention_block(
        q_i, k_block, v_block, block_mask, return_lse=True
    )
    o_i, lse_i = online_softmax_merge(
        o_i, lse_i, o_block, lse_block
    )
    k_block, v_block = recv_request.wait()

return output_projection(o_i)

Ring Attention 论文方法原图 — 论文 Figure 2 原图：上半部分是设备间传递 KV block，下半部分把一次设备处理画成 Query 外循环与 KV 内循环，并显式标出“计算当前块、发送下一设备”。[^9]

Ring Attention 最大上下文长度效果 — 论文 Table 3 原图：在论文给定 A100/TPU 配置上，最大上下文长度随设备规模显著扩展；这证明可扩展性，不代表任意网络都能隐藏通信。[^9]

实现与局限：官方 JAX 实现把通信 stream、块 Attention 和 online merge 放在同一循环中。²⁷ 它不切 head，因此没有 Ulysses 的头数上限；但实现明显更难，因果 mask 的负载均衡、通信与计算是否真能重叠、反向传播的块调度，都会影响结果。块太小或跨节点网络太慢时，“通信免费”并不成立。

Megatron CP¶

方案五：把 Ring Attention 变成框架能力。

起因：论文级 Ring 算法还不等于生产训练。框架需要让每一层都保持序列分片，同时兼容 TP/PP/DP、RoPE、因果负载均衡、GQA/MQA、Flash/Fused kernel 和 checkpoint。
小白直觉：Megatron 负责“分组、切序列、摆好数据”，Transformer Engine（TE）负责“边收 KV 边算 Attention”。用户通常只选 CP 度和通信后端，不必在模型层手写 Ring 循环。
与 Megatron SP 的差别：CP 让输入、Norm、Attention 输出、MLP 和残差都保持 [S/P,B,M]；只有 Attention 为取得全局上下文而跨 CP rank 通信。Megatron SP 则会在 TP Attention/MLP 前临时恢复完整序列。

Megatron CP P2P 逻辑、流程、时序与数据流 — 自绘四联图。右上强调职责边界：Megatron 配置 `cp_group/cp_comm_type`，TE 的 DotProductAttention 执行 P2P KV 交换和 Attention kernel。

# 框架层伪代码：Ring 内核由 Transformer Engine/Attention backend 执行
x_i = shard_sequence_and_load_balance(x, cp_rank, cp_size)
q_i, k_i, v_i = local_qkv_projection(x_i)      # 均为本地序列分片

cp_group = build_context_parallel_group(cp_size)
o_i = transformer_engine_attention(
    q_i, k_i, v_i,
    cp_group=cp_group,
    cp_global_ranks=global_ranks(cp_group),
    cp_stream=communication_stream(),
    cp_comm_type="p2p",
    attention_mask_type="causal",
)

y_i = local_output_projection(o_i)
y_i = local_mlp_norm_and_residual(y_i)         # 继续保持 [S/P,B,M]
return y_i

NVIDIA Megatron Core TP 与 CP 官方总览 — NVIDIA MCore 官方图：红色是 sequence chunk，绿色/蓝色是不同 TP rank；CP 使 Transformer 层中的激活持续按序列切分，而黄色、深蓝块标出 TP/CP 相关集合通信。[^20]

NVIDIA Megatron Core CP 相对全量重计算效果 — NVIDIA MCore 官方结果图：在其 BF16 MHA 与 FP8 GQA 示例中，若干 CP 配置相对 full recompute 获得加速；收益随序列、TP/CP 配比和精度变化。[^20]

实现与迁移：当前 MCore 配置把 p2p 明确标为可异步并与 Attention 重叠，并在 TE wrapper 传入 CP group、global ranks、stream 和通信类型。²⁸ 在支持的 GPT/LLaMA 路径中迁移成本低；若新模型改变 Attention 语义，例如特殊稀疏 mask、cross-attention 或自定义状态空间模块，仍需后端真正理解 CP，而不是只加一个配置开关。

USP¶

方案六：用二维进程网格组合 Ulysses 与 Ring。

起因：纯 Ulysses 集合通信效率高，却受 head 数约束；纯 Ring 不受 head 限制，却可能在节点内浪费 NVLink 集合通信能力，并在大 P 下产生很多 P2P 轮次。
小白直觉：先在“同一排”用 Ulysses 快速换轴，再在“同一列”用 Ring 扩展上下文。假设总并行度 P=P_u×P_r，只有 P_u 消耗可切分 heads，P_r 可以继续增大。
布局变化：初始是 [B,S/(P_uP_r),H,D]；Ulysses A2A 后是 [B,S/P_r,H/P_u,D]；Ring 在 P_r 组中流动 KV；最后逆 A2A 回到原序列分片。

USP 逻辑、流程、时序与数据流 — 自绘四联图。左上用 2×2 网格画出正交通信组；右上给出推荐映射：`P_u` 放节点内高带宽域，`P_r` 放跨节点方向。

u_group, r_group = build_2d_groups(world_size=P_u * P_r)
q_i, k_i, v_i = project_qkv(x_i)              # [B,S/(P_uP_r),H,D]

q_u = all_to_all(q_i, "head", "sequence", u_group)
k_u = all_to_all(k_i, "head", "sequence", u_group)
v_u = all_to_all(v_i, "head", "sequence", u_group)
                                                  # [B,S/P_r,H/P_u,D]
o_u = ring_attention(q_u, k_u, v_u, group=r_group)
                                                  # [B,S/P_r,H/P_u,D]
o_i = all_to_all(o_u, "sequence", "head", u_group)
                                                  # [B,S/(P_uP_r),H,D]
return output_projection(o_i)

反向与重计算¶

USP 的反向不是把前向图“左右翻转”这么简单，而是遵守链式法则的后进先出。前向是：

Q/K/V 局部序列 → A2A₁（序列换头）→ Ring Attention → A2A₂（头换序列）→ O_i。

因此反向依次为：

先穿过 A2A₂：本地序列布局的 dO_i 通过 A2A 变成 head-sharded 的 dO_u=[B,S/P_r,H/P_u,D]。
再做 Ring backward：固定本地 Q_u 和 dO_u，让 K/V block 走一圈；与此同时，属于各个 KV owner 的 dK/dV 部分梯度也沿 Ring 流动并累加。
最后穿过 A2A₁：head-sharded 的 dQ_u/dK_u/dV_u 通过逆布局 A2A，回到每个 rank 原始的局部序列分片。
普通投影层接管：本地 QKV Linear 的 autograd 根据 dQ_i/dK_i/dV_i 继续得到 dX_i 与参数梯度。

Ring backward 需要前向的 Q/K/V、最终输出 O 和 log-sum-exp LSE。LSE 是 online softmax 的归一化账本：没有它，就无法把每一个 KV block 的局部梯度还原到同一个全局 softmax。YunChang 的 RingFlashAttnFunc 正是在 forward context 中保存这些对象，并在自定义 backward() 中调用手写的 Ring FlashAttention backward。⁴⁴

def usp_backward(grad_o_i, saved):
    # 前向最后一个 A2A 的反向：局部序列布局 -> Ulysses 头切分布局
    grad_o_u = all_to_all(
        grad_o_i, scatter="head", gather="sequence", group=u_group
    )                                               # [B,S/P_r,H/P_u,D]

    q_u, k_owner, v_owner = saved.q, saved.k, saved.v
    out_u, lse_u = saved.out, saved.lse
    grad_q_u = zeros_like(q_u)
    grad_k_owner = zeros_like(k_owner)
    grad_v_owner = zeros_like(v_owner)

    for ring_step in range(P_r):
        pair_dq, pair_dk, pair_dv = flash_attention_block_backward(
            grad_o_u=grad_o_u,
            q_u=q_u,
            k_block=k_owner,
            v_block=v_owner,
            global_out=out_u,
            global_lse=lse_u,
            mask=causal_mask_for_ring_step(ring_step),
        )
        grad_q_u = grad_q_u + pair_dq
        grad_k_owner = grad_k_owner + pair_dk
        grad_v_owner = grad_v_owner + pair_dv

        k_owner, v_owner = ring_send_recv(k_owner, v_owner)
        grad_k_owner, grad_v_owner = ring_send_recv_and_accumulate(
            grad_k_owner, grad_v_owner
        )

    # 前向第一个 A2A 的反向：Ulysses 头切分布局 -> 原始局部序列布局
    grad_q_i = all_to_all(
        grad_q_u, scatter="sequence", gather="head", group=u_group
    )
    grad_k_i = all_to_all(
        grad_k_owner, scatter="sequence", gather="head", group=u_group
    )
    grad_v_i = all_to_all(
        grad_v_owner, scatter="sequence", gather="head", group=u_group
    )
    return qkv_projection_backward(grad_q_i, grad_k_i, grad_v_i)

上面把通信写成便于理解的同步形式；高性能实现会把下一轮 K/V 与 dK/dV 的收发提前挂到通信 stream，并与当前 block backward 重叠。关键不变：dQ 对所有可见 KV block 的贡献在 Q owner 本地相加；dK/dV 对所有 Q block 的贡献必须回到 KV owner 后才完整。 源码因此同时维护 K/V 通信器和梯度 K/V 通信器，而不是依赖一次普通 all_reduce。⁴⁴

如果把整个 USP block 放进 activation checkpoint，原始前向不会长期保留内部 A2A/Ring context。反向到达 checkpoint 时会发生：

def checkpointed_usp_backward(x_i_cpu, grad_y_i, rng_snapshot):
    x_i = copy_to_gpu(x_i_cpu)
    restore_rng_state(rng_snapshot)

    with enable_grad():
        y_i = usp_forward(x_i)                     # A2A + Ring + A2A 重跑
    autograd_backward(y_i, grad_y_i)                # 再进入上面的 USP backward
    return x_i.grad

所以启用 checkpoint 后，一层 USP 的通信不是只多一点 bookkeeping，而是原前向一次 + 重算前向一次 + 真正反向一次。这也是为什么 checkpoint 范围不能随意跨越不确定 collective：所有 rank 必须重放完全相同的 A2A/Ring 次序。Dropout 还要求恢复相同 RNG；否则原前向和重算前向的 mask 不一致，梯度就不再对应原来的 loss。

实现支持要看具体 backend

“USP 算法存在反向”不等于每个 kernel 分支都支持训练。YunChang 当前实现的 A2A 与 Ring FlashAttention 路径各自注册了 custom backward；但官方 README 仍标注部分 Torch efficient attention 分支没有 backward。接入新 NPU、稀疏 Attention 或自定义 mask 时，必须逐 backend 检查并做分布式数值对齐，不能只看 forward 能跑。¹⁹⁴⁴

USP 论文中的 Ulysses 与 Ring Attention 原图 — USP 论文 Figure 2 原图：左侧是 Ulysses 的 All-to-All 换轴，右侧是 Ring 的 KV block 轮转；USP 的关键不是发明第三种 Attention，而是把两者放进正交通信组。[^18]

USP 在 208K 序列长度下的性能实验 — USP 论文 Table 7 原图：2×8 A800、LLaMA3-8B 的 TP × Ulysses × Ring 组合；208K 序列达到 147.26 TFLOPS/GPU、47% MFU。[^18]

实现与局限：YunChang 先创建 Ulysses/Ring 两套 process group，LongContextAttention 再顺序执行 A2A、Ring Attention、逆 A2A。²⁹ 它缓解但没有消除 head 限制：P_u 仍必须与 Q/KV heads 兼容。二维 group、TP group 和物理拓扑一旦映射错误，通信会跨越慢链路，代码也比单一方案更难读。

AllGather CP¶

方案七：先复制完整 KV，再调用标准 Attention。

起因：不是每一种新 Attention kernel 都立刻有 Ring 版本。工程上需要一个最简单、最接近单卡语义的 CP 后端，用于先验证正确性或快速接入特殊 Attention。
小白直觉：Q 仍各管各的序列片段，但开考前每个人都复印一套完整 K/V。之后直接调用普通 Attention；反向时再把各人对 K/V 的梯度贡献相加并散回原 owner。
关键对象：Q_i/K_i/V_i 是 [B,S/P,Hq|Hkv,D]；AllGather 后每张卡临时拥有 K_full/V_full=[B,S,Hkv,D]；输出仍是本地 O_i=[B,S/P,Hq,D]。

AllGather CP 逻辑、流程、时序与数据流 — 自绘四联图。左上红字是代价：每张卡都会临时持有完整 KV；左下还展示了反向通过 ReduceScatter 把 `dK/dV` 归并回所属分片。

def forward(x_i):
    q_i, k_i, v_i = project_qkv(x_i)            # [B,S/P,Hq|Hkv,D]
    k_full = all_gather(k_i, dim="sequence", group=cp_group)
    v_full = all_gather(v_i, dim="sequence", group=cp_group)
                                                    # [B,S,Hkv,D]
    o_i = standard_attention(
        q_i, k_full, v_full,
        mask=global_causal_mask_for_local_queries(cp_rank, S, P),
    )                                              # [B,S/P,Hq,D]
    save_for_backward(q_i, k_full, v_full)
    return o_i

def backward(grad_o_i):
    q_i, k_full, v_full = load_saved_tensors()
    dq_i, grad_k_full, grad_v_full = standard_attention_backward(
        grad_o_i, q_i, k_full, v_full
    )
    dk_i = reduce_scatter_sum(grad_k_full, dim="sequence", group=cp_group)
    dv_i = reduce_scatter_sum(grad_v_full, dim="sequence", group=cp_group)
    return project_qkv_backward(dq_i, dk_i, dv_i)

为什么这里没有论文原图

AllGather CP 不是一篇独立论文提出的新 Attention 算法，而是 MCore 的 CP 通信后端。 因此这里不伪造“论文图”，证据应回到 TransformerConfig.cp_comm_type：当前源码明确说明 all_gather 是同步路径，不能与 Attention 计算重叠；同一枚举还列出 p2p、a2a 与 a2a+p2p。²¹²⁸

效果与局限：它不切 head，MHA/GQA/MQA 都能复用标准 kernel，代码最容易读，迁移成本也最低。但完整 KV 的峰值显存会随 S 增长，且当前 MCore 路径不能异步重叠，因此更适合 CP 度较小、正确性基线、Ring kernel 尚未适配或特殊 Attention 快速接入，不适合作为超长序列的大规模默认方案。Python 配置值是 cp_comm_type="all_gather"，部分命令行枚举写作 allgather。

Ulysses-Offload / FPDT¶

方案八：把一次完整 Ulysses Attention 变成 CPU—GPU 流水线。

起因：普通 Ulysses 把每卡输入切到 S/P，但 A2A 换轴后，一个 rank 仍要一次处理“完整 S、部分 heads”的 Q/K/V。序列变到百万 token 后，Attention 的 Q/K/V、反向保存状态、FFN 和 logits 工作区仍会造成峰值。FPDT（Fully Pipelined Distributed Transformer）因此不再把一个 rank 的局部序列当成一整块处理，而是在 Ulysses 内部继续分块。³⁷³⁸
小白直觉：普通 Ulysses 像把一本巨著按“章节”分给多个人；FPDT 又把每个人的章节切成小页包。GPU 桌面上只放当前页包和下一包，其他页包存在 CPU 仓库；GPU 计算当前包时，后台同时搬下一包，尽量把 PCIe 搬运藏在 Attention 计算后面。
关键对象：C 是配置的全局 Attention chunk 长度，u=S/C。每卡本地 X_i=[B,S/P,M] 会被切成 u 份 [B,C/P,M]；QKV 投影并 A2A 后，当前 rank 得到 [B,C,H/P,D]。这个 C/P → C 的换轴是读懂源码和估显存的关键。³⁹

Ulysses-Offload FPDT 的逻辑、流程、时序与数据流 — 自绘四联图。左上区分了 GPU 当前工作集与 CPU 保存量；左下展示 All-to-All、offload、双缓冲 prefetch 和在线 Attention 的重叠；右下明确 A2A 前后是 `C/P` 与全局 `C` 两种尺度。

FPDT 的前向 Attention 可以按下面的教学伪代码理解。真实实现还为反向维护 Q、K、V、输出和 log-sum-exp 等状态，并用独立 CUDA stream 协调 D2H、H2D 与计算：

# 每个 rank 的 x_i 为 [B,S/P,M]；C 是全局 chunk 长度
assert S % P == 0 and S % C == 0
num_chunks = S // C
local_chunks = reorder_then_split(x_i, chunks=num_chunks)
                                                    # 每块 [B,C/P,M]
cpu_qkv = []
local_outputs = []

for chunk_id in range(num_chunks):
    qkv_local = qkv_projection(local_chunks[chunk_id])
                                                    # [B,C/P,H,D]
    q_chunk, k_chunk, v_chunk = ulysses_all_to_all(qkv_local)
                                                    # 各为 [B,C,H/P,D]
    cpu_qkv.append(async_offload_to_pinned_cpu(q_chunk, k_chunk, v_chunk))

    out_chunk, lse_chunk = online_attention_init(q_chunk)
    for history_id in range(chunk_id + 1):          # causal：只看当前及历史 KV
        k_hist, v_hist = double_buffer_prefetch(cpu_qkv[history_id])
        block_out, block_lse = flash_attention_block(
            q_chunk, k_hist, v_hist, causal_block_mask(chunk_id, history_id)
        )
        out_chunk, lse_chunk = online_softmax_merge(
            out_chunk, lse_chunk, block_out, block_lse
        )

    save_backward_state_on_cpu(q_chunk, out_chunk, lse_chunk)
    local_outputs.append(inverse_ulysses_all_to_all(out_chunk))

attention_output_i = concat(local_outputs, dim="local_sequence")
ffn_output_i = tiled_ffn(attention_output_i)
loss = tiled_logits_and_cross_entropy(ffn_output_i, labels_i)

这里的‘任意序列长度’是什么意思

FPDT 让工作集由可调的 C 决定，因此显存不必一次容纳完整 Attention，可以把可训练长度继续推高；它不是说任意整数长度都无需处理。当前 DeepSpeed 实现要求 S 能被 P 与 C 整除，通常由 DataLoader pad 到共同倍数；它依赖 FlashAttention，当前自定义路径只支持无显式 attention_mask 的因果 Attention，并继承 Ulysses 的 head 切分约束。Megatron-DeepSpeed 集成还标注了 FPDT dropout 限制。³⁹⁴⁰

反向与卸载¶

FPDT 不能依赖 autograd 把一个巨大的标准 Attention 图原样反传，因为那会重新把所有 Q/K/V、中间输出和 workspace 同时拉回 GPU，抵消 offload 的意义。因此 DeepSpeed 把 Attention 整体实现成 torch.autograd.Function：前向在 no_grad 下执行并建立自己的 SequenceChunk 状态表；反向则由源码显式控制每一个 chunk 的 H2D、FlashAttention backward、累加、D2H 和 A2A。⁴⁶

小白可以先抓住一个“转置循环”：

前向固定一个 Q chunk，依次访问它能看到的历史 KV chunks，用 online softmax 合并输出。
反向固定一个 KV chunk，依次访问所有能看到它的当前及未来 Q chunks。这样 dK/dV 始终只需要当前 KV 的一份 GPU accumulator；每个 Q 的 dQ 部分和则写回相应的 CPU SequenceChunk。

因果 Attention 中，第 q_id 个 Q 只能看 kv_id ≤ q_id，所以反向外层固定 kv_id 后，内层从 q_id=kv_id 走到最后。下面的伪代码保留了源码中决定峰值的对象和搬运方向：

def fpdt_attention_backward(grad_output_i, state):
    num_chunks = state.num_chunks
    grad_q_cpu = allocate_pinned_accumulators(num_chunks, state.q_shape, fp32=True)
    grad_x_chunks = allocate_gpu_outputs(num_chunks, state.x_local_shape)
    grad_w_qkv = zeros_like(state.w_qkv, dtype="fp32")
    grad_b_qkv = zeros_like(state.b_qkv, dtype="fp32")

    # 输出前向做了 head -> sequence A2A；反向先恢复到 head-sharded 布局
    grad_o_head = [
        all_to_all(
            local_chunk(grad_output_i, q_id),
            scatter="head",
            gather="sequence",
            group=sp_group,
        )
        for q_id in range(num_chunks)
    ]

    for kv_id in range(num_chunks):
        k_gpu = prefetch(state.k_cpu[kv_id], stream="h2d")
        v_gpu = prefetch(state.v_cpu[kv_id], stream="h2d")
        grad_k_gpu = zeros_like(k_gpu, dtype="fp32")
        grad_v_gpu = zeros_like(v_gpu, dtype="fp32")

        for q_id in range(kv_id, num_chunks):
            q_gpu = prefetch(state.q_cpu[q_id], stream="h2d")
            out_gpu = prefetch(state.out_cpu[q_id], stream="h2d")
            lse_gpu = prefetch(state.lse_cpu[q_id], stream="h2d")
            grad_o_gpu = prefetch(grad_o_head[q_id], stream="h2d")

            pair_dq, pair_dk, pair_dv = flash_attention_backward(
                grad_o=grad_o_gpu,
                q=q_gpu,
                k=k_gpu,
                v=v_gpu,
                out=out_gpu,
                lse=lse_gpu,
                causal=(q_id == kv_id),
            )
            grad_q_cpu[q_id].add_from_gpu(pair_dq, stream="d2h")
            grad_k_gpu.add_(pair_dk)
            grad_v_gpu.add_(pair_dv)

            release_gpu(q_gpu, out_gpu, lse_gpu, grad_o_gpu)

        # 当前编号的 dQ、dK、dV 已完整；逆 A2A 回到本地序列布局
        grad_q_local = all_to_all(
            prefetch(grad_q_cpu[kv_id]),
            scatter="sequence",
            gather="head",
            group=sp_group,
        )
        grad_k_local = all_to_all(
            grad_k_gpu, scatter="sequence", gather="head", group=sp_group
        )
        grad_v_local = all_to_all(
            grad_v_gpu, scatter="sequence", gather="head", group=sp_group
        )

        x_local = prefetch(state.layer_input_cpu[kv_id], stream="h2d")
        grad_w_qkv.add_(
            qkv_weight_vjp(x_local, grad_q_local, grad_k_local, grad_v_local)
        )
        grad_b_qkv.add_(
            qkv_bias_vjp(grad_q_local, grad_k_local, grad_v_local)
        )
        grad_x_chunks[kv_id] = qkv_input_vjp(
            state.w_qkv, grad_q_local, grad_k_local, grad_v_local
        )

    return concat(grad_x_chunks), grad_w_qkv, grad_b_qkv

真实代码用 compute、general-offload 和专用 offload stream 做双缓冲：当前 pair 在 FlashAttention backward 中计算时，下一对 Q/O/LSE/dO 正在 H2D；当前 dQ 累加结束后又写回 pinned CPU。权重和 bias 梯度用 FP32 accumulator 跨 chunk 相加，最后一次性转回模型 dtype 并返回给 autograd。训练循环没有逐参数手改 .grad，但这些梯度公式、循环与 stream 都是 FPDT 作者手写的。⁴⁶

FPDT 的 FFN 与 logits/loss 又采用另一种方法：它们不把巨大的中间张量卸载后再取回，而是在反向按 Tile 局部重算：

def tiled_ffn_backward(x, grad_y, w1, b1, w2, b2, tile_size):
    grad_x = empty_like(x)
    grad_w1, grad_b1 = zeros_fp32(w1), zeros_fp32(b1)
    grad_w2, grad_b2 = zeros_fp32(w2), zeros_fp32(b2)

    for token_tile in split_tokens(x, tile_size):
        before_gelu = token_tile @ transpose(w1) + b1   # 反向时重算
        hidden = gelu(before_gelu)
        grad_y_tile = matching_tile(grad_y, token_tile)
        grad_w2.add_(weight_vjp(hidden, grad_y_tile))
        grad_b2.add_(sum_tokens_and_batch(grad_y_tile))
        grad_hidden = grad_y_tile @ w2
        grad_before = grad_hidden * gelu_derivative(before_gelu)
        grad_w1.add_(weight_vjp(token_tile, grad_before))
        grad_b1.add_(sum_tokens_and_batch(grad_before))
        write_matching_tile(grad_x, token_tile, grad_before @ w1)

    return grad_x, grad_w1, grad_b1, grad_w2, grad_b2

def tiled_logits_backward(lm_output_cpu, labels, grad_loss, lm_head, tile_size):
    grad_hidden_tiles = []
    grad_lm_head = zeros_like(lm_head, dtype="fp32")

    for hidden_tile, label_tile in paired_tiles(lm_output_cpu, labels, tile_size):
        hidden_gpu = copy_to_gpu(hidden_tile)
        logits = hidden_gpu @ transpose(lm_head)         # logits 反向时重算
        grad_logits = softmax_cross_entropy_vjp(logits, label_tile, grad_loss)
        grad_hidden_tiles.append(grad_logits @ lm_head)
        grad_lm_head.add_(weight_vjp(hidden_gpu, grad_logits))

    return concat(grad_hidden_tiles), grad_lm_head

这里的 before_gelu、logits 和 softmax 是 operator-local rematerialization：它们确实在 backward 里重算，但不是 PyTorch 自动选择的通用 checkpoint。FPDT 源码明确写出了重算公式与梯度累加。Attention 的 Q/K/V/O/LSE 则主要是 CPU offload/reload，而不是全部重算。⁴⁶

与 checkpoint 叠加¶

Megatron-DeepSpeed 的 FPDT 示例同时打开 activation checkpoint、checkpoint-in-cpu 与 FPDT offload。它们嵌套后的实际时序是：⁴⁵⁴⁷

第一次前向：外层 checkpoint 在 no_grad 下运行 FPDT，只保存 Transformer block 的边界输入，并可把它搬到 CPU；因为内部输入不要求梯度，这一次 FPDT 不保留供长期反向使用的完整自定义 context。
到达 checkpoint 反向：从 CPU 取回 block 输入，恢复 model-parallel RNG，在 enable_grad 下重跑整段 FPDT 前向。Ulysses A2A、chunk Attention 与 online merge 都会再次执行。
建立新 FPDT context：这次重跑会创建真正供反向使用的 Q/K/V/O/LSE SequenceChunk，并按 FPDT 规则卸载到 CPU。
执行 FPDT custom backward：autograd 调用上面的手写分块反向，逐块预取状态、算梯度、逆 A2A，并返回 dX/dW/db。
越过 checkpoint 边界：checkpoint wrapper 把重算输入的 .grad 交还给更早的层。

def checkpoint_plus_fpdt_backward(block_input_cpu, grad_block_output, rng):
    block_input = copy_to_gpu(block_input_cpu)
    restore_model_parallel_rng(rng)

    with enable_grad():
        block_output = fpdt_transformer_block(block_input)
        # 这次重跑会重新做 A2A，并建立 FPDT 自己的 CPU SequenceChunk

    autograd_backward(block_output, grad_block_output)
    # autograd 在内部调用 FPDT Attention、FFN、LogitsLoss 的手写 backward
    return block_input.grad

这套组合用更多前向 FLOPs、额外 A2A 和 PCIe 流量换更低 HBM。FPDT 把 SequenceChunk 等对象直接挂在 custom Function context 并自行管理，不能假设 PyTorch 的通用 saved_tensors_hooks 会自动接管这些私有对象；迁移到新框架时必须保留或重新实现它们的生命周期。⁴³⁴⁶

FPDT 论文中的 Attention offload 调度 — FPDT 论文 Figure 4 原图：红色块表示在 pinned CPU 与 GPU 间移动的 Attention 状态，计算与传输通过流水调度重叠。图示只表达机制，不代表传输在所有 PCIe/NUMA 拓扑上都能完全隐藏。[^38]

显存为什么不是常数：以每卡峰值做量级估算，FPDT 的序列相关部分可写成

M_peak_gpu ≈ M_model + κ0·b·B·S·M/P + κ1·b·B·C·M/P + M_buffer/workspace，

而 host 保存量约为 M_peak_cpu ∝ L·b·B·S·M/P。固定 C 只约束 chunk-local Q/K/V、双缓冲和 kernel 工作区；每卡当前层的本地 hidden state 等基座仍随 S/P 增长。因此 GPU 总峰值不与 S 完全无关，也不与 P 严格成反比；模型状态、通信 buffer、allocator 和某些复制对象会形成下限。

FPDT chunk 大小对激活显存与 MFU 的影响 — FPDT 论文 Figure 12(a) 原图：2.7B、4 GPU、256K 序列下，chunk 越小，激活显存通常越低，但 MFU 也可能因更碎的计算与更多调度下降；论文选 64K 作为该实验的折中，而不是通用最优值。[^38]

效果与局限：论文报告在相同硬件上最长序列最高提升 16×，8B 模型在 4 GPU 上训练 2M token 且 MFU 超过 55%；更大的 70B/4M 配置使用 32 GPU。它的代价是高额 pinned CPU RAM、PCIe/NUMA 敏感性、三条 stream 和自定义 autograd 带来的复杂度。C 太大容易 OOM，太小又使 kernel 变碎、通信更难隐藏，所以必须按模型、硬件和 S 实测峰值与 MFU。³⁷³⁸

FSDP2 迁移¶

FSDP2 负责参数/梯度分片，FPDT 负责 activation 的 chunk/offload，理论上互补，但当前 Ulysses-Offload 是 Megatron-DeepSpeed 的定制执行路径，不是可直接套到 FSDP2 的开关。建议分四步迁移：

先拆清 group 所有权：FSDP2 mesh 只负责参数 AllGather 与梯度 ReduceScatter；独立 sp_group 负责 FPDT 的 Ulysses A2A。两者可以包含同一批 rank，但 collective 的语义、顺序和 stream 不可混用。
把 FPDT Attention 做成显式模块边界：模块接收 sp_group、C 和 pinned-buffer 管理器，不读取 DeepSpeed 全局状态；先只移植 Attention chunk/offload，FFN 与 loss 保持普通实现，得到可对齐的最小版本。
审计参数生命周期：若自定义 backward 或重计算需要再次读取投影权重，必须确保 FSDP2 已在相应 pre-backward 阶段 unshard；按 Transformer block 自底向上 fully_shard，并验证 reshard_after_forward 不会让自定义 autograd 持有失效参数视图。FPDT activation offload 与 FSDP2 参数/梯度 CPU offload 是两套机制，不应共用 buffer 或含混统计。
最后恢复完整流水线：所有 rank 必须进入相同数量、相同顺序的 A2A 与参数 collective；再逐步加入双缓冲、tiled FFN/loss、变长 padding、checkpoint 和 torch.compile。验收应从 SP=1 数值基线、SP=2 固定长度开始，随后覆盖非整除长度、GQA/MQA、gradient accumulation、保存恢复，并分别记录 GPU 峰值、CPU 峰值和 PCIe 暴露时间。

FPDT 的 custom Function 把投影权重作为输入，并在 backward() 返回 dW/db；FSDP2 必须在这段 backward 运行前让对应参数处于可读的 unsharded 状态，并在梯度出现后仍能触发自己的 ReduceScatter/post-backward 生命周期。若为了 checkpoint 在反向中重跑 FPDT forward，参数还会被再次读取一次。因此移植时不能只验证最终 loss，还要给每个参数注册临时 hook，确认梯度只完成一次、shape 与 shard 一致、FSDP2 ReduceScatter 没有早于 custom backward。

这条迁移路线复用的是 FPDT 的对象生命周期与调度思想，不是复制 DeepSpeed 私有 autograd。它比 ALST 的 HF Ulysses wrapper 更难移植；在 FSDP2 上没有公开同条件基准前，应把它视为待验证工程设计，而非现成功能或性能承诺。

ALST¶

方案九：Ulysses 之后，继续处理 MLP、logits 与 checkpoint activation。

起因：旧版 DeepSpeed-Ulysses 主要把长序列分摊给多个 rank，并在 Attention 前后完成序列维与头维的换轴。序列继续变长后，即使每张卡只保留 S/P_s 个 token，[B,S/P_s,M] 的 MLP 中间量、[B,S/P_s,V] 的 logits、逐层保存的 checkpoint activation 和内存碎片仍会依次成为 OOM 来源。DeepSpeed 当前旧教程也已经把读者指向新的 ALST 集成。³⁰³¹
小白直觉：Ulysses 是把一条很长的流水线分给多张卡；ALST 又规定每张卡不要一次加工完自己手中的全部 token，而是把逐 token 工序切成小段，并把暂时不用的 checkpoint 搬到 CPU。它是一套建立在 Ulysses 上的长序列内存工具箱，不是新的 Attention 通信原语。
四个主要部件：HF Ulysses 在 SP rank 间分摊 Attention activation；Tiled Loss 避免物化完整 logits；Tiled MLP 只保留一个序列 Tile 的工作区；activation checkpoint offload 把跨层保存的 hidden states 移到 CPU。Liger fused loss 和 expandable_segments allocator 是额外的高收益实现选择。³²

ALST 的逻辑、流程、时序与张量数据流 — 自绘四联图。左上说明 ALST 是一组可叠加技术；右上和左下展示变长 batch、Ulysses A2A 与有效 token 加权 loss；右下展示 Attention 与逐 Tile MLP/loss 的两条数据路径。

# 教学伪代码：强调变长、边界标签和逐 Tile 工作集；不是可直接运行的训练脚本
install_hf_ulysses_attention(
    core_attention="flash_attention_2",
    sequence_parallel_size=P_s,
    seq_length_is_variable=True,
)

for raw_batch in dataloader:
    S = raw_batch["input_ids"].shape[1]           # 每个 batch 可以不同
    S_pad = ceil_div(S, P_s) * P_s                 # 当前 DS adapter 不会自动补齐
    batch = pad_sequence_fields(raw_batch, S_pad)
    batch["position_ids"] = build_global_positions(batch)
    batch["shift_labels"] = shift_before_sharding(
        batch.pop("labels"), ignore_index=-100
    )

    local = shard_on_sequence(batch, rank=sp_rank, parts=P_s)
    hidden_i = model_without_loss(local)            # Attention 内部执行 A2A → kernel → 逆 A2A

    loss_sum_i = 0.0
    valid_i = 0
    for hidden_tile, label_tile in sequence_tiles(hidden_i, local["shift_labels"]):
        # Liger 或通用 tiled compute：不创建完整 [B,S/P_s,V] logits
        tile_loss_sum, tile_valid = fused_or_tiled_loss(hidden_tile, label_tile)
        loss_sum_i += tile_loss_sum
        valid_i += tile_valid

    loss = differentiable_sp_sum(loss_sum_i) / differentiable_sp_sum(valid_i)
    loss.backward()

这里的任意长度不是无限长度

seq_length_is_variable=True 的含义是：UlyssesSPAttentionHF.forward() 每次从当前 Q 的局部长度重新计算 local_seq_length 与 global_seq_length，因此不同 batch 不必共享一个固定 S。它仍有四个边界：

长度整除：当前 UlyssesSPDataLoaderAdapter 会检查 S mod P_s == 0，不满足就报错。框架可在进入 adapter 前 pad 到 S'=ceil(S/P_s)P_s，并把补位 label 设为 -100；verl 已有这种 pad-and-slice 实现。³³³⁶
Head 约束：Hq mod P_s == 0；KV heads 还要满足 Hkv mod P_s == 0 或 P_s mod Hkv == 0，后一种情况通过复制 KV heads 支持 GQA/MQA。最大 SP 度仍不超过 Q heads。
资源上限：Tiling 把某个逐 token 算子的峰值工作区从“全序列”降到“单 Tile”，但不会让总显存、CPU RAM 或通信缓冲无限大。
计算上限：dense Attention 的时间仍是 O(S²)。论文的 15M token 配置一次迭代约 7 小时 25 分，说明“能跑”与“高吞吐”是两个目标。³²

为什么 label 必须先移位：自回归训练用位置 t 预测 t+1。若先把序列切给多个 rank，再让每个 rank 独立 shift，除最后一片外，每个分片末尾都会丢掉下一个 rank 的首 token 目标。ALST 因此先在完整样本上生成 shift_labels，再沿序列切分；各 rank 的有效 token 数可能不同，最终 loss 必须按 n_i=(shift_labels_i != -100).sum() 加权，而不能简单平均 rank loss。³¹³³

ALST 论文中激活显存随序列长度增长 — 论文 Figure 2 原图：以 Llama-8B 估算，activation memory 随序列长度近似线性增长。这解释了为什么只优化 Attention 通信还不够。[^32]

ALST 各组件逐步扩大最长可训练序列 — 论文 Figure 11 原图：8×H100、Llama-8B 的消融路径从 32K baseline，依次加入 tiled logits/loss、HF Ulysses、checkpoint CPU offload 和 Tiled MLP，最长序列达到 3.7M。[^32]

效果应该怎样读：论文在 Llama-8B、H100、ZeRO-3、batch size 1 的特定配置中实现单卡 500K、8 卡 3.7M、32 卡 15M，相对同文 32K baseline 分别是 16×、116×、469× 的最长可运行序列提升。32K 的训练 loss 与 baseline 几乎重合，但论文明确把“可运行的最长序列”放在吞吐之前；CPU offload 还可能把瓶颈转移到主机内存，例如大模型多层 checkpoint 会占用数百 GiB。³²

FSDP2 迁移¶

FSDP2 解决的是参数、梯度和优化器状态分片，ALST 解决的是长序列 activation 与工作区，两者目标互补。现有 verl 已有 FSDP2 + Ulysses 的真实计算路径，证明组合并非概念设想；但这只证明 Ulysses 与 FSDP2 能组合，不能自动推出 DeepSpeed 的 TiledMLP 与 checkpoint offload 已经可直接复用。³⁶

ALST 迁移到 FSDP2 的双 Mesh 设计 — 自绘迁移图。绿色 FSDP2 mesh 负责参数 AllGather/梯度 ReduceScatter；橙色 SP group 负责 Ulysses activation A2A。两套 group 可以覆盖同一批 rank，但不能把“参数分片轴”和“样本/序列分组轴”当成同一个概念。

建议按以下顺序实现：

建立两套逻辑 Mesh：fsdp_mesh 默认可以覆盖 WORLD，用于 fully_shard；另建 [dp,sp] Ulysses mesh，从中取得 sp_group 和数据并行 batch group。若设置较小 fsdp_size，PyTorch/框架可能建立 [ddp,fsdp] HSDP mesh，但它仍独立于 Ulysses 的 [dp,sp] 语义。³⁵
不要只在 DP 轴 shard 参数：若参数在 SP 轴复制，却没有 HSDP/DDP 等价的 SP 梯度同步，各 SP 副本会更新成不同模型。最简单的 ZeRO-3 对等起点是让 FSDP2 shard group 覆盖所有训练 rank；优化成 HSDP 前先证明梯度归约语义和 loss 缩放正确。
先移植 Attention 与数据边界：DeepSpeed 的 UlyssesSPAttentionHF 构造器已经接受外部 process_group，但 register_with_transformers() 会自行初始化 DeepSpeed SP group。FSDP2 后端应让注册函数接收 sp_group，在模型实例化前注入 HF Attention wrapper；DataLoader 负责 pad、position_ids、预移位 label 和 SP 分片。
按层 bottom-up 调用 fully_shard：FSDP2 在层前 AllGather 参数、层后释放未分片参数，并在反向后 ReduceScatter 梯度。Transformer block 和 root 分层 wrap 才能重叠通信并控制峰值；不要只 wrap root。³⁴
分阶段移植 Tiling：第一阶段先用 FSDP2 + HF Ulysses + Liger/tiled loss，验证单卡 reference；第二阶段再接 TiledMLP。当前 DeepSpeed TiledMLP.backward 通过 ZeRO 专属 param.ds_grad_is_ready 控制“最后一个 Tile 才归约梯度”，源码也明确留下 DDP/FSDP 需要等价机制的注释，因此不能原样复制。³³
区分两种 CPU offload：FSDP2 CPUOffloadPolicy 卸载参数/梯度，ALST checkpoint offload 卸载 activation；它们是两个独立开关。二者同时开启前要评估 PCIe 带宽、pinned memory、CPU RAM 和 gradient accumulation 的正确性。

# FSDP2 后端骨架：展示 group 所有权与 wrap 顺序
ulysses_mesh = init_device_mesh(
    "cuda", (dp_size, sp_size), mesh_dim_names=("dp", "sp")
)
sp_group = ulysses_mesh["sp"].get_group()
dp_group = ulysses_mesh["dp"].get_group()

fsdp_mesh = init_fsdp_mesh(world_size, fsdp_size=-1)  # MVP：WORLD 上 full shard
install_hf_ulysses_attention(
    process_group=sp_group,
    seq_length_is_variable=True,
    core_attention="flash_attention_2",
)
model = AutoModelForCausalLM.from_pretrained(model_config.model_path)

for block in model.model.layers:                    # bottom-up
    fully_shard(block, mesh=fsdp_mesh, reshard_after_forward=True)
fully_shard(model, mesh=fsdp_mesh, reshard_after_forward=False)

loader = VariableLengthSPAdapter(
    dataloader,
    sp_group=sp_group,
    dp_group=dp_group,
    pad_to_multiple=sp_size,
    pre_shift_labels=True,
)

FSDP2 迁移的验收顺序

先做 SP=1/FSDP2 与单卡 loss 对齐，再做 SP=2/FSDP2 的固定长度，随后测试不能整除的变长 batch、packed sequence、GQA/MQA、gradient accumulation、checkpoint save/resume。最后才打开 TiledMLP 和两类 CPU offload，并分别记录峰值 GPU/CPU 内存、step time 与通信暴露。目前没有 ALST-on-FSDP2 的论文基准，因此这一段是可验证的工程路线，不是性能承诺。

MindSpeed-MM Chunk Layer¶

方案十：不增加 SP rank，而是在一个 FSDP2 rank 内把 Decoder Layer 沿序列串行切块。

起因与一句话思路¶

起因：FSDP2 已经分摊参数、梯度和优化器状态，但一个 rank 仍可能一次性物化整段 [B,S,M] 的 Attention、MoE/MLP 和保存激活。序列增长时，参数能装下不等于这一层的 activation 峰值也能装下。
小白直觉：把一整车货物改成 N_c 趟运输，每趟只加工 C=S/N_c 个 token；为了不破坏因果关系，FA 把前面 token 的 K/V 当作“历史账本”，GDN 把卷积窗口与 recurrent delta state 当作“接力棒”传给下一块。
严格分类：这不是新的 Ulysses、Ring 或 USP 通信后端。调度器看到 ulysses_parallel_size>1 或 ring_attention_size>1 会直接报错，所以它是层内串行分块 + 重计算 + 激活卸载，不是跨 rank 序列并行。它可以运行在采用 FSDP2 的训练栈里，但当前不能与 SP/CP 叠加。⁴⁸
设计来源边界：第一版源码说明串行 sequence-chunking 语义跟随 XTuner bbd10c1；第二版在 MindSpeed-MM 中重做了 FSDP2 wrapper、FA/GDN cache、D2H 去重和 backward 生命周期。因此更准确的称呼是基于既有串行分块语义的 MindSpeed-MM FSDP2 工程实现，而不是一篇已经发表并完成跨系统评测的新 SP 算法。⁴⁹

MindSpeed-MM Chunk Layer 物理对照、逻辑、流程、时序与数据流 — 源码固定到 `chunk_layer` 分支提交 `9e9610a4` 的自绘五视图。左上是整层与分块后的物理对象，右上解释第二版为何把所选缓存 D2H 降为线性，中间是前向调度，左下是 NPU/CPU 组件时序，右下是 FA/GDN 张量流。图中同时保留了完整输出、线性 KV staging 与 `SP/CP=1` 三个关键边界。

前向与跨块状态¶

框架在 fully_shard 之前分两阶段改写模块。先把 module.forward 换成只计算一个 Chunk 的 decoder，再套上 recompute 和 activation-offload wrapper，最后把逐 Chunk scheduler 放回最外层。源码注释给出的实际嵌套关系是：⁴⁸

# MindSpeed-MM 9e9610a4, apply_features.py:96-106 的关键原文
prepare_chunk_layer_modules(model)

self.apply_recompute_models(model=model)
self.apply_activation_offload_modules(model=model)
finalize_chunk_layer_modules(model)
self.apply_chunk_mbs(model=model)

也就是：

ChunkMBS(ChunkLayer(Offload(Recompute(ChunkDecoder))))。

为什么 scheduler 必须在外层？如果先让 checkpoint 包住整个层，反向重算时仍会一次重新物化完整 S；现在框架只认识一个 [B,C,M] 的 ChunkDecoder，所以保存、卸载和重算的最小单元也随之变成一个 Chunk。

前向可归一化成下面的完整教学伪代码。它同时展示 FA 与 GDN 两类状态，并保留源码中的整除、layout 和训练态约束：

def chunk_layer_forward(x, position_embeddings, position_ids, cu_seqlens, Nc):
    B, S, M = x.shape
    assert not using_generation_kv_cache()
    assert ulysses_size() == 1 and ring_attention_size() == 1
    assert S % Nc == 0
    C = S // Nc

    layout = "tnd" if cu_seqlens is not None else "bsnd"
    if layout == "tnd":
        assert B == 1
        assert int(cu_seqlens[-1]) == S

    plan = build_chunk_plan(
        layout=layout,
        batch_size=B,
        sequence_length=S,
        num_chunks=Nc,
        cu_seqlens=cu_seqlens,
    )
    layer_cache = create_fa_or_gdn_cache(plan)
    cache_manager.register(layer_cache)

    x = mark_cache_backward_exit(x, plan.cache_key)
    output = empty_like(x)                         # [B,S,M] 仍完整存在
    history_1 = None
    history_2 = None

    for chunk_idx in range(Nc):
        start = chunk_idx * C
        end = start + C
        x_chunk = x[:, start:end, :].clone()       # [B,C,M]，独立 storage
        pos_chunk = slice_positions(
            position_embeddings, position_ids, start, end
        )
        local_meta = plan.chunk_infos[chunk_idx]

        if layer_type() == "full_attention":
            q, k_now, v_now = project_qkv_and_apply_rope(x_chunk, pos_chunk)
            k_ctx, v_ctx, kernel_kwargs = layer_cache.prepare_cache(
                cache=(k_now, v_now),
                history_cache=(history_1, history_2),
                chunk_idx=chunk_idx,
            )
            mixed = causal_flash_attention(
                q=q,
                k=k_ctx,
                v=v_ctx,
                query_cu_seqlens=local_meta["local_cu_seqlens"],
                **kernel_kwargs,
            )
            history_1, history_2 = layer_cache.append(
                cache=(k_ctx, v_ctx),
                history_cache=(history_1, history_2),
                chunk_idx=chunk_idx,
            )
        else:
            conv_0, delta_0 = layer_cache.prepare_cache(
                cache=empty_gdn_states(x_chunk),
                history_cache=(history_1, history_2),
                chunk_idx=chunk_idx,
            )[0]
            q, k, v, conv_1 = stateful_causal_conv_projection(
                x_chunk, initial_state=conv_0, output_final_state=True
            )
            mixed, delta_1 = gated_delta_rule(
                q, k, v, initial_state=delta_0, output_final_state=True
            )
            history_1, history_2 = layer_cache.append(
                cache=(conv_1, delta_1),
                history_cache=(history_1, history_2),
                chunk_idx=chunk_idx,
            )

        y_chunk = residual_moe_mlp(x_chunk, mixed)
        cache_manager.complete_chunk(plan.cache_key, chunk_idx)
        output[:, start:end, :].copy_(y_chunk)      # CopySlices 保留梯度边

    current_key, previous_key = cache_manager.finalize(plan.cache_key)
    return mark_cache_backward_enter(output, current_key, previous_key)

两类状态的语义不同：

FA 是增长的因果前缀：当前 Chunk 只生成 Q_c，但 K/V 是“历史 + 当前”。BSND 中每个 batch 样本的前缀持续增长；packed TND 中只携带恰好跨过 Chunk 边界的最后一篇文档，下一篇文档从空状态开始。
GDN 是固定尺寸的递归摘要：不保存全部历史 token，而是携带 causal-conv 窗口和 delta-rule recurrent matrix。状态本身不随文档长度增长，但反向需要保留每个 Chunk 边界的状态链。
完整层边界仍存在：输入 x 和 output_buffer 都是 [B,S,M]。分块封顶的是当前层内部工作区，不是把所有线性于 S 的对象都消灭。

反向、重计算与卸载¶

训练脚本仍只调用 loss.backward()，但“自动”的含义与 USP/FPDT 一样：autograd 自动按依赖调用已经注册的规则，源码作者仍要手写状态保存、H2D/D2H 时机和特殊算子的 VJP。

这版实现把责任拆成四层：

Chunk 重计算：ChunkDecoder 已被 checkpoint wrapper 包裹。反向到达某个 Chunk 时，框架取回该 Chunk 输入和跨块状态，重跑 Norm、FA/GDN、MoE/MLP，再对新图求导。
saved-tensor pack/unpack：带 _layer_cache_spec 的 K/V 或 GDN state 不走普通逐 tensor offload，而是变成 CacheHandle。前向先把新值放进 layer staging；反向一次 H2D 后，unpack() 用 tensor.set_() 把原 autograd tensor 重新绑定到恢复后的 view。
层级预取与释放：输出上的 _CacheBackwardEnter.backward() 先恢复当前层 cache，并预取前一层；输入上的 _CacheBackwardExit.backward() 在本层梯度全部通过后释放 device/host cache。
算子 VJP：FA backend 使用它已注册的 Attention backward；GDN 新增的 stateful causal-conv 是手写 torch.autograd.Function.backward()，显式计算 dx/dw/db/d_initial_state，还把 d_final_state 加回输入窗口。⁴⁹

下面的伪代码把这些分散在 saved-tensor hook、自定义 Function 和 manager 中的行为串成一条时间线：

def pack_saved_tensor(tensor):
    spec = getattr(tensor, "_layer_cache_spec", None)
    if spec is not None:
        return cache_manager.pack(tensor, spec)     # D2D 到唯一 staging slot
    if is_current_chunk_input(tensor):
        return ordinary_activation_offload(tensor) # Chunk 输入 D2H
    return tensor

def unpack_saved_tensor(saved):
    if isinstance(saved, CacheHandle):
        cache = cache_manager.get(saved.spec.cache_key)
        restored_view = cache.view_for(saved)      # 必要时先整层 H2D
        saved.tensor.set_(
            restored_view.untyped_storage(),
            restored_view.storage_offset(),
            restored_view.size(),
            restored_view.stride(),
        )
        return saved.tensor
    return ordinary_activation_reload(saved)

def cache_backward_enter(grad_output, current_key, previous_key):
    cache_manager.get(current_key)                  # 等待当前层 H2D 完成
    cache_manager.prefetch(previous_key)            # 与当前层 backward 重叠
    return grad_output

def recompute_one_chunk_and_backward(
    chunk_inputs,
    grad_chunk_output,
    grad_next_state_1,
    grad_next_state_2,
    rng_state,
):
    x_chunk, history_1, history_2 = reload_chunk_inputs(chunk_inputs)
    restore_rng_state(rng_state)
    with enable_grad(), saved_tensor_hooks(pack_saved_tensor, unpack_saved_tensor):
        y_chunk, next_1, next_2 = chunk_decoder(
            x_chunk, history_1, history_2
        )
    autograd_backward(
        outputs=(y_chunk, next_1, next_2),
        grad_outputs=(
            grad_chunk_output,
            grad_next_state_1,
            grad_next_state_2,
        ),
    )
    return (
        gradient_of(x_chunk),
        gradient_of_or_none(history_1),
        gradient_of_or_none(history_2),
    )

def cache_backward_exit(grad_input, current_key):
    cache_manager.release(current_key)
    return grad_input

Chunk 之间不是彼此独立的：后一个 Chunk 的 FA K/V 或 GDN state 来自前一个 Chunk，因此梯度沿这条状态链从 Chunk Nc-1 向 Chunk 0 传播。调度顺序由 autograd 自动得到，跨块状态为何可导、保存什么、何时恢复和特殊 kernel 如何返回 state gradient，都是实现者显式写出的。

第二版最重要的工程变化发生在 D2H 之前。第一版会把第 i 个 Chunk 的累计前缀再次当作 checkpoint 输入，FA 搬运量近似：

Σᵢ 2bBHkv(iC)D = Θ(bBHkvD·S²/C)。

第二版的 staging 为每条因果链预留位置，只复制 tensor[:, :, reused_tokens:, :]，旧 token 不再重复写入；层结束时把已用区间压紧并一次 D2H：

M_FA-D2H ≈ 2bBHkvSD = Θ(S)。

GDN 则保存 N_c-1 组边界状态，传输量约为：

M_GDN-D2H ≈ (N_c-1)b[B·Dconv·W + B·Hv·Dk·Dv]。

当 C 固定时 N_c=S/C，它同样随 S 线性增长。这里的 O(S) 只描述所选跨块 cache 的实际搬运量；每层 Chunk 输入的 activation offload、CPU 常驻量和 H2D 回读仍需单独统计。

与 USP、FPDT 的差异¶

三者都出现“切序列”和“状态搬运”，但解决问题的轴不同：

维度	USP	Ulysses-Offload / FPDT	MindSpeed-MM Chunk Layer
切分轴	空间并行：`P_u×P_r` 个 rank	先按 Ulysses 分 rank，再在 rank 内按全局 Chunk 流水	时间串行：一个 rank 内 `N_c` 个 Layer Chunk
Attention 全局上下文	A2A 换轴后用 Ring 流动 KV	每个 Chunk A2A 后，用在线 Attention 合并多个 Q/KV Chunk	FA 拼接当前 Q 所需的本地因果 K/V 前缀；GDN 传递递归状态
Rank 通信	Ulysses A2A + Ring P2P	Ulysses A2A，另有 CPU/NPU 传输	无 SP 通信；当前源码要求 Ulysses/Ring size 为 1
Head 限制	只有 `P_u` 消耗可切分 Q/KV heads	继承 Ulysses 的 head 限制	没有 head 整除限制，但也没有 `P` 卡显存分摊
GPU 工作集	主序列项近似 `S/(P_uP_r)`	当前全局 `C` 的 head shard，加本地 `S/P` 基座	当前 `[B,C,M]` 工作区，但 FA 最终前缀和 staging 可到 `S`
反向实现	手写 A2A 与 Ring Attention VJP	手写完整 FPDT custom backward 与流水	autograd 串起 Chunk；手写 cache hooks/sentinels 与 stateful-conv VJP
框架迁移	要管理二维 group 和 Attention backend	要移植 DeepSpeed 私有调度、自定义 backward 与 buffer	已原生接在 MindSpeed-MM FSDP2 pre-shard；迁新模型仍要定义正确的跨块充分状态
新模型适配	特殊 mask、GQA、cross-attention 需核对	Attention、FFN、loss、dropout 和 layout 都要核对	当前只接 Qwen3.5 MoE 的 full Attention / GDN；新 Token Mixer 必须证明所携状态足以重建完整因果语义
公开效果	有论文基准	有论文、教程和代码基准	只有分支代码；最新提交没有测试或性能结果

因此它不是“用更少代码实现了 USP”，而是走了另一条路线：不把序列分给更多设备，而是让一张设备分时处理。 当 rank 数有限、FSDP2 参数分片已就绪、当前瓶颈来自层内激活时，这个方向很有价值；当目标是让序列相关对象真正按 1/P 分摊，仍要使用 Ring、Ulysses、USP 或支持组合的 CP backend。

显存与“任意长度”的边界¶

对 full Attention 路径，源码可直接推出一个保守的单卡峰值分解：

M_peak ≳ M_FSDP-model + κ₀bBSM + 2bBHkvSD + κ₁bBCM + M_workspace + M_allocator。

κ₀bBSM：层输入和完整 output_buffer；κ₀ 取决于两者生命周期能否重叠释放。
2bBHkvSD：FA 的 K/V staging，或最后一个 Chunk 看到的因果前缀。
κbBCM：当前 Chunk 的 Norm、Q、MoE/MLP 和重计算工作区。
GDN 不需要增长的 KV 前缀，但 backward boundary cache 约为 (N_c-1) 份 conv/delta state。

所以它的结论不是“峰值与 S 无关”，而是：把一部分原来按整层 S 物化的高系数工作区改成按 C 物化，同时让保存状态的 D2H 去重为线性。 总峰值仍有线性 S 项，也完全没有随 SP size 线性下降的关系，因为当前 P_sp=1。

这里的任意长度不是无限长度

长度整除：必须满足 S mod N_c = 0；源码没有自动 pad 尾块。
并行互斥：当前明确拒绝 Ulysses 和 Ring，不能直接与 USP 组合。
Layout 边界：BSND 支持普通 batch；packed TND 仍要求 micro_batch_size=1。
训练专用：past_key_values 会报错，不能把训练用状态链当成生成 KV cache。
语义边界：scheduler 丢弃传入的 attention_mask，当前实现围绕因果 FA 和 Qwen3.5 GDN；任意稀疏、滑窗、cross-attention 或自定义 mask 都要重新证明等价性。
Kernel 边界：GDN chunk path 要求 causal-conv 与 GDN 都使用指定 Triton 实现；不是任意 backend 自动获得 state gradient。
资源与时间边界：pinned CPU、H2D/D2H、完整输出和 FA staging 仍随 S 增长；dense Attention 总计算仍是 O(S²)。

效果应该怎样读：提交 9e9610a4 的标题声称第二版把 D2H 降到 O(n)，源码也能验证“每个新增 K/V token 只进入 staging 一次、最终一次压紧 D2H”的机制；但这个 commit 没有测试文件、峰值曲线、最大序列、吞吐、MFU 或数值对齐结果。本文因此只把它评价为设计上支持把目标长度继续向上扩展的原型。要升级成“任意长度已支持”，至少应补四组证据：与未分块版本的 loss/gradient 对齐、随 S/N_c 的 HBM/CPU/D2H 曲线、FA 与 GDN 的 BSND/TND 矩阵、以及 FSDP2 多卡 step time 与失败边界。

工程联评¶

下表中的迁移成本与可读性是基于数据布局、通信组、mask/position 处理和 kernel 依赖做出的工程判断，不是论文测量值。显存列统一估算单卡同一时刻的活跃对象：

M_peak = M_model-state + M_saved-activation + max_t(M_operator-workspace(t) + M_comm-buffer(t)) + M_allocator-margin。

其中 b 是每个 activation 元素的字节数，C 是当前方法的 Chunk 长度，N_c=S/C 是层内 Chunk 数，T 是 MLP/loss 的本地 Tile 长度，V 是词表大小，κ 表示实现相关的 Q/K/V 份数、反向保存和 kernel workspace。不能把所有历史 tensor 简单相加，也不能只看渐进式就声称总峰值严格按 S/P 缩放。

方案	头数与拓扑限制	代码复杂度与可读性	新框架迁移成本	新模型是否重新适配	显存峰值估算；是否严格随 `S` 线性、随 `P` 反比	主要适用边界
Ring Self-Attention（早期）	无头数上限；需要 K、V 两个完整 Ring 阶段	较高；分数阶段与 Value 阶段分离，生命周期长	高	新 mask、位置编码和反向调度都需核对；现代 kernel 通常不会直接复用	每卡保留全部 owner 的分数/概率块时，主项约 `κ2·b·B·Hq·S²/P + κ1·b·B·S·M/P`。对 `S` 不是线性，而有二次项；对 `P` 也只是主项近似反比。	适合作为历史机制与正确性对照；通常不作为现代长上下文默认实现
Megatron SP	必须依附 TP；不负责 Attention 全局上下文	框架内高可读；通信藏在 TP Linear 边界	MCore 内低，跨框架中	通常不改 Attention；需保证 Norm、Dropout、残差和 RNG 的序列分片语义	长期保存的逐 token 激活约 `κs·b·B·S·M/P`，但 TP 子层前的临时 gathered activation 可达 `κg·b·B·S·M`，另有模型状态。保存项近似 `S/P`，总峰值不严格。	用于 TP 激活优化，可与 CP 叠加；不能替代长上下文 CP
Ulysses	`P_u` 受可切分头数，GQA 常受 KV heads 限制；偏好高带宽 All-to-All	中等；Attention 前后各一个换轴	中等	标准 MHA 较容易；GQA/MQA、TP 共存、packed sequence、cross-attention 需重新核对布局	Q/K/V、A2A 接收与 Attention 保存主项约 `κ·b·B·S·M/P_u`。固定 heads、kernel 与 buffer 策略时主项近似随 `S` 线性、随 `P_u` 反比；总峰值不是严格关系。	头数充足、节点内高带宽、希望复用现有 FlashAttention
Ring Attention	无头数上限；依赖 P2P 与足够大的计算块覆盖通信	较高；分块调度、online softmax 与反向较难读	高	新 mask、RoPE、变长序列、cross-attention 或特殊 Attention kernel 通常需适配	本地 Q、当前/下一 KV、输出和 online-softmax 状态约 `κ·b·B·S·M/P`。主项近似线性/反比，但双缓冲、kernel workspace 和模型状态使总峰值不严格。	GQA/MQA、SP 度大于头数、显存优先的超长序列
Megatron CP（P2P）	无头数上限；当前高性能路径依赖 Transformer Engine/支持的 Attention	应用层低，后端高	MCore GPT/LLaMA 路径低；移植到新框架高	共享 GPT 路径的模型多为配置接入；新 Attention 类型仍需实现 CP 语义	与工程化 Ring 类似，序列 activation 与环传 buffer 主项约 `κ·b·B·S·M/P`。主项近似线性/反比；TE workspace、因果负载均衡与重计算策略会改变实际峰值。	NVIDIA/Megatron 训练栈的默认长上下文方案
USP	仅 `P_u` 受头数限制；可按 NVLink/RDMA 拆二维组	较高；两套通信组和布局转换降低直观性	中高	标准 Attention 可集中在 wrapper；特殊 mask、GQA 映射和训练框架 group 管理仍需适配	令 `P=P_uP_r`，序列对象主项约 `κ·b·B·S·M/P`，再加 A2A 与 Ring 双通信工作区。主项近似线性/反比；二维 buffer 和拓扑使总峰值不严格。	多节点分层网络、GQA/MQA、既要集合通信效率又要扩展 SP 度
AllGather CP	无头数限制；要求每卡能临时容纳完整 K/V	最高可读性，最接近单卡 Attention	低到中	通常只需全局 KV、局部 Q 的 mask/position 与反向归并；新 kernel 适配最少	完整 KV 复制项约 `2·b·B·S·Hkv·D`，本地 Q/输出约 `κ·b·B·S·Hq·D/P`。对 `S` 主项线性，但 KV 主项与 `P` 无关，因此绝非严格反比。	正确性基线、小 CP 度、特殊 Attention 快速接入；超长序列受 KV 峰值显存限制
Ulysses-Offload / FPDT	继承 Ulysses head 约束；`S` 需被 `P`、`C` 整除；依赖 FlashAttention、pinned CPU 与可覆盖的 PCIe	很高；chunk 重排、A2A、双缓冲、online merge 与自定义反向交织	Megatron-DeepSpeed 内中；迁 FSDP2/新框架很高	因果 MHA/RoPE 路径较直接；显式 mask、dropout、GQA/MQA、packed/cross-attention 和新 autograd 后端均需验证	GPU 约 `M_model + κ0·b·B·S·M/P + κ1·b·B·C·M/P + buffer`；CPU 另约 `L·b·B·S·M/P`。固定 `C` 只封顶 chunk 工作区；GPU 基座仍随 `S/P`，总峰值不严格线性/反比。	Attention 状态成为峰值、CPU RAM 充足且传输能被计算覆盖的百万 token 训练；不适合弱 PCIe/NUMA 或追求最薄代码的框架
ALST（基于 Ulysses）	保留 Ulysses 的 Q/KV-head 兼容约束；变长仍需 pad 到 SP 倍数	高；同时改 Attention 注册、DataLoader、loss、MLP 与 offload	DeepSpeed/ArcticTraining 内中，迁到 FSDP2 为高	HF wrapper 降低 Attention 适配，但 TiledMLP、packed mask、loss 与 checkpoint offload 仍需逐模型/框架验证	序列相关 GPU 峰值近似取 `max(κatt·b·B·S·M/P, κmlp·b·B·T·M, κloss·b·B·T·V)`，CPU checkpoint 约 `L·b·B·S·M/P`。只有部分项近似 `S/P`；Tile、词表与模型状态项不是。	以“能训练超长单样本”为第一目标的 post-training；CPU RAM、`O(S²)` 时间和极小有效 batch 是主要边界
MindSpeed-MM Chunk Layer	无 head 整除限制；但当前强制 Ulysses/Ring size 为 1，`S mod N_c=0`，packed TND 要求 `B=1`	很高；wrapper 顺序、saved-tensor hook、storage rebind、FA/GDN 两套 cache 与状态 VJP 交织	MindSpeed-MM FSDP2 路径内低到中；迁新框架高	当前只适配 Qwen3.5 MoE 的 full Attention / GDN；新 Token Mixer、mask、layout 和 kernel 都要定义并验证跨块充分状态	FA 路径约 `M_FSDP + κ0·bBSM + 2bBHkvSD + κ1·bBCM + buffer`；GDN 另存 `(N_c-1)` 份边界状态。D2H 主项可随 `S` 线性，但 HBM 仍有线性 `S` 项；因为 `P_sp=1`，不存在随 `P` 反比。	FSDP2 单 rank 层内 activation 削峰原型；适合先验证 Qwen3.5 FA/GDN，不能当作 USP/CP 替代，且尚无公开性能证据

选型建议¶

先区分目标：只想减少 TP 的冗余激活，使用 Megatron SP；真正被上下文长度卡住，才增加 CP/SP Attention。
Megatron + Transformer Engine 的常规模型：优先从 P2P CP 开始；AllGather CP 用作正确性对照、较小 CP 度或尚未支持 Ring 的 Attention 类型。
头数充足且通信域带宽高：纯 Ulysses 通常代码更薄、kernel 粒度更大；但需要验证 P_u 能整除 Q/KV heads，并检查 TP 是否也在切头维。
GQA/MQA 或总 SP 度超过 KV head 数：优先 Ring 或 USP；多节点时可把 Ulysses 放在节点内、Ring 放在节点间，再用 profile 选择二维比例。
不要只看最长序列：固定相同模型与 batch 语义，同时记录峰值显存、吞吐/MFU、通信暴露、因果负载均衡、数值一致性，以及变长/packed sequence 的有效 token 比例。
Ulysses 的 Attention Q/K/V 或反向状态先 OOM：考虑 Ulysses-Offload/FPDT。先 profile 出 C 的可行区间，同时监测 pinned CPU RAM 与 PCIe 暴露；FSDP2 后端应先移植最小 Attention 模块，不要直接照搬 DeepSpeed 私有 autograd。
已经用 Ulysses 但仍在非 Attention 位置 OOM：再引入 ALST。先启用 tiled/fused loss，再按 profile 决定 TiledMLP 和 checkpoint CPU offload；迁移到 FSDP2 时先完成 Ulysses 与 loss 的正确性闭环，不要一次打开所有内存技巧。
已经采用 MindSpeed-MM FSDP2，且暂时不需要 SP/CP：可以把 Chunk Layer 作为 Qwen3.5 FA/GDN 的实验性削峰路径。先做 N_c=1 与 N_c>1 的 loss/gradient 对齐，再测 D2H/H2D、HBM 和 CPU RAM；若目标要求 1/P 分摊或多节点扩展，应回到 Ulysses、Ring、USP 或 FPDT，而不是继续增大 N_c。

流水线并行 PP¶

Pipeline Model Parallelism (层间切分的模型并行)

思想：AI训练是重复的有依赖长过程，可以打散成有依赖的基本单元micro-batch进行流水线调度，提高设备的利用率。
难点：依赖基本单元间的数据传输时间，如何隐藏。流水线并行的方式更复杂，并且micro-batch的方式减少了单节点计算密集度，增加了节点间的信息传递频率，使得取得一个好的加速比成为一个难题。
优点：解决了数据并行显存利用率低的问题，其通过对模型的切分，每个节点只需要放置一部分的模型参数，从而使得其可以部署更大的模型。

GPipe¶

³

1F1B¶

1F1B

PipeDream¶

⁴

VPP¶

假定当前模型网络共16层（编号 0-15），4个Device，
前述GPipe模式和PipeDream是分成4个stage，按编号0-3层放Device1，4-7层放Device2，并以此类推。
virtual pipeline则是按照文中提出virtual_pipeline_stage概念减小切分粒度，
以virtaul_pipeline_stage=2为例，将0-1层放Device1，2-3层放在Device2，...，6-7层放到Device4，8-9层继续放在Device1，10-11层放在Device2，...，14-15层放在Device4。

DualPipe¶

idea

传统流水线并行卡之间串行，效率低。可以0号和7号卡同时都是layer0和layer7,这样做的目的：真个流水线能同时运行两批次的数据！比如第1批从gpu0卡的layer0开始forward，第10批数据从gpu7卡的layer0开始forward，提升效率！

DualPipe 是一种创新的双向管道并行算法，在 DeepSeek-V3 技术报告中提出。实现了正向和反向计算-通信阶段的完全重叠，同时也减少了管道气泡时间。

混合并行¶

2021年10月，微软和英伟达联合提出了 PTD-P(Inter-node Pipeline Parallelism, Intra-node Tensor Parallelism, and Data Parallelism)训练加速方法，
通过数据并行、张量并行和 Pipeline 并行“三管齐下”的方式，将模型的吞吐量提高 10%以上。
该并行方法可以在3072个GPU 上，以502P的算力对一万亿参数的GPT 架构模型进行训练，实现单GPU吞吐量52%的性能提升。

Bloom-176B

DP8，TP4, PP12：

megatron 并行默认维度

采用的顺序是tp-cp-ep-dp-pp，我们认为越靠前的并行组，通讯量越大，所以尽量安排在一台机器内的更近的维度。¹²

专家并行 MoE(EP)¶

MOE的概念

思想：一种基于稀疏 MoE(Mixture-of-Experts) 层的深度学习模型架构被提出，即将大模型拆分成多个小模型(专家，expert)，每轮迭代根据样本决定激活一部分专家用于计算，
优点：只计算一部分，达到了节省计算资源的效果；
实现：MoE 将模型的某一层扩展为多个具有相同结构的专家网络(expert)，并由门(gate)网络决定激活哪些 expert 用于计算，从而实现超大规模稀疏模型的训练。

MOE 与 EP的关联

¹⁷

MOE层如果不开EP， MoE的结构由多个 expert 构成，每次只选择一个专家执行；
MOE层开启EP，MOE层的多个专家就能同时并行，处理不同的tokens，从而达到并行的效果。为了实现这点需要对topk的token来dispatch和combine

EP 与 TP的区别

TP开启时：每个 EP rank 上只包含一部分 expert，而每个 EP rank 上的 token（即 token 对应的 hidden state）会根据 gating 结果分发到其他 EP rank 上的 expert。这个过程通过 all-to-all 通信完成。
只开启EP时，GPU会拥有多个完整的专家，每个矩阵计算都是完整的，是大矩阵计算。
不同情况各有优劣。

EP

红色和绿色方块表示非-MOE层，e0、e1表示MOE层（总共4n个专家）

一共16块GPU：

ep_world_size = 4：表示我们希望用4块GPU装下一套完整的专家。确定这个数值后，我们就能确认ep_groups
local_expert_num：expert_num / ep_world_size，其中expert_num表示每层专家的总数。
- 假设每层专家数量是4，那么1块gpu上就放一个专家；
- 假设每层专家数量是8，那么1块gpu上就放2个专家。
- 所以图中的e0等符号并不绝对表示这里只有1个专家，只是对local_expert的统称。
ep_dp_world_size：类比于non-MoE层，MoE层同样也有数据并行的概念。例如图中[g0, g4, g8, g12]上都维护着e0，所以它们构成一个ep_dp_group。这个group的作用是当我们在计算bwd时，它们之间是需要做梯度的allreduce通讯的，我们会在下文详细图解这一点。另外需要注意的是，构成ep_dp_group的条件不仅是e相同，还需要每个e吃的batch的数据不同（类比于一个普通的dp_group，组内的每张卡吃的是不同的小batch）。现在你可能无法具象化感受这点，我们在后文将ep+tp+dp并行的时候再细说。
ep_tp_world_size：类比于non-MoE层，MoE层同样也有张量并行的概念，即一个专家可以纵向切割成若干份.

代码实现

参考DeepSeek-VL2实现，当前代码实现有几个特点：

EP=x，就将world_size拆分成几份，每份里的机器上拥有 experts/EP 的专家数的副本。
self.shared_experts = MLP, forward里和路由专家结果相加。

ZeRO¶

ZeRO通过在多个设备上分片优化器状态、梯度和参数来减少每个设备的存储需求。
Zero 优化方法有三个层次，分别是 ZeRO-1、ZeRO-2 和 ZeRO-3。¹³
它们是由微软提出的 ZeRO（Zero Redundancy Optimizer） 优化技术的不同阶段，旨在减少大规模分布式训练中的内存占用。

¶

1. ZeRO-1：优化器状态分区¶

目标：减少优化器状态的内存占用。
实现方式：将优化器状态（如动量、梯度方差等）分布在不同的 GPU 上，而不是在每个 GPU 上保存完整的副本。
优点：显著减少内存占用，同时通信开销较小。
适用场景：适合中等规模的模型训练。

2. ZeRO-2：梯度分区¶

目标：进一步减少梯度存储的内存占用。
实现方式：将梯度分区存储在不同的 GPU 上，每个 GPU 只保存一部分梯度。
优点：内存占用进一步降低，但通信开销有所增加，因为需要在反向传播后聚合梯度。
适用场景：适合大规模模型训练。

3. ZeRO-3：参数分区¶

目标：最大化内存节省，支持超大规模模型训练。
实现方式：将模型参数分区存储在不同的 GPU 上，每个 GPU 只保存一部分参数。
优点：内存占用大幅降低，可以训练非常大的模型，但通信开销最大，因为需要在每次前向和反向传播时聚合参数。
适用场景：适合超大规模模型训练（如 GPT、BERT 等）。

对比总结¶

特性	ZeRO-1	ZeRO-2	ZeRO-3
分区对象	优化器状态	梯度	模型参数
内存节省	中等	较大	最大
通信开销	最小	中等	最大
适用场景	中等规模模型	大规模模型	超大规模模型

分层ZeRo¶

分层Zero（Hierarchical Zero）是一种用于大规模AI模型训练的优化方法，旨在解决传统Zero（Zero Redundancy Optimizer）在大规模分布式训练中的局限性。
它通过分层通信和计算优化，提升训练效率和扩展性。

核心思想是将计算和通信任务分层处理，减少通信开销，提高资源利用率。具体包括：

分层通信：将通信任务分为多个层次，优先在低层次（如节点内）完成，减少高层次（如跨节点）的通信。
分层计算：将计算任务分层处理，优先在低层次完成，减少高层次的计算负担。

主要特点

减少通信开销：通过分层通信，降低跨节点通信频率，提升效率。
提高资源利用率：分层计算使资源分配更合理，减少闲置。
增强扩展性：优化后的方法更适合大规模分布式训练，支持更多计算节点。

异构系统的并行¶

人们思考为什么 CPU 内存没有被用于分布式训练。

AI Training Parallelism

数据并行 DP¶

张量并行 TP¶

序列并行 SP¶

两条技术路线¶

演进脉络¶

反向、重计算与卸载¶

先分清三种内存策略¶

Ring Self-Attention¶

Megatron SP¶

DeepSpeed-Ulysses¶

Ring Attention¶

Megatron CP¶

USP¶

反向与重计算¶

AllGather CP¶

Ulysses-Offload / FPDT¶

反向与卸载¶

与 checkpoint 叠加¶

FSDP2 迁移¶

ALST¶

FSDP2 迁移¶

MindSpeed-MM Chunk Layer¶

起因与一句话思路¶

前向与跨块状态¶

反向、重计算与卸载¶

与 USP、FPDT 的差异¶

显存与“任意长度”的边界¶

工程联评¶

选型建议¶

流水线并行 PP¶

GPipe¶

1F1B¶

PipeDream¶

VPP¶

DualPipe¶

混合并行¶

专家并行 MoE(EP)¶

ZeRO¶

¶

1. ZeRO-1：优化器状态分区¶

2. ZeRO-2：梯度分区¶

3. ZeRO-3：参数分区¶

对比总结¶

分层ZeRo¶

异构系统的并行¶

参考文献¶

评论