multi_card_cp

鱿鱼圈 Lv4

FlashKDA 多卡 Context Parallelism (CP) 数学原理详解

本文面向"小白"读者,从最基础的概念讲起,一步步推导多卡 CP 的 forward 和 backward 数学原理,重点讲清线性分解这个核心。 每个公式都配文字解释,最后用可手算的数值例子把整条链走一遍。

阅读顺序建议:§1 直觉 → §2 记号 → §3 forward → §4 backward → §5 跨卡编排 → §6 例子。 如果只想看结论,直接跳 §7 一页速查表。


1. 先建立直觉:CP 到底在切什么?

1.1 一句话版本

FlashKDA 是一个带循环状态的注意力算子。你可以把它想象成一个"读书做笔记"的过程:

  • 从左到右一个 token 一个 token 地读;
  • 每读一段(一个 chunk = 16 个 token),就更新一次"笔记本" (状态矩阵);
  • 笔记本 会一直往后传,后面的输出依赖前面积累的笔记。

问题是:这个"往后传"是串行的。序列长到 64K、256K 时,GPU 上千个核心大部分在干等前一段算完。

Context Parallelism(上下文并行) 就是把这条长序列沿时间维 切成好几段,让不同的段(不同的 GPU / 不同的 SM)同时开工。

  • 单卡 CP:在一张 GPU 内部,把序列切成多个 sub-segment,让空闲的 SM 并行跑(见 cp_optimization.md)。
  • 多卡 CP:把序列切到多张 GPU上,每张卡(rank)负责一段 token,卡之间只传很小的"边界信息"。

本文讲的是多卡。

1.2 切开之后的麻烦:接缝处的状态

假设有 3 张卡,序列 切成三段:

1
2
3
    rank 0            rank 1            rank 2
[----------) [----------) [----------)
0 P1 P1 P2 P2 T
  • forward 的麻烦:rank 1 要正确工作,需要知道"读到 时笔记本长什么样"——而这个笔记本是 rank 0 算出来的。状态 从前往后流。
  • backward 的麻烦:反向传播时,梯度 从后往前流。rank 1 的右边界梯度,来自 rank 2 的左边界。

CP 的全部技巧,就是如何用最小的通信,把这个"接缝处的状态/梯度"接好。

而 FlashKDA 能把通信压到极小的根本原因,是下面要讲的线性性质:状态对初始值是线性的,所以接缝可以用一个小矩阵一次接好,不用真的把整段重算一遍。

1.3 为什么 KDA 特别适合 CP?

普通 softmax 注意力有因果三角,切开后每段计算量不均(前面的段轻、后面的段重),需要复杂的负载均衡(zigzag 重排)。

KDA 是线性注意力:每个 chunk 的计算量恒定。所以等长切分 = 天然负载均衡,实现上省掉一大块麻烦。


2. 记号与基本设定

先把符号定清楚,后面推导都用这套。

符号 含义 形状
序列总长度(token 数) 标量
head 数 标量
key / value 维度,FlashKDA 固定 标量
chunk 大小,固定 标量
rank 数(GPU 数) 标量
状态矩阵(“笔记本”)
进入第 个 chunk 前的状态
初始状态
query/key/value/gate 输入
beta 门控(logit)

说明:下面所有推导都是**“每个 head 独立”**的。 个 head 之间没有耦合, 所以我们只写一个 head 的公式,实际代码对 维并行即可。同理 batch 维在 varlen 下 。

2.1 一个 chunk 的 forward(来自 autograd.py)

FlashKDA 每个 chunk(16 个 token)做的事,数学上是(全部 fp32 计算):

其中 ,, 是 sigmoid。

别被吓到。对 CP 而言,这一大堆公式里我们只关心一件事:

状态 是怎么从 变成 的?而且这个变换对 是不是线性的?

下一节就把这一坨压缩成一个干净的仿射式子。


3. Forward 的线性分解

3.1 把一个 chunk 写成仿射变换

看 §2.1 最后两行,把含 的部分和不含 的部分分开。

里含 :,所以

代入状态更新 :

于是每个 chunk 就是一个仿射变换:

  • 是 矩阵,叫转移矩阵(transition matrix),只跟 有关,跟 无关。
  • 是 矩阵,是数据项,含 。

小白理解: 是"旧笔记怎么衰减/改写", 是"这一段新读到的内容"。 新笔记 = 旧笔记按 变换一下 + 新内容 。

3.2 一整段的链:

一个 rank 负责一段,里面有 个 chunk(编号 )。反复套用仿射式:

定义两个量:

则整段的输入输出关系是一个干净的仿射式:

这就是 forward 线性分解,是整个 CP 的数学基石。含义:

  • (转移矩阵):初始状态 穿过这一整段后,被"缩放/旋转"成 。它只跟这一段的 有关,跟 都无关。
  • (数据驱动终态):假设 ,纯靠这一段的 数据能积累出的终态。

为什么这个分解如此重要? 因为它把"整段递推"这个串行过程,浓缩成了两个可以提前、独立算好的量 。 一旦有了 ,任意给一个 ,我都能一步(一次矩阵乘 + 一次加法)算出 , 不用把整段 个 chunk 重新串一遍。跨卡时,卡之间只需要传 这两个小矩阵。

3.3 怎么算 和 ?(state_only kernel)

state_only kernel 分别从两个不同的初值出发,跑同一条递推,就能"测量"出 和 :

  • 算 :初始 ,跑 ,终态就是 (因为 )。
  • 算 :初始 (单位阵),且去掉数据项 (因为 跟 无关), 跑 ,终态就是 (因为 )。

对应 fwd_state_only.cuh 里的 CalcMt 模板参数:CalcMt=false 从 0 出发算 , CalcMt=true 从 出发、更新式带减号( 的减号)算 。

3.4 关键简化:强 decay 下 (以真实 warmup kernel 为准)

强 decay 时, 是一串 连乘、指数级衰减到接近零矩阵,于是

含义:decay 强时,初始状态 的影响被彻底遗忘,终态几乎只由本段数据决定。 这就是为什么生产环境下 CP 的跨段修正大多数段几乎不需要 :只要往回扫少数几个 chunk, 残差就压到可忽略,跨卡只传一个边界量 即可。

但"往回扫几个 chunk 就够"这件事不能拍脑袋定死,得逐段、逐 head、逐维度地实测。 真正决定它的是 csrc/warmup_chunks.cu(get_warmup_chunks_kernel),逻辑如下—— 注意它和旧文档 cp_optimization.md 里"mean over D + 0.25 + 中心极限"的说法不一样, 以下这版是当前 kernel 的真实实现。

(1) 衰减在 空间累加。 每个 chunk 往回扫时,先按维度算 sigmoid 门控,再乘一个 已经含 因子的 gate_scale:

越扫越负

于是残差幅度是 ,不是 ——因为 已经折进 gate_scale 把自然底换成了 2 底:,两者相等, 但若既写 又套 就重复乘了一次,会把数量级算错。 (旧文档写的 就是犯了这个双重计数: 已经是 单位,正确残差是 ,不是 。结论""没错,数值差了 8 个量级。)

(2) 维度上取 min,不是 mean。 kernel 对 个通道逐个算 , 然后取 (kernel 注释原话:“min over D … to guarantee the weakest dimension converges”)。 越小 decay 越接近 0 衰减越慢 残差越大。取 min 就是盯住最难收敛的那个维度, 是保守下界;用 0.5/mean 会低估最慢维度、提前判定收敛,是不安全的。

(3) head 上取 max。 各 head 的 g_cumsum 做树形归约取 max(最不负 = 最慢的 head), 只有当最慢的 head 都收敛才算整段收敛。

(4) 阈值 + 逐段自适应 + fallback。 往回扫直到 (例如 残差 ),此时记 num_warmup = c+1(该段只需这么多 chunk)。 若扫完整段仍没到阈值(weak decay, 接近 0,或序列太短),置 fallback=true—— 这一段不能用 近似,要老老实实算精确 走完整前缀/后缀 scan。

所以"强 decay 下 "不是一个全局假设,而是 kernel 逐段实测出来的结论: 强 decay 段自动只扫少数 chunk、丢掉 ;弱 decay 段自动 fallback 补精确 。近似与精确的边界由数据决定,误差始终有界(见 §8)。


4. Backward 的线性分解(forward 的镜像)

这一节是本文重点。我们要证明:backward 也有一个和 forward 完全对称的线性分解, 只是方向反过来(从后往前),转移矩阵用转置 。

4.1 什么是 ?

训练时有一个 loss 。定义

即"进入第 个 chunk 的状态 "对 loss 的梯度,形状还是 。 反向传播要做的就是:已知后面的 ,往前推出前面的 ,同时收集各参数的梯度。

4.2 的递推:从后往前的仿射链

状态 在 forward 里影响了两个地方(看 §2.1):

  1. 本 chunk 的输出 ,而 里也含 (通过 )。
  2. 下一个状态 ,进而影响之后所有 chunk。

按链式法则, 是这两条路径的贡献之和:

路径:本地回传路径:状态递推回传

第二项 从哪来? 因为 ,这是一个线性映射,雅可比矩阵就是 。 梯度反向穿过线性映射 时,。所以 的梯度 回传到 就是 。注意是转置 。

第一项 是什么? 它是"本 chunk 输出 收到的梯度 ,经由 直接回传到 的部分"。 展开(来自 两条子路径):

  • 第一块 :来自 ,直接对 求导。
  • 第二块:来自 ,而 经 又依赖 ,把这条路回传回来(带负号,因为 里是 )。

的关键性质:它只依赖本 chunk 的 和 forward 已经存好的 workspace(),完全是本地量,不含跨 chunk 的 。

于是 backward 就是一条从后往前的仿射链:

把它和 forward 的 对照:

  • forward:往后走,用 ,数据项 。
  • backward:往前走,用 ,本地源项 。 一模一样的结构,只是方向和转置不同。这就是"镜像"。

这条链对应 bwd_kernel2.cuh 的主循环:for t = t_tiles-1 downto 0(倒序), Step 8b "dS 向前传播(为上一个 chunk)"就是在算 , Step 5/6 里 驱动的项就是 。

4.3 一整段的 backward 分解:

现在把这条链在一段 上展开(记右边界梯度 , 左边界梯度 ):

注意第一项的矩阵连乘:

这正是 forward 里那个同一个 的转置! 定义左边界的本地源

这个求和号是怎么来的?( 去哪了)

很多人到这里会困惑:上面展开时明明写的是一个个具体的 , 怎么最后一行突然变成了带哑变量 的 ?

答案: 并没有"被替换掉",而是全部被求和号 收了进去。 里的 是求和的哑变量(dummy index),它依次取遍 , 所以 就是 那一项、 就是 那一项……一个都没丢。

把递推 从 一路代到 , 本地源项那部分逐行长这样:

观察规律:带 的那一行,前面乘的正好是"从 连乘到 "的转置矩阵积 。 ( 时是空乘积 ,所以 前面没有 ,对得上。)于是这一整串就压成一个 :

求和里 取值 展开成的具体项

所以 只是"通项/一般项"的写法, 已经把 到 全部含进去了, 就是这一整串的和。

和 forward 的 是同一个套路: 也是把 用一个 收起来。backward 完全镜像,只是把 换成本地源 、 换成 ,连乘方向从" 在 右边"变成" 在 左边"(因为链方向反了)。

就得到 backward 线性分解:

含义完全镜像 forward:

  • :右边界梯度 穿过整段传到左边界,被 变换。 它就是 forward 那个 的转置,免费复用,不用重新算。
  • :假设 ,纯由本段 数据积累、传到左边界的梯度。 它是 backward 版的"数据驱动项",对应 forward 的 。

4.4 对照表:forward ↔ backward 完美镜像

维度 forward backward
边界量 状态 状态梯度
链方向 前缀(前 → 后) 后缀(后 → 前)
单 chunk 递推
转移矩阵 (转置)
数据项 (含 ) (含 )
整段分解
转移分量 (同一个 转置)
数据分量 ( 的终态) ( 的左边界梯度)
入口边界 只在第一段(rank 0) 只在最后一段(rank )
出口 在最后一段 在第一段
强 decay

一句话总结:backward 就是把 forward 的图倒过来走一遍, 把 换成 、 换成 、 换成 、“前缀"换成"后缀”、 (左入口)换成 (右入口)。数学结构一模一样。

4.5 强 decay 下 backward 也退化

backward 用的是同一个 的转置,所以 §3.4 的 warmup 结论原封不动搬过来: warmup kernel 逐段测出某段强 decay()时 ,转置后 :

含义:右边界梯度 穿过强 decay 段后被彻底压没, 每个 rank 的反向几乎只依赖本地 ,跨卡耦合可忽略。 和 forward 的 完全对称,跨卡开销同样极低。 而 fallback 段(弱 decay)会算精确 ,转置后 精确参与后缀 scan,误差始终有界。


5. 跨卡编排:把线性分解用起来

现在把序列切成 个 rank,rank 负责 。相邻 rank 的接缝处, 就是 forward 传的边界状态 (rank 的 final = rank 的 init), 它的梯度记 。每个 rank 用 §3、§4 的整段分解,把自己浓缩成 和 。

5.1 Forward:跨卡前缀 scan

边界条件:(只有 rank 0 有,其余从 0)。 用 §3.2 的分解,跨 rank 链式求每个 rank 的正确初始状态:

  • = rank 的正确 ;
  • = 全局 final_state(在最后一个 rank)。

这是一个前缀 scan(从 rank 0 往后)。

5.2 Backward:跨卡后缀 scan

边界条件:(只有 rank 有,其余从 0)。 用 §4.3 的分解,跨 rank 链式求每个 rank 的右边界梯度:

  • = rank 的右边界梯度(反向的输入);
  • = 全局 d_initial_state(在 rank 0)。

这是一个后缀 scan(从 rank 往前)。

5.3 两种落地路线(forward / backward 都适用)

路线 A(精确串行,做对拍基线用) rank 按顺序跑,前一个 rank 的边界算完传给下一个。正确但完全串行, 退化成没有并行收益,只用来验证正确性。

路线 B(并行,推荐,即 fla 的 build_cp_context 思路) 以 backward 为例:

1
2
3
4
5
6
7
8
1. 每个 rank 本地反向,右边界 dS 先置 0
-> 同时得到 d_ht_r(本地源)和"暂定"的本地 dq/dk/dv/dg/dbeta。
mt_r 直接复用 forward 存下的 state_only 结果,不重算。
2. all-gather 所有 rank 的 (d_ht_r, mt_r)。
通信量 = R * H * D * D 的 fp32,非常小。
3. 每张卡本地冗余做后缀 scan(§5.2 的公式),得到真正的 dS_{r+1}。
4. 若 mt_r^T 可忽略(强 decay)-> dS_{r+1} 贡献≈0,本地梯度已正确,直接用;
否则用修正后的右边界 dS_{r+1} 重跑一次 K2_bwd(或对本地梯度做线性修正)。

forward 的路线 B 完全对称:本地跑 state_only 得 → all-gather → 本地前缀 scan 得 → 用 当 并行跑完整 forward。

5.4 参数梯度的 all-reduce

反向产出的梯度分两类:

梯度 性质 跨卡处理
per-token,按 rank 天然分片 不需要 reduce,各卡各自持有自己 token 的梯度
, 全局参数,forward 里被所有 token 用 atomicAdd 累加 每卡只累加了自己 token 的贡献 → all-reduce SUM

为什么 、 要 all-reduce? 因为 (每 head 一个)和 (每 head、每维一个)是所有 token 共享的参数。 完整的梯度 = 所有 token 贡献之和。切卡后每张卡只看到一部分 token, 算出的只是部分和,必须跨卡加起来才是完整梯度。 而 等是每个 token 各自一份,rank 之间不重叠,所以不用加。

5.5 插曲:集合通信原语(all-gather / all-to-all / all-reduce)

上面反复用到 “all-gather”、“all-reduce” 这些词。它们是集合通信(collective communication)原语 —— 多张卡之间约定好的标准数据交换动作。这里用 4 张卡 (rank 0/1/2/3)、每张卡起始各持一块数据的例子讲清楚,重点是 CP 里最常用的 all-gather 和容易混淆的 all-to-all 的区别。

all-gather:人人贡献一块,人人拿到全部(聚合,结果一致)

1
2
3
4
5
        起始                       结束(all-gather 后)
rank0: [A] -> [A][B][C][D]
rank1: [B] -> [A][B][C][D]
rank2: [C] -> [A][B][C][D]
rank3: [D] -> [A][B][C][D]
  • 每张卡贡献 1 块,最后每张卡都拿到全部 4 块(拼接),所有卡结果完全相同。
  • CP 里的用途:all-gather 各 rank 的 / , 让每张卡都凑齐所有 rank 的边界量,好在本地各自冗余跑那条前缀/后缀 scan(§5.3 路线 B)。

all-to-all:每卡按目标把数据切块分发(转置,结果各异)

起始时每张卡持有 4 块,第 块是"要发给 rank 的"(记 = rank 上发往 的块):

1
2
3
4
5
        起始                        结束(all-to-all 后)
rank0: [A0][A1][A2][A3] -> [A0][B0][C0][D0]
rank1: [B0][B1][B2][B3] -> [A1][B1][C1][D1]
rank2: [C0][C1][C2][C3] -> [A2][B2][C2][D2]
rank3: [D0][D1][D2][D3] -> [A3][B3][C3][D3]
  • rank 的第 块送到 rank ,相当于把这个 数据矩阵做一次转置。
  • 每张卡发出去的是不同的块,收到的也来自各卡的不同块,所有卡结果各不相同。
  • 典型用途:MoE 的 token dispatch(按专家路由分发)、序列并行 张量并行重排、分布式转置。 FlashKDA 的 CP 边界交换用不到 all-to-all(我们要的是"人人凑齐全量",不是"按收件人分发")。

核心区别一览

all-gather all-to-all
每卡发的内容 同一份(自己那块) 按目标切成 份,份份不同
每卡收的结果 全量拼接,人人相同 各收各的,人人不同
本质 收集 / 广播式聚合 转置 / 重分布
类比 “把大家的名片收齐,人手一套” “每人把一叠信按收件人分发出去”

顺带区分 all-reduce 和 reduce-scatter

  • all-reduce = all-gather 再加一步逐元素归约(求和 / 求 max 等), 最后人人拿到的是同一个归约结果(不是拼接)。§5.4 里 、 就是用 all-reduce SUM 把各卡的部分和加成完整梯度。
  • reduce-scatter = all-reduce 的前半段:归约后每卡只留结果的一片(不是全量)。
1
2
3
4
all-gather:      各卡[片] -> 各卡[全量拼接]          (结果一致, 无计算)
all-reduce: 各卡[片] -> 各卡[逐元素求和的全量] (结果一致, 带归约)
reduce-scatter: 各卡[全量] -> 各卡[求和后的一片] (结果分片, 带归约)
all-to-all: 各卡[N份] -> 各卡[转置后的N份] (结果各异, 无计算)

一句话记忆:all-gather 收集拼接、all-reduce 收集再求和、all-to-all 是转置分发。 CP 边界交换用 all-gather,参数梯度用 all-reduce SUM,两者都用不到 all-to-all。

5.6 通信量小结

  • 每次跨卡只传 个 边界矩阵。 时约 (fp32), 相对 activation 微不足道。
  • 强 decay 生产参数下 ,连 都不用传,只传 (一个边界量)。
  • 主计算(完整 forward / backward)全部本地并行,跨卡只有轻量 warmup + 边界 scan + 两个参数梯度的 all-reduce,开销和单卡 CP 的 ~5–10% 一个量级。

6. 手算例子:把整条链走一遍

为了让链看得见摸得着,我们把状态退化成标量()。此时 是一个数, 转移矩阵 也是一个数(),转置 。取 个 rank, 每个 rank 整段归并成一次转移。给定下面这组数(假想已由 kernel 算好):

rank 转移 (=) forward 数据项 backward 本地源
0 0.5 2.0 0.3
1 0.4 1.0 0.5
2 0.3 3.0 0.7

6.1 Forward 前缀 scan(从 rank 0 往后)

初始状态 。用 :

的的,卡卡传卡卡传全局

每张卡拿到自己的 ()后,就能并行跑各自那段完整 forward。

6.2 Backward 后缀 scan(从 rank 2 往前)

右边界梯度 。用 (标量下 ):

卡算出,回传给卡卡算出,回传给卡,卡输出

每张卡实际干的事(路线 B):

1
2
3
4
卡2: 本地反向(右边界 = dfinal = 1.0) -> 直接得 dS_2 = 1.0, 本地 dq/dk/... 已精确; 传 dS_2
卡1: 本地反向(右边界 = 收到的 dS_2 = 1.0) -> dS_1 = 0.9; 传 dS_1
卡0: 本地反向(右边界 = 收到的 dS_1 = 0.9) -> dS_0 = 0.75 = d_initial_state
参数梯度: 三卡各自的 dA_log / ddt_bias 求和 all-reduce

6.3 验证:CP 版和单卡版数值完全一致

如果不切卡,单卡把这 3 段连成一条链一次反向, 一路传到头, 结果也正好是 ,而且每个 chunk 的局部梯度逐一相同。

为什么是精确相等,而不是近似? 因为 是精确的代数恒等式(§4.3 严格推出来的), 不是任何近似。所以多卡 CP 反向的结果和单卡反向逐比特意义上等价(忽略浮点重排的舍入), 不引入任何数学误差。

6.4 强 decay 演示:跨卡项被压没

把 全换成 (模拟 的强 decay), 不变:

forward:

backward:

跨卡回传项全被 压掉,每张卡的 forward/backward 退化成互相独立、 只用本地数据,一次并行即得(近似精确)结果。这就是路线 B 在生产参数下几乎零跨卡开销的根本原因。

6.5 从标量回到矩阵()要注意什么

上面的例子是标量 。真实 时:

  • 都是 矩阵, 不再交换,连乘顺序不能颠倒。
  • forward 前缀 scan 是 (矩阵乘向量化到 ); backward 后缀 scan 是 ,注意转置。
  • 强 decay 下 是接近零矩阵, 也接近零矩阵,退化结论不变。
  • 每个 head 独立,实际是 份 的 scan 并行。

7. 一页速查表

Forward(前缀 scan,前 → 后)

Backward(后缀 scan,后 → 前)

单 chunk

跨卡通信

  • forward:all-gather → 本地前缀 scan → 用 当 并行 forward。
  • backward:all-gather → 本地后缀 scan → 用 修正并行 backward。
  • 参数梯度: 天然分片不 reduce; 需 all-reduce SUM。

强 decay()退化

跨卡耦合可忽略,每张卡 forward/backward 近似独立并行,开销 ~5–10%。

工程要点(多卡特有)

  1. conv1d halo:段首需前一 rank 段尾 个 token,一次小通信。
  2. 切分点对齐到 chunk 边界(),避免破坏 workspace 的 per-tile 布局。
  3. varlen:cu_seqlens 要叠一层 rank→段 映射,处理一条序列被切到多卡。
  4. workspace / all_states 是本地前向产物、本地反向消费,不跨卡(这是通信只剩边界量的前提)。
  5. 可与单卡 CP 叠加:rank 内再切 sub-segment,两层 scan 可组合(先卡内、边界项再进卡间)。

8. 训练中能用吗?会不会有精度问题?

结论:能用,而且在数学上和单卡训练等价——只要 warmup 的 fallback 打开。

分两层看:

8.1 线性分解本身是精确恒等式,不引入近似

(forward)和 (backward) 都是严格代数恒等式(§3、§4.3 一步步推出来的),不是任何近似。 所以多卡 CP 的前向/反向结果,与单卡跑完整序列逐比特意义上等价(只差浮点求和重排的舍入)。 训练时反传给 的梯度,和单卡是同一套值—— CP 只是把同一份计算按 切开并行,梯度语义不变,优化器看到的东西完全一样。

8.2 唯一的近似来自"丢 ",而它被 warmup + fallback 兜底

真正可能引入误差的只有一个动作:强 decay 段把 当 0 丢掉、只传 。 这一步的误差恰好是被丢掉的 项,幅度 。 get_warmup_chunks_kernel 用 min-over-D(最慢维度)+ max-over-head(最慢 head) 的保守估计, 只在 (残差 ,如 )时才敢丢, 且丢不掉时 fallback=true 强制算精确 。所以:

  • 训练想要零近似:把 threshold 调得足够负(或直接强制 fallback / 走 §5 的路线 A 精确 scan), 此时 CP 训练与单卡没有任何数学差异,可放心用。
  • 训练想要跨卡近零开销:用默认 threshold,误差被 有界压住, 且只落在强 decay( 本就被遗忘)的段上——被丢的项本身量级就 梯度噪声/bf16 舍入,实践中对收敛无影响。

一句话:丢 的近似发生在"它本来就趋于零"的地方,误差有解析上界;不满足条件就自动退回精确路径。

8.3 精度注意事项(和单卡内核同源,CP 不新增)

CP 不改变内核的数值精度特性,训练需要注意的与单卡一致:

  1. 状态/中间量 bf16,累加走 fp32:backward 的 workspace、 累加、INV 都在 fp32 下做, 用 fp32 atomicAdd 累加(跨卡再 all-reduce SUM),避免 bf16 累加掉精度。
  2. 切分点对齐 chunk 边界():错位会破坏 per-tile workspace 布局,属于正确性问题不是精度问题,必须对齐。
  3. 必须 all-reduce SUM:这两个是所有 token 共享的参数梯度, 漏掉跨卡归约会让梯度偏小(等价于只用了本卡那段 token),这是 CP 训练最容易踩的坑,和精度近似无关。
  4. fallback 段的开销:weak decay( 接近 0)会大量触发 fallback → 算精确 → 跨卡通信与计算变重。这是"精度换开销"的旋钮,不是精度损失。

附录:符号速记

符号 forward 含义 backward 对应
状态(笔记本) 状态梯度
单 chunk 转移矩阵
单 chunk 数据项(含 ) 本地源(含 )
整段转移矩阵
整段数据驱动终态() ()
初始状态(左入口) (右入口)
终态(右出口) (左出口)
  • 标题: multi_card_cp
  • 作者: 鱿鱼圈
  • 创建于 : 2026-09-12 22:20:00
  • 更新于 : 2026-09-12 22:26:40
  • 链接: https://yuyanqi.com/2026/09/12/multi_card_cp/
  • 版权声明: 本文章采用 CC BY-NC-SA 4.0 进行许可。
评论
目录
multi_card_cp