AIGC 基本功|视频 DiT 里的 3D RoPE 与分辨率外推-3D-RoPE 视频 DiT 里的 3D RoPE 与分辨率外推 所属方向：注意力与核心零件 ｜ 难度：进阶 ｜ 前置知识：RoPE 的原理与实现 关键词：3D RoPE、时空位置编码、视频 DiT、分辨率外推、逐轴位置插值、帧数外推 01. 为什么需要它 先看一对位置：视频格点 A=(0,0,0)，B=(0,1,0)。B 在 A 的正下方一格。将每帧按行展开，再按时间串接，普通 1D RoPE 看到的序列位置是 m=tHW+hW+w。当每帧宽 W=4 时，A 到 B 的序列距离是 4；仅把宽改为 W=8，距离变成 8。两块仍上下相邻，旋转角却按 1D 的位置差变了。时间上的相邻帧也会受 HW 改变影响。固定画幅下，这种展开次序可以用，模型也可能学到空间规律；问题是在画幅或时长改变时，序列距离不再稳定代表同一种时空关系。这里的 1D RoPE 是为了隔离坐标问题而构造的诊断基线，不是在声称所有早期视频模型都按这一公式编码位置。 3D RoPE 给 A、B 保留 (t,h,w)，所以无论宽度是多少，两者都只在高轴相差一格。它解决的是展平编号把画幅尺寸混进位置差的问题；画质外推还涉及训练分辨率、物体尺度和注意力成本。后文用 CogVideoX 的实验 分开讨论位置表直接扩展与逐轴插值，不把“坐标可计算”误当成“视频一定生成得好”。 02. 最小可用理解 基础 RoPE 已解释二维旋转与相对位置性质，这里直接看视频新增的两步： 换坐标：展平版给每个 token 一个 m，相对位移是 Δm；三轴版保留 (t,h,w)，相对位移是 (Δt,Δh,Δw)。图 1 的上下邻居在画幅变宽后，只有三轴版仍把它识别为“高轴相差一格”。 分特征与频率：一个注意力头的特征分为时间、高、宽三段，每段用自己的坐标和频率表生成位置角。三段随后仍在同一次注意力内积里汇总；“3D”说的是坐标轴数，不是把视频物体在欧氏空间旋转。 图 1：固定 A、B 的时空坐标，只改每帧格点宽度。展平位置差从 4 变 8，三轴位移始终是 (0,1,0)；真实模型的内容特征也会随画幅变化。 具体比例由模型选择。CogVideoX 给时间 1/4、高 3/8、宽 3/8；64 维头对应 (16,24,24)，不是所有 3D RoPE 的固定标准。按这个分配，高轴的坐标映射可以单独调整，而时间、宽轴保持原样。实际换模型还要核对每轴维度、坐标单位、频率基数、配对次序和文本 token 的处理，不能只照搬比例。 图 2：64 维注意力头的示例。每个轴的维度都成对构成旋转平面；三段共享一个 token，但不会互相混成单一序列编号。图由文末代码生成。 03. 数学推导 3.1 从视频格点到三组旋转 设视频 patch token 为 $z_{t,h,w}$。$t$ 是压缩后的时间格点，$h$ 和 $w$ 是 patch 后的高、宽格点；它们不是原始帧号或像素坐标。以下只讨论注意力打分前，如何给对应的 query/key 构造位置相位。 先写出基础版到底把什么送进旋转。固定一个按行优先展开的视频网格 $(T,H,W)$，位置编号 $m=tHW+hW+w$。任意两格 $p=(t,h,w)$、$r=(t',h',w')$ 的编号差为： $$m_r-m_p=(t'-t)HW+(h'-h)W+(w'-w)$$ 上式说明 1D 并非丢了时空信息：固定 $H,W$ 时，每个格点仍有唯一编号。但同样一个高轴位移会乘上当前的 $W$，同样一个时间位移会乘上当前的 $HW$。换画幅后，“向下 1 格”的位置相位因此变化；这只比较固定内容特征的位置项，不推断整个模型的注意力分数或画质。 令头维度 $d=d_t+d_h+d_w$，每段维度都必须为偶数。把 $q$ 和 $k$ 拆为 $q^{(t)},q^{(h)},q^{(w)}$ 与对应的三段 $k$。对轴 $a\in\{t,h,w\}$ 的第 $i$ 个二维旋转平面，选频率： $$\omega_{a,i}=\Theta^{-2i/d_a},\qquad i=0,1,\ldots,d_a/2-1$$ $\Theta$ 是频率基数，CogVideoX 的源码默认 10000；$d_a$ 是该轴分到的维度，而不是整个头维度。时间与空间段各自从最高频重新开始，不是把一条 1D 频率数组切成三段。沿用先修文章的旋转记号 $R$，用 $R_a(p_a)$ 表示轴 $a$ 上所有旋转平面的组合。完整位置操作是三个互不混合的块拼接： $$\widetilde q_{t,h,w}=\operatorname{diag}(R_t(t),R_h(h),R_w(w))q_{t,h,w}$$ key 用同一个位置规则得到 $\widetilde k$。分轴的是位置相位的构造；注意力仍对所有可见视频 token 计算一次打分。 可以把区别压到一行：1D 的相对角是 Δm × 频率，其中 Δm=Δt·HW+Δh·W+Δw；3D 的三组相对角分别是 Δt × 时间频率、Δh × 高频率、Δw × 宽频率。三组贡献在同一次内积里相加。这里的 Δ 都是 token 格点差，不是秒数或像素差。 3.2 三个位移怎样汇入一次打分 基础 RoPE 已推过两个旋转角相减的恒等式。这里新增的是：每个轴只用自己的坐标差，整头把三段贡献相加： $$s(p,r)=\sum_{a\in\{t,h,w\}}\langle q^{(a)},R_a(r_a-p_a)k^{(a)}\rangle$$ $p=(t,h,w)$ 和 $r=(t',h',w')$ 是两个格点。只有高差为 1 时，时间段与宽段的位置差都是零，高段单独改变打分；换宽度不会把这一高差改写成另一个数。公式描述固定 $q,k$ 时的位置作用，文末代码用上下邻居和下一帧两种位移分别检查它。 再看一个跨帧例子：某个物体的投影从 (0,5,5) 到 (1,5,6)，位置差是 (1,0,1)；时间段和宽段的相位会变，高段不会。但静止物体遇到相机运动也可能产生完全相同的屏幕投影轨迹。对这两个例子，3D RoPE 的位置部分完全一样；它没有标记“这是同一物体”，也没有记录相机姿态或两个时间格点之间的秒数。模型仍可从图像内容、条件和后续层学习运动关系，不能把三轴相位本身当作光流或物理速度。 式子还有一个可检查的边界：三段打分相加，位置操作本身没有 Δt×Δw 这样的跨轴乘积项。这不妨碍注意力通过 $q,k$ 的内容组合时空线索；它只说明“分别编码三个轴”并不等于“直接编码一条物体轨迹”。 3.3 直接外推与逐轴插值到底改了什么 假设训练高轴有 $H_0$ 个格点，推理变成 $H_1$ 个格点，零起点时两者的最大坐标分别是 $L_0=H_0-1$、$L_1=H_1-1$。直接外推使用新坐标 $h'=h$，相邻高 token 仍相差 1，角差仍为 $\omega_{h,i}$，但远端坐标超过旧范围。若要让两个端点严格对齐，逐轴位置插值应使用 $h'=h/s_h$、$s_h=L_1/L_0=(H_1-1)/(H_0-1)$；相邻角差变为 $\omega_{h,i}/s_h$，最远坐标正好被压回 $H_0-1$。常见的 $H_1/H_0$ 是大网格下的近似，网格很短时两者不可混写。宽与时间可以分别选择 $s_w,s_t$： $$\phi_{a,i}(p_a)=\frac{p_a}{s_a}\omega_{a,i},\qquad s_t,s_h,s_w>0$$ 这条式子说明“分辨率外推”不是只能一起拉伸三个轴。若只想演示高轴坐标压缩两倍，可以选 $(s_t,s_h,s_w)=(1,2,1)$；若要复现某条真实管线，则应按它的端点、crop 坐标和取样约定计算，不能只看形状比。保持局部细节对应相邻格点的相位差尽量不变；覆盖新范围对应最远端相位别离开已训练的分布。单一全频率插值无法同时严格满足两者。图 3 用 scale=2 隔离这个机制：高度插值后高频振荡慢了一半，低频在这 32 个格点上原本就几乎不动；它不是某个 CogVideoX 推理配置的逐参数复刻。读图时也不要把 cos 的偶然重合当成“两个远处位置相同”。 图 3：高维度 24、基数 10000 的示例。左边用原坐标，右边只把高坐标除以 2；时间与宽并未压缩。图中画的是 cos 相位而非生成质量。 3.4 维度预算不是按视频长宽直接分蛋糕 为什么 CogVideoX 把时间分到 1/4、两个空间轴各分到 3/8，而不是三等分？论文给出了这个比例与三轴独立 RoPE 的实现，却没有证明它对所有数据和画幅最优。可以从编码需求理解：空间中每一帧都包含大量局部邻接与大范围结构，两个空间轴合计拿到较多旋转平面；压缩后的时间格点通常少于空间格点，但仍需能分辨前后帧。这个解释是设计动机的推断，不能代替原论文消融，更不能据此宣布时间维度永远够用。 64 维示例中，时间 16 维对应 8 个旋转平面，高与宽的 24 维各对应 12 个。最高频平面的角增量都是 1 弧度，因为 $i=0$ 时 $\omega=1$，所以“空间分得多”并不是说高、宽的每一步比时间转得快。区别落在可用的频率层级数量和最慢频率：时间最后一组在坐标 15 的角约 0.004743，高度最后一组约 0.003232。两组数字都很小，它们在短网格上几乎不提供相邻区分，却可能给更长距离提供不同的相位尺度。把更多维度分给某轴，得到的是更丰富的频谱，不是对该轴长度的直接承诺。 还要区分物理帧数与时间 RoPE 格点。假设 VAE 把若干帧压缩到一个 latent 时间格点，视频从 49 帧变成 97 帧，RoPE 的时间长度未必简单从 49 变到 97；VAE 的首帧保留或补帧规则甚至会让两个公式对不上。头维度预算必须在真正进入注意力的 token 网格上讨论。读实现时先打印 patchify 输出和 (T,H,W)，再看时间频率有没有被拉伸。 帧率变化不等于时长外推。 同样是把原始帧数翻倍，可能是同一段动作以两倍帧率采样，也可能是原帧率下播放两倍时间。若 VAE 与 patchify 最终都把它们变成两倍的时间格点，整数时间 RoPE 看到的范围一样，却不能仅凭位置角判断“相邻格点代表多少秒”。前者主要考察更密的动作采样，后者还考察更长的主体与场景记忆；测试时间缩放前应固定并记录帧率、实际秒数、latent 时间压缩和条件时间戳。这是视频轴特有的语义问题，空间高宽轴没有对应的播放时钟。 3.5 NTK 变基数为什么不是“插值”的同义词 另一类策略不缩坐标，而是逐轴更改频率基数 $\Theta_a$。若 $\Theta_a$ 变大，$i=0$ 的最高频 $\omega_{a,0}=1$ 完全不变，较低频平面转得更慢；于是近邻分辨率较容易保留，远距离相位覆盖可以加宽。它与 $p_a/s_a$ 把所有频率一起压低不同。NTK-aware、YaRN 等方法在语言长上下文中发展，迁移到视频时还要决定给哪个轴、哪些频段用、是否微调和怎样处理训练画幅。仅把 rope_theta 调大，不能替代逐轴分辨率训练或保证画质。 “频率变慢”也不等于“完全不损局部”。注意力打分汇总多个平面，改低频仍会改变其对近处和远处匹配的贡献。任何外推方法都必须用目标模型、目标画幅与目标时长做消融。 04. 代码实现 code/rope_3d_minimal.py 只依赖 NumPy。核心代码先给时间、高、宽各生成角度表，再广播到 (T,H,W)，沿最后一维拼接成每个 token 的角度；rotate 把相邻两维作为一个二维平面旋转。它保留 float64 以让代数测试更容易观察，不模拟工业模型的 BF16 或 CUDA 性能。 dt, dh, dw = axis_dims(head_dim) at = axis_angles(np.arange(frames), dt, scale=scales[0])[:, None, None, :] ah = axis_angles(np.arange(height), dh, scale=scales[1])[None, :, None, :] aw = axis_angles(np.arange(width), dw, scale=scales[2])[None, None, :, :] angles = np.concatenate( [np.broadcast_to(a, (frames, height, width, d // 2)) for a, d in zip((at, ah, aw), (dt, dh, dw))], axis=-1 ).reshape(frames * height * width, head_dim // 2) cos = np.repeat(np.cos(angles), 2, axis=-1) sin = np.repeat(np.sin(angles), 2, axis=-1) 这里 angles 的最后一维只有 d/2 个旋转平面；repeat(...,2) 才让 cos/sin 与 d 维特征逐元素相乘。把代码复制到本地运行 python3 rope_3d_minimal.py，本次真实输出如下： axis_dims= (16, 24, 24) cache_shape= (24, 64) input_shape= (24, 64) rotated_shape= (24, 64) norm_error= 1.78e-15 same_offset_scores= 2.208374 2.208374 difference= 1.33e-15 vertical_neighbor_1d_offsets= 4 8 vertical_neighbor_1d_scores= 1.678217 3.749697 vertical_neighbor_3d_offsets= (0, 1, 0) (0, 1, 0) vertical_neighbor_3d_scores= 0.419453 0.419453 height_phase_at_15_first_last= [15. 0.003232] height_PI_x2_first_last= [7.5 0.001616] time_phase_at_15_first_last= [15. 0.004743] token_count_13x30x45= 17550 token_count_13x60x90= 70200 dense_attention_ratio= 16.0 naive_score_matrix_fp16_gib_per_head= 0.574 9.179 四条 vertical_neighbor_* 输出固定同一组随机 $q,k$ 与同一对上下相邻 token，只切换展平时的 W=4/8。1D 的 Δm 从 4 变 8，内积分数从 1.678217 变 3.749697；3D 的逐轴距离始终 (0,1,0)，位置项分数都为 0.419453。1D 与 3D 的绝对分数本身不该互比：两套频率配置不同；这个小实验只比较各自规则在宽度变化前后是否稳定。模型实际改宽后 $q,k$ 也会变，所以它不是画质或整层注意力不变的证据。 cache_shape=(24,64) 来自 2×3×4=24 个视频格点。旋转前后范数误差接近浮点舍入，因为旋转矩阵保长度。高坐标 15 的最高频角为 15 弧度，插值后变成 7.5；高轴最慢一组角从 0.003232 变 0.001616。时间角仍是 15 和 0.004743，证明只改高轴没有动时间表。最慢高频与最慢时频数值不同，因为两轴拿到的维度数分别是 24 与 16，指数分母不同。 再看代码里一个表面上很普通的选择：scales 被明确写成 (time,height,width) 三元组，而不是一个全局的 scale=2。这个接口让实验可以只改一个轴，也强迫调用者记录每轴真实扩张比例。假如训练网格是 T0×H0×W0、目标网格是 T1×H1×W1，对格点数大于 1 的轴，端点对齐的起点应分别按 (T1-1)/(T0-1)、(H1-1)/(H0-1)、(W1-1)/(W0-1) 计算；只有一个格点的轴没有可缩放的坐标跨度，应单独保持在零点。但这些比例仍不能机械照搬：训练数据往往包含多个时长和长宽比，没有单一的 T0,H0,W0。若宽度原本就见过目标值，宽轴可以维持原坐标；若时间增长跨越训练分布，更慢的低频策略可能更值得先测试。参数名称应表达“坐标映射”，而不是含糊写成“画质增强”。 norm_error、same_offset_scores 和宽度变化前后的 3D 分数在脚本里都有断言；若配对、广播或相对位置性质被改坏，脚本会直接失败，而不是只打印一串看似正常的数字。断言检查的是代数不变量，仍不能判断位置映射对视频是否有用，真实效果还要在固定权重的 DiT 上测。 请注意：示例里的 13×30×45 是格点形状。不能仅凭像素分辨率就猜它。真实模型要逐层根据 VAE 的时空压缩、首帧特殊处理、patch size 和 padding 算出 (T,H,W)；错一个维度，位置表和 token 顺序都会错位。 05. 工业级实现对照 以 2026-09 核对的 CogVideo SAT 源码 Rotary3DPositionEmbeddingMixin 为例，类中 dim_t = hidden_size_head // 4、dim_h = dim_w = hidden_size_head // 8 * 3，每轴各自建立频率与格点，再广播、拼接、取 sin/cos。它在注意力函数中只旋转视频 query/key，先切出 text_length 个文本 token 原样放回；可选 rot_v 才旋转 value，默认并非必要。rope_T/H/W 从预计算的三维表里截取当前视频尺寸，排平为 (T*H*W,D)。这与最小代码的数组顺序一致，也便于逐轴检查实际相位。 一个容易踩的坑是：这个 SAT 类的构造函数虽然接收 height_interpolation、width_interpolation 和 time_interpolation 参数，在本次读取的 main 文件中，角度表仍直接由 torch.arange 计算，这些参数没有参与坐标缩放。因此不能因为配置里写了 height_interpolation=1.875，就声称该 SAT 旋转实现真的进行了高度插值。文章讨论的“逐轴插值”是可选的设计实验，不是这一段 CogVideo SAT 代码的实际行为。作者原论文采用扩表取前段，恰好与这里的原坐标逻辑相符。 Diffusers 的 get_3d_rotary_pos_embed 也有时间 1/4、空间各 3/8 的分配，但它提供 grid_type="linspace" 和 "slice" 两种格点构造。linspace 可把某个目标空间范围映射到当前高宽格点；slice 在最大表上用原索引再取当前尺寸。二者的差别首先是坐标选取，不是旋转公式变化。该函数在本次核对中返回 (cos,sin)，供注意力处理器作用在 query/key 上。具体使用哪个 grid_type、裁剪坐标是什么，应沿管线调用栈核对，不能只看函数定义就给整个模型下结论。 写工业调用栈时建议按四个具体问题逐层追：输入是像素帧还是 VAE latent；patchify 后一条样本的 (T,H,W) 分别是多少；三轴坐标从零开始、按最大表切片还是重新映射到训练范围；cos/sin 最后究竟旋转的是视频 query/key、文本 token 是否另有位置处理。单独找到 get_3d_rotary_pos_embed 的名字，只能证明功能存在，不能证明这一模型版本已经启用它。CogVideoX 系列可能根据配置选择旋转位置或学习位置，不同 checkpoint 更不能一概而论。 维度排列也值得实际打印。本文的 rotate 假设相邻两维组成 (实部, 虚部)，因此角表用 repeat_interleave(2) 形式复写到两维。若上游先把所有实部放前半、虚部放后半，旋转公式仍可成立，但 rotate_half 与 cos/sin 的广播顺序必须一起换。仅把缓存从另一仓库复制过来，形状同为 [N,D] 也可能在不报错的情况下给完全不同的旋转平面。检查一个位置 1、一个二维平面 (1,0) 的旋转结果应近似 (cos 1,sin 1)，比只核对 shape 更可靠。 生产实现还会处理变长视频、文本与视频拼接、packing、缓存、设备和精度。最小脚本只面向一个规则视频网格；真实 batch 如果把不同视频压成一条序列，必须保存每段的起点与独立 (T,H,W)，否则第二个视频的时间会错误地延续第一个视频的末帧。cos/sin 表可以复用，但格点顺序和真实 token 排列必须逐项对齐。很多“换分辨率立即崩”其实是这种形状或坐标实现错误，而非外推理论失败。 06. 代价与边界 纯 3D RoPE 本身不引入新的可训练位置参数，换形状时也可以按轴计算新表；具体模型仍可能另外叠加可学习的位置向量。这个性质让多分辨率训练与不同长度输入更自然，但没有把注意力从二次复杂度变成线性。高、宽各两倍时，视频 token 四倍、稠密注意力项十六倍：示例从 13×30×45=17,550 个 token 增至 13×60×90=70,200 个 token，若天真地为单个头物化一张 FP16 分数矩阵，理论容量会从约 0.574 GiB 涨到 9.179 GiB，尚未计梯度、softmax 和其他头。FlashAttention 不必把整张矩阵常驻显存，所以这不是实际峰值预测；它仍准确反映稠密注意力的二次计算量。把空间 RoPE 频率调得再聪明也不会消除这笔账，仍需 VAE 压缩、稀疏或局部注意力、分块计算等手段。 外推的另一重代价是位置区分与训练分布之间的冲突。直接延长表，邻近 token 的相位关系不变，却把远端坐标送到没见过的位置；全部插值，把远端压回旧尺度，却弱化每一步的相位变化。轴越短、分到的旋转平面越少，可用的频率跨度通常越窄。CogVideoX 的时间 16 维并不意味着“时间必然先坏”或“空间必然先坏”：结果还取决于训练长度、VAE 时间压缩、运动跨度与内容建模。 分辨率外推尤其要保住物理尺度解释。同一人物在原图占 10 个 latent 格点，放大画幅后占 20 个，坐标间隔与模型学过的视觉尺度同时变了；只插值位置并不能把纹理、构图与物体大小的分布一起搬过去。长视频类似：更多帧要求模型记住人物、场景和动作因果，位置相位不是记忆网络。若输出失败，先分辨是表索引越界、token 排列错位、局部重复、全局模糊还是跨帧身份漂移，再决定该调位置、数据、注意力或去噪。 相同 token 数也不意味着相同空间外推。 假设训练格点为 (T,H,W)=(8,32,32)：目标 (8,64,32) 与 (8,32,64) 都把视频 token 翻倍，但前者只扩高轴，后者只扩宽轴。在 3D RoPE 中，它们分别改变高段或宽段的最远相位；在展平 1D RoPE 中，改宽还会重新标定所有行间与帧间的序列距离。实验可以保持 token 数、权重、prompt 和种子一致，只互换扩张轴，再分别检查竖向重复、横向重复与主体比例。若两个结果不同，先核对训练画幅分布和 VAE/patchify，才讨论是否来自轴频率。 排查失败时可以先做一个不动权重的三步实验。第一步固定 VAE latent 与 patch token，只换位置表构造，在同一坐标对上比较 q·k；若结果突然变化，检查表索引、轴顺序或维度配对。第二步固定画幅和视频内容，单独改变时间格点长度，看主体漂移是否只在跨越训练时长后出现。第三步固定时间格点，单独变高、宽并分别观察全局构图和局部纹理；若高宽都扩大导致局部重复，不能直接断言是 rope_theta 错，还要看是否出现 token 数暴涨引发的注意力截断、分块边界或 VAE tiling 伪影。把这些变量一次全改掉，会失去定位能力。 若比较位置插值与直接外推，评分也至少拆成两类：近邻细节，例如边缘、纹理与细小物体是否保持；远距结构，例如人物是否只有一个、左右场景是否连贯、同一主体跨帧是否一致。CogVideoX 附录里“一张模糊的大画面”与“多个清楚的小画面”的差异正说明：一个单一的美学分数可能掩盖局部与全局朝相反方向变化。画质指标之外还要记录 token 长度、时空网格与推理显存，否则方法之间可能使用了不同计算预算。 这套分轴设计也不表达“这是第 1 帧”这类绝对锚点——它主要编码相对位移。首帧图像条件、相机轨迹和剪辑时间戳需要通过额外条件或掩码表达。多图参考、回放状态与不同视频段拼接时，不同段若都从 (0,0,0) 重启，还需要明确段 ID 或相机等身份信息，免得同坐标误当同内容。 更远一步，镜头绕场景移动后再次拍到同一处墙面，这处墙面的世界位置没变，屏幕上的 (h,w) 却可能不同；反过来，相同 (h,w) 也可能指向另一块墙面。视频 3D RoPE 的“3D”是时间加二维屏幕格点，不是场景的三维坐标。面向相机可控的视频 DiT，SCoPE 将相机视线作为另一种位置线索加入注意力，并保留原有 RoPE；这是研究者针对屏幕坐标局限提出的扩展，不是普通 3D RoPE 自动拥有的几何能力。 07. 经典论文脉络 RoFormer、Position Interpolation 与 YaRN 的旋转及语言长上下文扩展已在先修文章讲过；对视频的问题是哪一轴、哪一频段需要缩放，语言模型结果不能直接当画质证据。 CogVideoX 将三维相对位置用于视频 DiT，并用多分辨率 Frame Pack 与渐进训练支持不同尺寸；论文附录展示的初始生成状态对比揭示了“局部清晰”与“全局成形”的矛盾。它是该模型上的定性证据，不是所有视频 RoPE 的通用胜负结论。 Diffusers 的 CogVideoX 实现 提供可核对的三轴维度分配、格点生成和注意力接口；论文、SAT 与 Diffusers 版本未必走同一坐标路径。 SCoPE 在视频 DiT 中保留时空格点 RoPE，另加相机视线线索以改善相机轨迹下的场景一致性；它讨论的是屏幕位置之外的几何信息，并非改写基础 RoPE 的旋转恒等式。 08. 常见误解 误解一：3D RoPE 就是把三个坐标加起来再做 1D RoPE。 若先算 $t+h+w$，(1,0,0) 与 (0,1,0) 就变成同一个位置值，时间变化和上下位移无法区分。正确做法是三段特征独立旋转，最后只在注意力内积里汇总贡献。 误解二：相对位置恒等式成立就证明任意分辨率都能生成。 恒等式对没见过的坐标同样成立；模型从未学过如何在新相位、更多 token 和新物体尺度上使用它。数学可计算性只保证不会因公式本身报错。 误解三：把所有坐标除以两倍没有副作用。 位置插值让新端点回到旧范围，却也让邻近格点的相位差减半，可能伤害局部细节。若目标只换高度，就不要无故压缩时间和宽度。 误解四：调大 rope_theta 等于 Position Interpolation。 更改基数主要影响较低频平面，最高频不动；位置除以缩放因子让所有频率一起慢下来。两者可能组合，作用机制不同。 误解五：CogVideoX 配置里有 height_interpolation 参数，就一定在使用插值。 本次读到的 SAT Rotary3DPositionEmbeddingMixin 接受该参数，却没有用它来计算 grid_h。要按真实执行路径确认；不同实现不能只凭参数名互相代替。 误解六：空间轴拿了 75% 维度，所以换分辨率导致的所有错误都来自空间 RoPE。 维度比例只是一种频率预算，无法排除注意力规模、训练分辨率、VAE 压缩和内容尺度的影响。看输出症状、做逐项消融，才有因果结论。 09. 动手验证 先运行 python3 rope_3d_minimal.py，看 vertical_neighbor_* 四行。把 lower=(0,1,0) 改成 (1,0,0)，即从“下一行”改为“下一帧”：3D 位移由 (0,1,0) 改成 (1,0,0)，1D 位移则由 W 改成 H×W。再单独改网格宽度，检查三轴版的下一帧相位是否仍只由时间坐标决定。 调用 rope_3d_cache(2,6,4,64,scales=(1,2,1)) 与 scales=(2,1,1)，分别只压缩高、时间坐标；检查同一高坐标的高段相位和同一时间坐标的时间段相位。最后把网格 2×3×4 改成 2×6×4：缓存行数应从 24 变 48，列数仍为 64，只有高坐标的取值范围增加。 如果要评估生成模型的真实外推质量，不能只跑这段 NumPy 脚本。固定模型权重、prompt、种子与去噪设置，比较原训练画幅、直接扩表、逐轴插值和不同频段缩放；同时记录局部纹理、重复图案、全局构图、跨帧主体一致性与实际显存/延迟。只有这样才能分清“位置编码问题”与“高分辨率本来就需要新训练”的差异。 10. 延伸阅读 继续向下可以看视频 DiT 的稀疏注意力、3D VAE/patchify 和长视频状态记忆：位置只负责“在哪里”，这几项决定“看哪些 token、用什么尺度和怎样记住内容”。 本篇用到的具体外部依据以 CogVideoX 论文、CogVideo SAT 源码 与 Diffusers 位置编码源码 为准；它们分别支持论文中的训练取舍与两种工程实现，文中的 NumPy 输出仅证明所写公式和示例代码一致。 附录：完整代码 09 节用到的脚本全文如下（rope_3d_minimal.py、make_figures.py）。复制到本地存成同名文件，按各脚本开头的依赖说明准备环境后即可运行。 rope_3d_minimal.py """NumPy-only 3D RoPE and per-axis position extrapolation demo.""" import numpy as np def axis_dims(head_dim): """CogVideoX's time/height/width split for a head divisible by sixteen (all three axis sizes must be even).""" assert head_dim > 0 and head_dim % 16 == 0 dims = (head_dim // 4, head_dim * 3 // 8, head_dim * 3 // 8) assert all(d % 2 == 0 for d in dims) return dims def axis_angles(positions, dim, base=10000.0, scale=1.0): """Angles [positions, dim/2]; scale>1 compresses this axis only.""" assert dim % 2 == 0 and scale > 0 inv_freq = base ** (-np.arange(0, dim, 2, dtype=np.float64) / dim) return np.asarray(positions, dtype=np.float64)[:, None] * inv_freq[None, :] / scale def rope_3d_cache(frames, height, width, head_dim=64, scales=(1, 1, 1)): """Return interleaved cos/sin, one row per (t,h,w) token.""" assert frames > 0 and height > 0 and width > 0 assert len(scales) == 3 and all(scale > 0 for scale in scales) dt, dh, dw = axis_dims(head_dim) at = axis_angles(np.arange(frames), dt, scale=scales[0])[:, None, None, :] ah = axis_angles(np.arange(height), dh, scale=scales[1])[None, :, None, :] aw = axis_angles(np.arange(width), dw, scale=scales[2])[None, None, :, :] shape = (frames, height, width) angles = np.concatenate( [np.broadcast_to(a, shape + (d // 2,)) for a, d in zip((at, ah, aw), (dt, dh, dw))], axis=-1, ).reshape(frames * height * width, head_dim // 2) return np.repeat(np.cos(angles), 2, axis=-1), np.repeat(np.sin(angles), 2, axis=-1) def rotate(x, cos, sin): """Rotate adjacent dimension pairs; x and cache both [tokens, head_dim].""" paired = x.reshape(x.shape[0], -1, 2) quarter_turn = np.stack((-paired[..., 1], paired[..., 0]), axis=-1).reshape(x.shape) return x * cos + quarter_turn * sin def score_for_positions(q, k, pos_q, pos_k, scales=(1, 1, 1)): """A single dot product without building an attention matrix.""" dims = axis_dims(q.size) out_q, out_k = [], [] start = 0 for axis, dim in enumerate(dims): aq = axis_angles([pos_q[axis]], dim, scale=scales[axis]) ak = axis_angles([pos_k[axis]], dim, scale=scales[axis]) cq, sq = np.repeat(np.cos(aq), 2, axis=-1), np.repeat(np.sin(aq), 2, axis=-1) ck, sk = np.repeat(np.cos(ak), 2, axis=-1), np.repeat(np.sin(ak), 2, axis=-1) out_q.append(rotate(q[start:start + dim][None], cq, sq)[0]) out_k.append(rotate(k[start:start + dim][None], ck, sk)[0]) start += dim return np.dot(np.concatenate(out_q), np.concatenate(out_k)) def flatten_index(pos, height, width): """The row-major 1D index of a video token at (time, height, width).""" t, h, w = pos return t * height * width + h * width + w def score_for_flattened_positions(q, k, pos_q, pos_k, height, width): """Baseline: one 1D RoPE angle table for all head channels.""" mq = flatten_index(pos_q, height, width) mk = flatten_index(pos_k, height, width) aq = axis_angles([mq], q.size) ak = axis_angles([mk], k.size) cq, sq = np.repeat(np.cos(aq), 2, axis=-1), np.repeat(np.sin(aq), 2, axis=-1) ck, sk = np.repeat(np.cos(ak), 2, axis=-1), np.repeat(np.sin(ak), 2, axis=-1) return float(np.dot(rotate(q[None], cq, sq)[0], rotate(k[None], ck, sk)[0])) def naive_score_matrix_gib(token_count, bytes_per_element=2): """GiB for one materialized dense attention-score matrix of one head.""" return token_count * token_count * bytes_per_element / 2**30 if __name__ == "__main__": np.set_printoptions(precision=6, suppress=True) dims = axis_dims(64) cos, sin = rope_3d_cache(2, 3, 4, 64) rng = np.random.default_rng(7) tokens = rng.normal(size=cos.shape) rotated = rotate(tokens, cos, sin) norm_error = np.max(np.abs(np.linalg.norm(tokens, axis=1) - np.linalg.norm(rotated, axis=1))) print("axis_dims=", dims) print("cache_shape=", cos.shape, "input_shape=", tokens.shape) print("rotated_shape=", rotated.shape) print("norm_error=", f"{norm_error:.2e}") q, k = rng.normal(size=64), rng.normal(size=64) a = score_for_positions(q, k, (0, 1, 2), (1, 2, 3)) b = score_for_positions(q, k, (5, 6, 7), (6, 7, 8)) print("same_offset_scores=", f"{a:.6f}", f"{b:.6f}", "difference=", f"{abs(a-b):.2e}") upper, lower = (0, 0, 0), (0, 1, 0) flat4 = flatten_index(lower, 2, 4) - flatten_index(upper, 2, 4) flat8 = flatten_index(lower, 2, 8) - flatten_index(upper, 2, 8) one_d4 = score_for_flattened_positions(q, k, upper, lower, 2, 4) one_d8 = score_for_flattened_positions(q, k, upper, lower, 2, 8) three_d4 = score_for_positions(q, k, upper, lower) three_d8 = score_for_positions(q, k, upper, lower) three_d_offset = tuple(b - a for a, b in zip(upper, lower)) print("vertical_neighbor_1d_offsets=", flat4, flat8) print("vertical_neighbor_1d_scores=", f"{one_d4:.6f}", f"{one_d8:.6f}") print("vertical_neighbor_3d_offsets=", three_d_offset, three_d_offset) print("vertical_neighbor_3d_scores=", f"{three_d4:.6f}", f"{three_d8:.6f}") # These are algebraic invariants, not video-quality tests. Keeping them as # assertions turns the example into a small regression test for pairing, # broadcasting and relative-position behavior. assert norm_error < 1e-12 assert abs(a - b) < 1e-12 assert abs(three_d4 - three_d8) < 1e-12 assert not np.isclose(one_d4, one_d8) raw = axis_angles([15], dims[1])[0] pi = axis_angles([15], dims[1], scale=2)[0] print("height_phase_at_15_first_last=", np.round(raw[[0, -1]], 6)) print("height_PI_x2_first_last=", np.round(pi[[0, -1]], 6)) print("time_phase_at_15_first_last=", np.round(axis_angles([15], dims[0])[0][[0, -1]], 6)) small_tokens = 13 * 30 * 45 large_tokens = 13 * 60 * 90 print("token_count_13x30x45=", small_tokens) print("token_count_13x60x90=", large_tokens) print("dense_attention_ratio=", (large_tokens / small_tokens) ** 2) print( "naive_score_matrix_fp16_gib_per_head=", f"{naive_score_matrix_gib(small_tokens):.3f}", f"{naive_score_matrix_gib(large_tokens):.3f}", ) make_figures.py """Generate 3D RoPE figures. Requires NumPy, Matplotlib and Pillow.""" from pathlib import Path import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt import numpy as np from PIL import Image, ImageDraw, ImageFont from rope_3d_minimal import axis_angles, axis_dims OUT = Path(__file__).resolve().parents[1] / "figures" OUT.mkdir(exist_ok=True) def _chinese_font(size): for path in ( "/System/Library/Fonts/Hiragino Sans GB.ttc", "/System/Library/Fonts/PingFang.ttc", "/usr/share/fonts/truetype/noto/NotoSansCJK-Regular.ttc", ): try: return ImageFont.truetype(path, size) except OSError: pass raise RuntimeError("请安装中文字体后再生成 1D/3D 对照图") def one_d_vs_three_d(): """Show the same vertical neighbor under two video widths.""" image = Image.new("RGB", (1800, 950), "#f7f9fe") draw = ImageDraw.Draw(image) ink, muted = "#1e293b", "#475569" draw.text((80, 50), "同一对上下相邻的 token，换宽度会怎样？", font=_chinese_font(68), fill=ink) draw.text((82, 150), "A=(0,0,0) → B=(0,1,0)；只改变视频格点宽度 W", font=_chinese_font(42), fill=muted) cards = ( (70, 260, 865, 790, "#e8efff", "#1d4ed8", "1D RoPE：一个序列距离", "W=4 → Δm=4", "W=8 → Δm=8", "上下相邻仍是 1 格，但旋转相位变了"), (935, 260, 1730, 790, "#e3f7f2", "#047857", "3D RoPE：三个轴分别计数", "W=4 → (0,1,0)", "W=8 → (0,1,0)", "高轴始终转 1 格，时间和宽轴不动"), ) for x0, y0, x1, y1, bg, accent, title, first, second, note in cards: draw.rounded_rectangle((x0, y0, x1, y1), radius=38, fill=bg) draw.rectangle((x0 + 38, y0 + 48, x0 + 50, y0 + 130), fill=accent) draw.text((x0 + 77, y0 + 51), title, font=_chinese_font(51), fill=ink) draw.text((x0 + 75, y0 + 210), first, font=_chinese_font(68), fill=accent) draw.text((x0 + 75, y0 + 322), second, font=_chinese_font(68), fill=accent) draw.text((x0 + 75, y0 + 435), note, font=_chinese_font(35), fill=muted) draw.text((80, 844), "旋转和相对位置性质相同；变化的是坐标系与每轴的特征维度。", font=_chinese_font(42), fill=ink) image.save(OUT / "one_d_vs_three_d.png", optimize=True) def axis_budget(): dims = axis_dims(64) fig, ax = plt.subplots(figsize=(9, 2.8)) colors = ["#1d4ed8", "#0d9488", "#ea580c"] start = 0 for name, dim, color in zip(("Time 25%", "Height 37.5%", "Width 37.5%"), dims, colors): ax.barh([0], [dim], left=start, color=color, height=0.55) ax.text(start + dim / 2, 0, f"{name}\n{dim} dims", ha="center", va="center", color="white", weight="bold") start += dim ax.set_xlim(0, 64) ax.set_ylim(-0.55, 0.55) ax.set_xlabel("Head channels (64 total)") ax.set_yticks([]) ax.set_title("CogVideoX 3D RoPE: independent rotary planes per axis") fig.tight_layout() fig.savefig(OUT / "axis_budget.png", dpi=180) plt.close(fig) def phase_tradeoff(): dim_h = axis_dims(64)[1] pos = np.arange(32) original = axis_angles(pos, dim_h) compressed = axis_angles(pos, dim_h, scale=2) fig, axes = plt.subplots(1, 2, figsize=(9, 3.5), sharey=True) for ax, values, title in zip(axes, (original, compressed), ("Raw coordinates", "Height PI: positions / 2")): ax.plot(pos, np.cos(values[:, 0]), color="#ea580c", linewidth=2, label="highest frequency") ax.plot(pos, np.cos(values[:, -1]), color="#1d4ed8", linewidth=2, label="lowest frequency") ax.set_title(title) ax.set_xlabel("Height-token index") ax.grid(alpha=0.18) axes[0].set_ylabel("cos(angle)") axes[1].legend(frameon=False, loc="lower right") fig.suptitle("Interpolation halves phase change on this axis at every frequency") fig.tight_layout() fig.savefig(OUT / "phase_tradeoff.png", dpi=180) plt.close(fig) if __name__ == "__main__": one_d_vs_three_d() axis_budget() phase_tradeoff() for name in ("one_d_vs_three_d.png", "axis_budget.png", "phase_tradeoff.png"): print(OUT / name) 更多 AIGC 论文解读，关注微信公众号「人工智能炼丹君」 每日更新 · 论文精选 · 深度解读 · 技术脉络 微信搜索 人工智能炼丹君 或扫描下方二维码关注

AIGC 基本功|从 DPO 到 GRPO：去掉价值网络-GRPO 从 DPO 到 GRPO：去掉价值网络 所属方向：对齐与强化学习 ｜ 难度：进阶 ｜ 前置知识：PPO 关键词：DPO、GRPO、偏好优化、组内归一化、免价值网络、参考模型 01. 为什么需要它 把一个 7B 语言模型接进 PPO-RLHF，训练栈很快会变成四个模型：负责生成的 actor、预测每个 token 未来回报的 critic、限制策略漂移的 reference model，以及把整条回答变成分数的 reward model。真正更新参数的是 actor 和 critic，后两者虽然通常冻结，前向和显存仍要付钱。 critic 是最难解释的一笔成本。语言模型往往只在回答结束时拿到一个总分，但 value head 必须给回答中每个 token 估计未来回报。它不但要和 actor 同量级地跑前向、反向，还要从稀疏的终局信号中学习密集价值。DeepSeekMath 指出，LLM 的 value model 通常与 policy model 规模相当，这正是 GRPO 要删除的部件。 以全参数 AdamW 的粗略账本为例：7B 参数的 BF16 权重约 14 GB，BF16 梯度约 14 GB，FP32 主权重和两份动量约 84 GB，一个可训练模型合计约 112 GB 逻辑状态，实际会被 ZeRO/FSDP 分片。再训练一套同规模 critic，等于再背一次参数、梯度、优化器状态和激活。删掉 critic 的收益很大，但不能机械地说“所有训练都便宜十倍”：GRPO 每个 prompt 要在线生成一组回答，DeepSeekMath 的实验组大小是 64，采样和奖励打分可能成为新的主成本。 DPO 与 GRPO 给了两条不同的减法路线。DPO 把带 KL 约束的 RLHF 目标改写成偏好二分类，训练时不再在线采样，也不需要显式奖励模型；GRPO 保留在线探索和奖励信号，用同一道题的组内相对分数替代 critic。二者名字相邻，解决的却不是同一个问题。 图 1：DPO 省掉显式奖励模型和在线 RL 循环，GRPO 省掉 critic。Policy 与 reference 是否共享底座、reward 是规则还是模型，会继续改变真实成本。 02. 最小可用理解 先记住五句话： DPO 吃离线偏好对：同一 prompt 下给一个 chosen 和一个 rejected，直接提高 chosen 相对 reference 的概率优势。 GRPO 吃在线样本组：同一 prompt 采样 $G$ 个回答，打分后用组均值当 baseline，不训练 value model。 GRPO 沿用 PPO 的概率比与 clip：它只换了优势估计，没有丢掉“旧策略采样、新策略更新”这套框架。 reference 不是 critic：reference 负责限制长期漂移，critic 负责估计当前状态的未来回报；删掉 critic 后 reference 通常还在。 组内相对分数只回答谁更好：如果一组答案同分，GRPO 在这道题上没有方向；如果奖励尺度跨题不一致，组内归一化反而能消掉尺度差。 因此，从 DPO 到 GRPO 不是把一个损失函数逐项变形为另一个，而是从“离线比较学习”切换到“在线生成、在线评价”。选择算法时先问数据来自哪里，再问显存够不够。 03. 数学推导 3.1 PPO-RLHF 的起点 对 prompt $x$ 和回答 $y$，标准 KL 正则化目标可以写成： $$\max_{\pi}\ \mathbb{E}_{y\sim\pi(\cdot\mid x)}[r(x,y)]-\beta D_{KL}[\pi(\cdot\mid x)\|\pi_{ref}(\cdot\mid x)]$$ $\pi$ 是待训练策略，$\pi_{ref}$ 是固定参考策略，$r(x,y)$ 是回答奖励，$\beta$ 控制策略离参考模型多远。PPO 用采样轨迹、critic 优势和裁剪目标近似优化它。这里至少要维护 actor、critic 和 reference；若奖励由神经网络给出，还要运行 reward model。 3.2 DPO 如何把奖励塞回策略 上面的目标对每个 $x$ 都有闭式最优策略： $$\pi^*(y\mid x)=\frac{1}{Z(x)}\pi_{ref}(y\mid x)\exp\left(\frac{r(x,y)}{\beta}\right)$$ $Z(x)$ 是归一化常数。把式子移项，可把奖励写成策略与参考策略的对数概率比： $$r(x,y)=\beta\log\frac{\pi^*(y\mid x)}{\pi_{ref}(y\mid x)}+\beta\log Z(x)$$ 偏好数据给出同一 prompt 下的胜者 $y_w$ 与败者 $y_l$。用 Bradley-Terry 模型表示“胜者更好”的概率： $$P(y_w\succ y_l\mid x)=\sigma(r(x,y_w)-r(x,y_l))$$ 把上一式代入后，同一道题共享的 $\beta\log Z(x)$ 自动相消。用可训练策略 $\pi_\theta$ 代替未知最优策略，得到 DPO 损失： $$\mathcal{L}_{DPO}=-\mathbb{E}\log\sigma\left(\beta\left[\log\frac{\pi_\theta(y_w\mid x)}{\pi_{ref}(y_w\mid x)}-\log\frac{\pi_\theta(y_l\mid x)}{\pi_{ref}(y_l\mid x)}\right]\right)$$ 括号里有两层差。第一层是 chosen 与 rejected 的差；第二层是当前 policy 相对 reference 的差。只让 policy 提高 chosen 概率还不够，它提高的幅度必须超过 reference 原本已有的偏好。 DPO 的省钱来自训练协议：偏好对提前收集好，训练循环只做似然计算和二分类，不需要当前策略在线生成回答，也不需要在循环里调用显式 reward model。但代价也在这里：数据分布被固定。策略更新后会产生哪些新错误，旧偏好集未必覆盖。 3.3 PPO 的 critic 到底提供了什么 PPO 的优势通常来自 critic： $$A_t=Q(s_t,a_t)-V_\phi(s_t)$$ $V_\phi(s_t)$ 预测从当前 token 前缀继续生成的平均回报。减去它不会改变策略梯度的期望，却能降低方差。问题是，回答只在末尾拿到 $r(x,y)$，critic 却要在每个位置给出估计；它要学习的目标本身就很噪。 如果暂时不追求每个 token 的状态价值，而是对同一道题一次采 $G$ 个完整回答，就能用样本组构造另一种 baseline。 3.4 GRPO 的组内相对优势 从旧策略 $\pi_{old}$ 对同一 prompt $x$ 采样回答 $\{y_1,\ldots,y_G\}$，奖励为 $\{r_1,\ldots,r_G\}$。Outcome supervision 下，GRPO 定义： $$\hat A_i=\frac{r_i-\mathrm{mean}(r_1,\ldots,r_G)}{\mathrm{std}(r_1,\ldots,r_G)+\epsilon}$$ $\epsilon$ 防止标准差为零时除零。均值把“这道题本来容易还是难”扣掉，标准差把不同题目的奖励尺度拉到接近一致。对一条回答里的每个有效 token，原始 GRPO 都广播同一个 $\hat A_i$：得分高于组均值的回答整体增概率，低于均值的整体降概率。 图 2：两道题的原始 reward 尺度相差很大，组内标准化后各自在零均值参照系里比较。被保留的是同题回答间的相对好坏。 组均值包含样本自身，并不满足“baseline 与本次动作独立”的条件。对固定 prompt 的独立同分布样本，若只减组均值而不除标准差，有 $E[(r_i-\bar r)\nabla\log\pi(y_i)]=(1-1/G)\nabla J$；leave-one-out 均值可去掉这项缩放。再除随机组标准差会改变样本及 prompt 权重，因此完整 GRPO 不能宣称是原始期望奖励梯度的无偏估计。 这不是免费午餐。$G$ 太小时，组均值和标准差很不稳定；奖励全相等时，优势全部为零；同一回答的所有 token 共用一个结果优势时，算法知道“整条回答好”，却不知道哪一步推理真正立功。 3.5 GRPO 仍然是 PPO 家族 对回答 $y_i$ 的第 $t$ 个 token，定义新旧策略概率比： $$\rho_{i,t}(\theta)=\frac{\pi_\theta(y_{i,t}\mid x,y_{i,<t})}{\pi_{old}(y_{i,t}\mid x,y_{i,<t})}$$ GRPO 沿用 PPO-Clip 的保守替代目标，并直接在 loss 上加 reference KL： $$\mathcal{J}_{GRPO}=\mathbb{E}\left[\frac{1}{G}\sum_i\frac{1}{|y_i|}\sum_t\left(\min\left(\rho_{i,t}\hat A_i,\mathrm{clip}(\rho_{i,t},1-\varepsilon,1+\varepsilon)\hat A_i\right)-\beta D_{KL}(\pi_\theta\|\pi_{ref})\right)\right]$$ $\pi_{old}$ 是本轮 rollout 的采样快照，用于纠正同一批数据被重复更新后的分布偏移；$\pi_{ref}$ 是行为锚点，用于限制训练全过程的长期漂移。两者在训练刚开始可能数值相同，但职责不同，不能混用。 DeepSeekMath 使用的单样本 KL 正值估计器是： $$D_{KL}=\exp(\Delta)-\Delta-1,\quad \Delta=\log\pi_{ref}-\log\pi_\theta$$ 因为对任意实数 $\Delta$ 都有 $e^\Delta\geq1+\Delta$，这个量非负。在实现里它可以直接由采到 token 的 policy/reference logprob 计算，不必枚举整个词表。它的期望等于 $D_{KL}(\pi_\theta\|\pi_{ref})$ 的前提是动作从当前 $\pi_\theta$ 采样；实际 rollout 来自 $\pi_{old}$，更新后直接平均得到的是旧采样分布下的代理量，需要相应重要性校正才有同一无偏解释。 3.6 DPO 与 GRPO 的统一观察 DPO 和 GRPO 都在做“相对比较”，但参照物不同： 维度 DPO GRPO 数据 固定的 chosen/rejected 对 当前策略在线生成的一组回答 相对量 policy 相对 reference 的偏好间隔 回答奖励相对同组均值 删除的部件 显式 reward model 与在线 RL 循环 critic/value model 仍需保留 policy、reference、偏好数据 policy、old policy logprob、reference、奖励函数 主要风险 离线数据覆盖不足 采样昂贵、组内同分、奖励投机 这张表也解释了为什么“DPO 之后再做 GRPO”可以成立：先用便宜的离线偏好对把策略推到合理区域，再用在线 GRPO 探索当前策略真正会生成的回答。但这是一种训练安排，不是数学上必须的前后继关系。 3.7 手算一组 DPO 和 GRPO 先看 DPO。设第一对样本中，当前 policy 对 chosen 和 rejected 的整句 logprob 分别为 -2.2 和 -2.8，因此 policy 的偏好间隔是 0.6；reference 对二者的 logprob 分别为 -2.0 和 -2.4，偏好间隔是 0.4。当前策略相对 reference 多出来的间隔只有 0.2。当 $\beta=0.1$ 时，送进 sigmoid 的 logit 是 0.02，胜者概率只略高于 0.5，loss 接近随机猜测的 $\log 2$。 这解释了一个常见现象：chosen 的 logprob 即使很低，也不代表 DPO 一定给它很大梯度。DPO 看的是四个 logprob 组成的相对差。若 reference 已经强烈偏爱 chosen，而 policy 只复制了这种偏好，DPO 不会把它误判成新的进步；若 policy 相比 reference 反而缩小了 chosen 的优势，logit 会变负，损失上升。 再看 GRPO。第一组奖励是 [0.2, 0.8, 1.4, 0.6]，均值为 0.75，标准差约为 0.433。标准化后得到 [-1.2699, 0.1154, 1.5008, -0.3463]。第二个回答虽然原始分数 0.8 看上去不高，但它略高于同题平均，因此得到很小的正优势；第三个回答比平均高很多，得到最大的正优势。 假设这四条回答更新后的 token ratio 依次是 [1.2840, 1.1052, 0.7788, 0.9048]，裁剪区间为 [0.8,1.2]。第一条优势为负，而 ratio 大于 1，说明坏回答变得更可能，目标必须继续惩罚，不能裁掉。第二条优势为正且 ratio 在区间内，正常提供正向梯度。第三条优势为正但 ratio 低于 0.8，说明好回答被错误降权，目标同样继续惩罚。第四条优势为负且 ratio 在区间内，也正常提供降低概率的梯度。 真正被截住的是“已经朝正确方向走过头”的两种情况：正优势回答的 ratio 超过 1.2，或负优势回答的 ratio 低于 0.8。PPO/GRPO 的 clip 从来不是把所有 ratio 强制塞进区间，而是停止奖励过度乐观的改进。只看 np.clip(ratio) 而不把优势符号放进来，几乎一定会误读目标。 04. 代码实现 code/dpo_grpo_minimal.py 只依赖 NumPy，包含三部分：DPO 二分类损失、GRPO 的组内优势与裁剪损失，以及一个四动作 bandit 的在线训练。实际运行输出如下： 两张解释图由 code/make_figures.py 从文中的固定数值直接生成，图和公式使用同一组奖励，避免手工绘图与代码结果不一致。 === DPO === pair_shape=(2,) preference_logits=[0.02 0.03] loss=0.6807 === GRPO group baseline === rewards.shape=(2, 4) advantages= [[-1.2699 0.1154 1.5008 -0.3463] [-1.2648 -0.6324 0.6324 1.2648]] row_mean=[-0. 0.] === GRPO clipped objective === token_logps.shape=(4, 3) ratio_first_token=[1.284 1.1052 0.7788 0.9048] mean_kl=0.018550, loss=0.1626 奖励张量形状是 [prompt, generation] = [2, 4]。每行优势的均值都为零，说明每道题自己形成 baseline。裁剪例子里第一条优势为负而 ratio 为 1.284，第三条优势为正而 ratio 为 0.7788：二者都朝错误方向越界，因此 min 保留未裁剪项以继续惩罚。另两条 ratio 在区间内。这四条回答都没有进入裁剪收益的平台分支，与 3.7 节一致。 bandit 有四个回答动作，真实奖励分别是 [-0.2, 0.3, 1.0, 0.0]。脚本不训练 critic，每轮只采 32 个动作、组内标准化奖励，再用两轮 PPO 式更新复用这组样本： step P(a0) P(a1) P(a2=best) P(a3) 1 0.2288 0.2494 0.2883 0.2335 20 0.0210 0.0346 0.9179 0.0265 40 0.0048 0.0064 0.9833 0.0055 80 0.0026 0.0035 0.9906 0.0033 第 80 步时，最优动作概率从 0.25 升到 0.9906。这个实验展示的是组 baseline 能提供正确梯度，不代表大模型会如此平滑：真实回答空间巨大，奖励噪声、长度与采样温度都会改变结果。 05. 工业级实现对照 本文人工核对了 Hugging Face TRL 的 GRPOTrainer，以 2026-09-13 的主分支提交 cd2c528 为准。生产实现比论文公式多出几层防护： 奖励函数全返回空值的回答会标成 NaN，从组均值中排除，随后把优势置零，避免“不可评分”被误当成零奖励。 scale_rewards 支持按 group、batch 或不缩放；多目标奖励还可以“先加权求和再归一化”或“各目标先归一化再求和”。两者表达的偏好并不一样。 importance sampling 可以按 token 或 sequence 计算。长回答里 token 级 ratio 连乘会很不稳定，序列级方案则改变了权重粒度。 loss 已不止原始 grpo，还包括 bnpo、dr_grpo、dapo 等不同归一化方式。它们主要在回答长度和 batch/token 归一化上修补偏差。 使用 vLLM 生成时，推理引擎与训练模型的 logprob 可能有数值差异，框架提供额外 importance-sampling correction。 PEFT 场景可在同一底座上挂训练 adapter 与 reference adapter，减少复制完整模型的显存。 最值得监控的日志不是单独一个 reward，而是：组内 reward 标准差为零的比例、clip ratio、KL、回答长度、每种 reward 的均值与方差。如果 reward 上升同时长度暴涨、KL 激增或同分组变多，训练可能在钻评分规则的空子。 工程上还要区分三种 mask：prompt token 不应进入 completion loss，padding token 不应参与平均，工具调用或环境交互产生的非策略 token 也不该伪装成策略动作。mask 错一个维度，代码仍可能正常运行，但 loss 的分母已经变了。 5.1 一次工业 GRPO update 的数据流 第一步是挑 prompt。随机抽题看似公平，但大量过易题会产生“全对组”，大量过难题会产生“全错组”，两者都没有组内方差。生产系统通常会维护题目通过率、奖励方差和最近采样时间，让训练 batch 更多覆盖当前策略刚好有分歧的区域。 第二步是冻结采样快照。系统保存 $\pi_{old}$ 的版本标识，并为每个 prompt 生成 $G$ 条 completion。若用独立 vLLM 服务生成，必须记录推理引擎返回的 token id、mask 和 logprob；只保存文本再重新 tokenize，特殊 token、空格或聊天模板差异都可能让动作序列错位。 第三步是评分。规则验证器适合数学答案、代码测试和工具任务；神经 reward model 适合帮助性、风格和主观质量。多个奖励合并时要先决定是“先加权再组内归一化”，还是“每种奖励各自归一化再加权”。前者保留权重指定的绝对尺度，后者让每种奖励先取得相近话语权。 第四步计算 group statistics。分布式训练里，同一 prompt 的 $G$ 条回答可能散在不同 GPU 上，均值和标准差必须在完整组上聚合。若每张卡只看本地子组，baseline 会随数据切分改变；同一实验换一个 GPU 数就可能得到不同梯度。 第五步计算当前 policy 与 reference 的 token logprob。old_logprob 必须保持 rollout 时的值，不能在每个优化 epoch 重算；reference_logprob 可以预计算并缓存，也可以通过冻结模型或禁用 LoRA adapter 得到。三种 logprob 看着形状相同，来源却完全不同。 第六步形成 per-token loss。Outcome reward 把回答级优势广播到有效 completion token，再乘 ratio、做 clip、加 KL。随后按 mask 归一化。先对每条序列取平均再对 batch 平均，会让长短回答权重相近；直接对全 batch token 求平均，则长回答贡献更多。两种都能运行，但优化目标不同。 第七步更新并监控。一次 rollout 可做有限个优化 epoch，随后必须重新采样。应同时记录 reward、组内标准差、zero-std 比例、KL、clip fraction、回答长度和吞吐量。只有 reward 上升而这些量失控时，不能把结果叫作训练成功。 5.2 六个不会报错却会训歪的实现细节 分组边界错位。 假设 batch 排列本应是“题目 A 的八条回答、题目 B 的八条回答”，数据加载器却在聚合前打乱了 completion，view(-1, G) 仍然合法，均值却混合了不同题目。最可靠的做法是保留 prompt id，计算优势前断言每个连续组的 id 完全相同，分布式 gather 后再断言每题正好有 $G$ 条记录。 把样本标准差和总体标准差混着用。 NumPy 的 std 默认分母是 $G$，PyTorch 某些接口曾默认用 $G-1$。组很大时差别不明显，$G=2$ 或 4 时会明显改变优势幅度。公式、最小实现和训练框架必须约定同一个定义，日志里也应记录原始 reward，而不是只保留标准化后的值。 先逐卡归一化再全局聚合。 数据并行下每张卡只拿到同一题的一部分回答，如果先在本地求均值，结果取决于回答被分到哪张卡。正确顺序是先收集完整组的 reward，再统一计算统计量，最后切回各进程需要的那一段优势。 生成文本与训练 token 不一致。 推理服务可能自动补聊天模板、合并空白、提前截断，训练端重新编码字符串后得到另一串 token。此时旧策略 logprob 与当前策略 logprob 对应的动作不同，ratio 已经失去意义。rollout 应把 token ids 视作主数据，文本只用于评分和审计；恢复训练时也要保留 tokenizer 与 chat template 版本。 长度归一化悄悄改变偏好。 若每条回答先按自身长度求 token loss 均值，短回答和长回答各占一个样本权重；若整个 batch 按 token 总数求均值，长回答天然贡献更多。再叠加“回答越长越可能碰到一次正确步骤”的 reward，模型可能学会延长推理。应同时画 reward 对长度的散点图，并用长度分桶比较通过率。 奖励器和策略一起漂移。 在线迭代训练 reward model 时，今天的 0.8 与上周的 0.8 可能不是同一尺度。组内归一化能消掉部分仿射漂移，却消不掉排序规则改变。需要保留固定校准集和冻结的旧奖励器，周期性比较新旧排序一致率；否则策略看似持续提高，其实只是追着变化的裁判跑。 最后还要保存可恢复状态：policy、optimizer、scheduler、reference 版本、rollout 生成参数、reward 版本和数据游标缺一不可。只保存 actor 权重可以用于推理，却无法保证训练恢复后仍在优化同一个目标。GRPO 少了 critic checkpoint，但并没有把实验追踪简化成一个模型文件。 06. 代价与边界 删掉 critic，省下的是真实成本。 全参数训练时少一套同规模可训练状态和激活；LoRA 或 ZeRO 下绝对节省会变化，但工程链路仍少了 value target、value loss、value clipping 与 critic checkpoint。 组采样把成本转移到生成端。 GRPO 必须对同一 prompt 采多个回答。推理较长、组大小较大或 reward model 很重时，rollout 可能比 critic 更贵。适合有便宜可验证奖励的数学、代码和工具任务；对主观写作、开放式对话，可靠打分本身就是难题。 组内归一化丢掉跨题绝对难度。 一道题的最高分只有 0.2，另一道题四个答案都超过 0.9，标准化后仍可能产生相似幅度的优势。前者也许是奖励器不可靠或题目不可解，却会得到同等更新权重。 同分会让梯度消失。 二值正确性奖励在策略已经很强或很弱时尤其明显：一组全对或全错，标准差接近零。这时增加组大小、提高采样温度、设计更细奖励或主动挑选处在学习边界的 prompt，比盲目调学习率更有效。 Outcome reward 的信用分配很粗。 一条 2000-token 推理只拿一个终局优势，每个 token 被同方向推动。过程奖励可以细化信号，但会引入步骤标注、奖励模型偏差与“奖励器教模型写格式”的新问题。 reference KL 不能修复坏奖励。 KL 只限制走多远，不判断方向对不对。奖励函数偏爱冗长、套话或特定格式时，较小 KL 只是让投机变慢。 6.1 什么时候选 DPO，什么时候选 GRPO 如果手里已经有高质量 chosen/rejected 数据，生成和人工标注又很贵，DPO 通常是更稳的第一步。它像普通监督训练一样容易批处理，失败时也容易回查是哪一对偏好造成梯度。它尤其适合风格、安全边界、回答格式这类可以提前整理成对的数据。 如果任务能自动验答案，而且策略必须靠探索发现新解法，GRPO 更合适。数学最终答案、代码单元测试、工具调用结果和可执行环境都能提供相对便宜的在线奖励。此时 critic 很难从稀疏终局分数学习每个 token 的价值，组 baseline 的简化更有吸引力。 如果奖励高度主观、每条回答又很长，两者都不会自动解决问题。DPO 受限于离线偏好覆盖，GRPO 受限于奖励模型可靠性和 rollout 成本。常见组合是 SFT 打基础、DPO 吸收稳定的人类偏好、GRPO 在可验证子任务上继续在线探索，最后再用独立评测检查能力回退。 做预算时至少拆成四项：训练模型状态、冻结模型前向、在线生成 token 数、奖励评估次数。只比较“有几个模型”会漏掉 GRPO 的采样账；只比较“每步生成多少 token”又会漏掉 PPO critic 的反向和优化器状态。把四项分别量出来，才能判断删掉 critic 是否真的让当前系统便宜了一个量级。 还有一笔常被忽略的是通信。PPO 的 critic 若与 actor 分开部署，需要在 rollout、价值推理和更新之间搬运 token、mask、value 与 advantage；若与 actor 共卡，又会争夺显存和计算流。GRPO 把 value 请求换成组统计，统计量只是一组标量，通信链更短。可是一旦奖励模型远程部署，$G$ 条回答仍要跨服务传输并等待评分。优化时应测端到端每步墙钟时间，把生成、评分、训练和等待分别打点；单看 GPU 峰值显存，很容易把“省显存”误写成“省总成本”。 评估也要和训练奖励解耦。至少准备一套不参与更新的题目、一个不同实现的验证器，以及人工抽查样本。训练奖励与独立评测同时上升，才说明策略能力可能真的提高；只有训练奖励上涨，最多能证明优化器成功找到了当前评分函数偏爱的输出。对推理模型还应分别统计最终答案、推理有效性、格式合规和长度，避免一个聚合分数掩盖能力回退。 这些记录也是定位回归、比较算法和复现实验的最低证据。 07. 经典论文脉络 Proximal Policy Optimization Algorithms（arXiv:1707.06347，2017）提出可多轮复用 rollout 的裁剪替代目标，是 GRPO 概率比与 clip 的直接来源。 Direct Preference Optimization（arXiv:2305.18290，2023）利用 KL 正则化 RLHF 的闭式最优策略，把偏好学习变成二分类，训练阶段绕开在线 RL。 DeepSeekMath（arXiv:2402.03300，2024）正式提出 GRPO，用同题多回答的组分数替代 value model；7B 实验中 GSM8K 从 82.9% 提升到 88.2%，MATH 从 46.8% 提升到 51.7%。 DeepSeek-R1（arXiv:2501.12948，2025）把大规模 GRPO 推到长链推理与可验证奖励场景，也让组内相对优化成为推理模型训练的常用基线。 这些数字来自各论文自己的设置，不能跨模型直接比较。DPO 的优势在训练简单，GRPO 的优势在能跟随当前策略在线探索；算法名本身不决定最终能力。 08. 常见误解 “GRPO 不需要奖励模型。” 错。GRPO 不需要 critic，但必须有奖励信号。奖励可以是神经 reward model、规则验证器、单元测试或多个评分器的组合。 “DPO 和 GRPO 都去掉 value，所以是一回事。” DPO 通过重参数化绕开整套在线 RL；GRPO 仍是 policy-gradient 方法，只把优势 baseline 从 learned value 换成 group statistics。 “组均值就是无偏的真实价值函数。” 组均值只是当前策略、当前 prompt、当前采样温度下的 Monte Carlo baseline。它会随组大小和采样分布变化，不是对所有状态都可查询的 $V(s)$。 “归一化后奖励尺度不重要。” 组内仿射变换会被抵消，但排序错误、同分、异常值和多目标权重仍会改变梯度。reward design 依然是核心工作。 “去掉 critic 就一定便宜十倍。” 不成立。删掉一个同规模可训练模型可能显著降低显存和通信，但 GRPO 增加了组采样。总成本取决于序列长度、组大小、reward 推理、并行方式和是否使用 LoRA。 “clip 能保证 KL 很小。” clip 只截断采样动作上的替代收益；共享参数更新仍可能让未采到 token 的分布大幅变化，所以工程实现仍单独监控或惩罚 KL。 09. 动手验证 在仓库目录运行： python3 outputs/fundamentals_files/dpo_grpo/code/dpo_grpo_minimal.py 可以做四个改动观察边界： 把 group_size=32 改成 2，多换几个 seed。预期收敛抖动明显增大，有时一组采不到最优动作。 把 reward_table 改成四个相同值。预期标准差为零，优势接近零，策略不再学习。 把最好动作奖励从 1.0 改成 0.31，使它只比第二名高 0.01。预期需要更多样本才能稳定区分。 把 beta 从 0.02 提到 0.5。预期策略更贴近均匀 reference，最优动作概率上升更慢、上限更低。 再做一个 DPO 对照：交换 policy_chosen 与 policy_rejected，观察 preference logit 变负、loss 变大。它说明 DPO 优化的是相对 reference 的 chosen/rejected 间隔，而不是 chosen 的绝对 logprob。 10. 延伸阅读 读这篇之前建议先看： 策略梯度与 PPO 基础：优势估计、概率比和 clip 的完整推导。 读完这篇可以继续看： 把 RL 用到扩散模型上（DiffusionRL）：把多步去噪改写成 MDP 轨迹。 视频生成中的强化学习与奖励模型：GRPO 进入视觉生成后，采样成本、时序奖励与 reward hacking 如何变化。 附录：完整代码 09 节用到的脚本全文如下（dpo_grpo_minimal.py、make_figures.py）。复制到本地存成同名文件，按各脚本开头的依赖说明准备环境后即可运行。 dpo_grpo_minimal.py """DPO 与 GRPO 的最小可运行实现。 运行：python3 dpo_grpo_minimal.py 只依赖 NumPy，不需要模型权重或外部数据。 """ import numpy as np def log_sigmoid(x): """数值稳定的 log(sigmoid(x))。""" return -np.logaddexp(0.0, -x) def dpo_loss(policy_chosen, policy_rejected, ref_chosen, ref_rejected, beta=0.1): """DPO 二分类损失；四个输入都是整条回答的对数概率。""" policy_gap = policy_chosen - policy_rejected reference_gap = ref_chosen - ref_rejected logits = beta * (policy_gap - reference_gap) return -log_sigmoid(logits).mean(), logits def group_advantages(rewards, eps=1e-4): """对同一 prompt 的 G 个奖励做组内中心化和标准化。""" mean = rewards.mean(axis=1, keepdims=True) std = rewards.std(axis=1, keepdims=True) return (rewards - mean) / (std + eps) def grpo_loss(new_logps, old_logps, ref_logps, advantages, mask, clip_eps=0.2, beta=0.04): """Outcome-supervision GRPO：一条回答的优势广播给它的全部 token。""" log_ratio = new_logps - old_logps ratio = np.exp(log_ratio) clipped_ratio = np.clip(ratio, 1 - clip_eps, 1 + clip_eps) advantage_per_token = advantages[:, None] surrogate = np.minimum( ratio * advantage_per_token, clipped_ratio * advantage_per_token, ) # DeepSeekMath 使用的正值 KL 估计器：exp(x) - x - 1，其中 x=log pi_ref-log pi。 ref_gap = ref_logps - new_logps per_token_kl = np.exp(ref_gap) - ref_gap - 1 per_token_loss = -surrogate + beta * per_token_kl sequence_loss = (per_token_loss * mask).sum(axis=-1) / mask.sum(axis=-1) return sequence_loss.mean(), ratio, per_token_kl def softmax(logits): shifted = logits - logits.max() exp_logits = np.exp(shifted) return exp_logits / exp_logits.sum() def train_group_bandit(steps=80, group_size=32, seed=7): """用 GRPO 训练四选一 bandit，展示无 critic 的在线更新。""" rng = np.random.default_rng(seed) logits = np.zeros(4, dtype=np.float64) reference_logits = np.zeros(4, dtype=np.float64) reward_table = np.array([-0.2, 0.3, 1.0, 0.0]) lr, clip_eps, beta = 0.18, 0.2, 0.02 snapshots = [] for step in range(steps): old_logits = logits.copy() old_probs = softmax(old_logits) actions = rng.choice(4, size=group_size, p=old_probs) rewards = reward_table[actions] advantages = (rewards - rewards.mean()) / (rewards.std() + 1e-4) old_selected = np.log(old_probs[actions]) ref_probs = softmax(reference_logits) ref_selected = np.log(ref_probs[actions]) # 同一组 rollout 做两轮 PPO 式复用；old_logprob 始终来自采样快照。 for _ in range(2): probs = softmax(logits) new_selected = np.log(probs[actions]) ratio = np.exp(new_selected - old_selected) active = ~(((advantages >= 0) & (ratio > 1 + clip_eps)) | ((advantages < 0) & (ratio < 1 - clip_eps))) # d[-min(rA, clip(r)A)]/d log pi，以及采样 KL 项的导数。 coefficient = -active.astype(np.float64) * ratio * advantages ref_gap = ref_selected - new_selected coefficient += beta * (1 - np.exp(ref_gap)) one_hot = np.eye(4)[actions] grad_logp = one_hot - probs[None, :] grad_logits = (coefficient[:, None] * grad_logp).mean(axis=0) logits -= lr * grad_logits if step in {0, 19, 39, 79}: snapshots.append((step + 1, softmax(logits).copy())) return snapshots if __name__ == "__main__": np.set_printoptions(precision=4, suppress=True) policy_chosen = np.array([-2.2, -1.7]) policy_rejected = np.array([-2.8, -2.1]) ref_chosen = np.array([-2.0, -1.9]) ref_rejected = np.array([-2.4, -2.0]) loss_dpo, logits_dpo = dpo_loss( policy_chosen, policy_rejected, ref_chosen, ref_rejected ) print("=== DPO ===") print(f"pair_shape={policy_chosen.shape}") print(f"preference_logits={logits_dpo}") print(f"loss={loss_dpo:.4f}") rewards = np.array([[0.2, 0.8, 1.4, 0.6], [8.0, 9.0, 11.0, 12.0]]) advantages = group_advantages(rewards) print("\n=== GRPO group baseline ===") print(f"rewards.shape={rewards.shape}") print(f"advantages=\n{advantages}") print(f"row_mean={advantages.mean(axis=1)}") old_logps = np.log(np.array([ [0.30, 0.25, 0.20], [0.18, 0.22, 0.20], [0.12, 0.16, 0.20], [0.25, 0.20, 0.18], ])) new_logps = old_logps + np.array([ [0.25, 0.25, 0.25], [0.10, 0.10, 0.10], [-0.25, -0.25, -0.25], [-0.10, -0.10, -0.10], ]) ref_logps = old_logps - 0.05 mask = np.ones_like(old_logps) one_group_adv = group_advantages(np.array([[0.2, 0.8, 1.4, 0.6]]))[0] loss_grpo, ratio, kl = grpo_loss( new_logps, old_logps, ref_logps, one_group_adv, mask ) print("\n=== GRPO clipped objective ===") print(f"token_logps.shape={new_logps.shape}") print(f"ratio_first_token={ratio[:, 0]}") print(f"mean_kl={kl.mean():.6f}, loss={loss_grpo:.4f}") print("\n=== Online group bandit ===") print("step P(a0) P(a1) P(a2=best) P(a3)") for step, probs in train_group_bandit(): print(f"{step:>4} " + " ".join(f"{p:.4f}" for p in probs)) make_figures.py """生成 DPO/GRPO 教程的两张解释图。""" from pathlib import Path import matplotlib.pyplot as plt import numpy as np from matplotlib.patches import FancyBboxPatch ROOT = Path(__file__).resolve().parents[1] OUT = ROOT / "figures" OUT.mkdir(exist_ok=True) plt.rcParams.update({ "font.sans-serif": ["PingFang SC", "Arial Unicode MS", "DejaVu Sans"], "axes.unicode_minus": False, "figure.dpi": 160, }) def box(ax, x, y, w, h, text, color): patch = FancyBboxPatch( (x, y), w, h, boxstyle="round,pad=0.02,rounding_size=0.03", facecolor=color, edgecolor="white", linewidth=1.5, ) ax.add_patch(patch) ax.text(x + w / 2, y + h / 2, text, ha="center", va="center", fontsize=12, weight="bold", color="#182238") def method_map(): fig, axes = plt.subplots(1, 3, figsize=(14, 4.6), facecolor="#f7f9fc") methods = [ ("PPO-RLHF", [("偏好对", "#dbeafe"), ("奖励模型", "#fde68a"), ("在线采样", "#ddd6fe"), ("Actor + Critic", "#fecaca")], "奖励模型 + critic 都要维护"), ("DPO", [("离线偏好对", "#dbeafe"), ("二分类损失", "#bbf7d0"), ("Policy + Reference", "#ddd6fe")], "省掉显式奖励模型和在线 RL"), ("GRPO", [("同题多回答", "#dbeafe"), ("规则/奖励打分", "#fde68a"), ("组内相对优势", "#bbf7d0"), ("Policy + Reference", "#ddd6fe")], "保留在线采样，省掉 critic"), ] for ax, (title, blocks, note) in zip(axes, methods): ax.set_xlim(0, 1) ax.set_ylim(0, 1) ax.axis("off") ax.set_title(title, fontsize=18, weight="bold", color="#172554", pad=14) gap = 0.08 h = (0.72 - gap * (len(blocks) - 1)) / len(blocks) y = 0.85 - h for text, color in blocks: box(ax, 0.12, y, 0.76, h, text, color) y -= h + gap ax.text(0.5, 0.03, note, ha="center", va="bottom", fontsize=11, color="#475569") fig.suptitle("三种对齐路线省掉的部件不同", fontsize=22, weight="bold", color="#0f172a", y=1.02) fig.tight_layout() fig.savefig(OUT / "alignment_method_map.png", bbox_inches="tight", facecolor=fig.get_facecolor()) plt.close(fig) def group_relative_plot(): rewards = np.array([[0.2, 0.8, 1.4, 0.6], [8.0, 9.0, 11.0, 12.0]]) advantages = (rewards - rewards.mean(1, keepdims=True)) / rewards.std(1, keepdims=True) fig, axes = plt.subplots(1, 2, figsize=(12, 4.8), facecolor="#f7f9fc") colors = ["#60a5fa", "#a78bfa", "#34d399", "#fb7185"] x = np.arange(4) for row, name in enumerate(["题目 A", "题目 B"]): axes[0].bar(x + (row - 0.5) * 0.18, rewards[row], width=0.18, label=name, color=colors[row]) axes[0].set_title("原始奖励不可跨题直接比较", fontsize=16, weight="bold") axes[0].set_xticks(x, [f"回答 {i+1}" for i in x]) axes[0].set_ylabel("reward") axes[0].legend(frameon=False) width = 0.34 axes[1].bar(x - width / 2, advantages[0], width=width, label="题目 A", color="#60a5fa") axes[1].bar(x + width / 2, advantages[1], width=width, label="题目 B", color="#a78bfa") axes[1].axhline(0, color="#334155", linewidth=1) axes[1].set_title("组内标准化只保留相对名次", fontsize=16, weight="bold") axes[1].set_xticks(x, [f"回答 {i+1}" for i in x]) axes[1].set_ylabel("advantage") axes[1].legend(frameon=False) for ax in axes: ax.spines[["top", "right"]].set_visible(False) ax.grid(axis="y", alpha=0.18) fig.suptitle("GRPO：每道题自己组成一个参照系", fontsize=21, weight="bold", color="#0f172a", y=1.02) fig.tight_layout() fig.savefig(OUT / "group_relative_advantage.png", bbox_inches="tight", facecolor=fig.get_facecolor()) plt.close(fig) if __name__ == "__main__": method_map() group_relative_plot() print(OUT / "alignment_method_map.png") print(OUT / "group_relative_advantage.png") 更多 AIGC 论文解读，关注微信公众号「人工智能炼丹君」 每日更新 · 论文精选 · 深度解读 · 技术脉络 微信搜索 人工智能炼丹君 或扫描下方二维码关注

AIGC 基本功|策略梯度与 PPO 基础-PPO 策略梯度与 PPO 基础 所属方向：对齐与强化学习 ｜ 难度：核心必修 ｜ 前置知识：无 关键词：策略梯度、REINFORCE、优势函数、GAE、PPO、裁剪、KL 约束 01. 为什么需要它 先看一个很容易把模型训坏的场景。 你有一个已经做过监督微调的语言模型，给同一个问题采样了若干回答，奖励模型认为其中一条更好。最直接的想法是：让模型提高这条回答的概率。于是你把它的对数概率乘上奖励，连续更新十轮。训练日志里的 reward 不断上升，重新采样时却发现模型开始反复输出奖励模型偏爱的句式，内容变长，语言能力下降，最后连原来会答的问题也答不好。 问题不只在奖励。至少有三笔账同时失控了： 信用分配：一个回答只得到一个总分，到底哪些 token 值得鼓励？ 估计噪声：这条回答得 8 分，是动作真的好，还是题目本来就容易？ 更新幅度：数据由旧策略采出，新策略反复使用同一批数据后已经变了，为什么还能继续把旧结论当真？ 策略梯度解决第一笔账：把“提高期望奖励”变成对可采样策略的梯度。价值函数和优势估计解决第二笔账：把状态本身的难易扣掉，只保留“这个动作比预期好多少”。PPO 的裁剪目标处理第三笔账：允许一批昂贵样本做多轮更新，同时阻止单个样本把新旧概率比推得过远。 这三层关系决定了后面的算法谱系。DPO 改写优化问题，GRPO 换掉价值网络，RLOO 换基线，很多 RLHF 系统再加入参考模型 KL、奖励归一化和长度处理。名字不断变，问题仍可追溯到三个量：优势怎么估、概率比怎么算、策略走多远。 02. 最小可用理解 先记住四句话： 策略梯度会提高“优势为正”的动作概率，降低“优势为负”的动作概率。 优势 $A(s,a)$ 是采取动作 $a$ 后的回报，相对状态 $s$ 下平均预期的超额部分。 GAE 用参数 $\lambda$ 在低方差但依赖 critic 的一步 TD，与高方差但依赖较少的完整回报之间插值。 PPO-Clip 比较新旧策略的动作概率比 $r_t$；一旦更新已经朝有利方向超过 $1\pm\epsilon$，该样本不再继续提供奖励。 把一次语言模型生成看成一条轨迹也很直观：状态是 prompt 加已经生成的 token，动作是下一个 token，策略是词表上的概率分布，终局奖励来自奖励模型。critic 预测“从当前前缀继续生成，平均还能拿多少分”，优势则问“刚才选的这个 token，让结局比平均预期好还是差”。 PPO 的 clip 不是把梯度值裁小，也不是保证 KL 一定不超过某个阈值。它裁的是替代目标里的概率比收益。这个区别是理解所有变体的入口。 03. 数学推导 3.1 从期望回报到策略梯度 考虑一个马尔可夫决策过程。$s_t$ 是时刻 $t$ 的状态，$a_t$ 是动作，$r_t$ 是环境在这一步给出的奖励，$\gamma\in[0,1]$ 是折扣因子。策略 $\pi_\theta(a_t\mid s_t)$ 由参数 $\theta$ 控制。一条长度为 $T$ 的轨迹记作 $\tau=(s_0,a_0,r_0,\ldots,s_T)$，折扣回报为： $$R(\tau)=\sum_{t=0}^{T-1}\gamma^t r_t$$ 训练目标是让轨迹的期望回报最大： $$J(\theta)=\mathbb{E}_{\tau\sim\pi_\theta}[R(\tau)]$$ 轨迹概率由初始状态分布 $p(s_0)$、策略和环境转移 $p(s_{t+1}\mid s_t,a_t)$ 连乘得到： $$p_\theta(\tau)=p(s_0)\prod_{t=0}^{T-1}\pi_\theta(a_t\mid s_t)p(s_{t+1}\mid s_t,a_t)$$ 对 $J$ 求梯度，并使用恒等式 $\nabla p=p\nabla\log p$： $$\nabla_\theta J(\theta)=\mathbb{E}_{\tau\sim\pi_\theta}[R(\tau)\nabla_\theta\log p_\theta(\tau)]$$ 环境转移不依赖 $\theta$，因此它在对数梯度中消失，只剩策略项： $$\nabla_\theta J(\theta)=\mathbb{E}\left[R(\tau)\sum_{t=0}^{T-1}\nabla_\theta\log\pi_\theta(a_t\mid s_t)\right]$$ 这就是 REINFORCE 的骨架。还可以利用“未来动作不能影响过去奖励”，把整条轨迹回报换成从 $t$ 开始的 reward-to-go： $$G_t=\sum_{l=0}^{T-t-1}\gamma^l r_{t+l}$$ 于是每个动作只为它之后的结果负责： $$\nabla_\theta J(\theta)=\mathbb{E}\left[\sum_{t=0}^{T-1}\gamma^t G_t\nabla_\theta\log\pi_\theta(a_t\mid s_t)\right]$$ 这里的 $\gamma^t$ 不能漏：本文从固定初始状态、折扣总回报 $J$ 出发，$G_t$ 的折扣却从当前步重新计起。只有 $\gamma=1$，或把该权重吸收到折扣状态访问分布时，才能省略显式 $\gamma^t$。实际 PPO 常把 rollout 的时间步均匀平均，这是常用替代目标，不应与上面的精确有限时域梯度混同。 如果 $G_t$ 总为正，采到的每个动作都会被提高概率，区别只是力度大小。这种估计虽然无偏，方差却很大。我们需要一个不改变期望梯度的参照物。 3.2 基线为什么可以减 从回报减去任何只依赖状态、不依赖本次动作的基线 $b(s_t)$，期望梯度不变。因为对固定状态 $s$： $$\mathbb{E}_{a\sim\pi_\theta}[b(s)\nabla_\theta\log\pi_\theta(a\mid s)]=b(s)\nabla_\theta\sum_a\pi_\theta(a\mid s)=0$$ 最常见的基线是状态价值函数： $$V^\pi(s_t)=\mathbb{E}_\pi[G_t\mid s_t]$$ 动作价值函数把本次动作也作为条件： $$Q^\pi(s_t,a_t)=\mathbb{E}_\pi[G_t\mid s_t,a_t]$$ 两者之差就是优势函数： $$A^\pi(s_t,a_t)=Q^\pi(s_t,a_t)-V^\pi(s_t)$$ $A>0$ 表示这个动作比该状态下的平均动作好，$A<0$ 表示更差。Actor 通过优势更新策略，critic 通过回归回报学习 $V$，所以这类方法叫 actor-critic。 3.3 从 TD 残差到 GAE 真实 $Q$ 和 $V$ 都不知道，只能估计。最短视的估计是一步 TD 残差： $$\delta_t=r_t+\gamma(1-d_t)V_\phi(s_{t+1})-V_\phi(s_t)$$ $V_\phi$ 是参数为 $\phi$ 的 critic，$d_t\in\{0,1\}$ 表示这一步后轨迹是否真正终止。若 $d_t=1$，就不能再 bootstrap 到下一个状态。critic 准确时，$\delta_t$ 是优势的低方差估计；critic 有系统误差时，它也会把误差直接传给 actor。 把未来多个 TD 残差加进来，会得到不同步数的优势估计。GAE 用指数权重把它们合在一起： $$\hat A_t^{\mathrm{GAE}(\gamma,\lambda)}=\sum_{l=0}^{T-t-1}(\gamma\lambda)^l\delta_{t+l}$$ 实际代码不需要为每个 $t$ 再套一层求和。由上式直接得到从后向前的递推： $$\hat A_t=\delta_t+\gamma\lambda(1-d_t)\hat A_{t+1}$$ 两个极端很关键。$\lambda=0$ 时，$\hat A_t=\delta_t$，只看一步，方差小但强依赖 critic。$\lambda=1$ 且轨迹正确终止时，TD 项会望远镜式相消，得到 $G_t-V(s_t)$，对未来 critic 误差不再敏感，但把后续所有随机奖励都带了进来。常见的 $\gamma=0.99,\lambda=0.95$ 是经验起点，不是定律。 训练 critic 时通常使用 value target： $$\hat V_t^{\mathrm{target}}=\hat A_t+V_\phi(s_t)$$ 实现时要区分 terminated 和因时间上限触发的 truncated。真正终止的状态没有未来价值；时间截断通常仍应 bootstrap。把两者都当 done=1，会在每段 rollout 尾部制造系统性低估。 3.4 为什么需要新旧策略概率比 采样成本高，所以 PPO 会冻结采样策略 $\pi_{\theta_{\mathrm{old}}}$，用同一批轨迹对当前策略 $\pi_\theta$ 做多轮小批量更新。数据来自旧策略，目标却要描述新策略，动作重要性比在固定旧策略状态上校正动作分布；它不校正整条轨迹的状态访问分布，所以以下是局部替代目标： $$r_t(\theta)=\frac{\pi_\theta(a_t\mid s_t)}{\pi_{\theta_{\mathrm{old}}}(a_t\mid s_t)}=\exp(\log\pi_\theta-\log\pi_{\theta_{\mathrm{old}}})$$ 未裁剪的替代目标是： $$L^{\mathrm{PG}}(\theta)=\mathbb{E}_t[r_t(\theta)\hat A_t]$$ 刚开始新旧策略相同，$r_t=1$。如果 $\hat A_t>0$，增大该动作概率会使目标变大；如果 $\hat A_t<0$，减小概率会使目标变大。但在同一批数据上更新多轮后，某些比率可能冲到 1.8、0.2，旧数据已不能可靠描述新策略附近的行为。 3.5 PPO-Clip 的 min 到底裁了哪一边 PPO-Clip 的核心目标只有一行： $$L^{\mathrm{CLIP}}(\theta)=\mathbb{E}_t\left[\min\left(r_t\hat A_t,\mathrm{clip}(r_t,1-\epsilon,1+\epsilon)\hat A_t\right)\right]$$ 这里的 min 取较悲观的收益。它必须与优势的符号一起看： 当 $\hat A_t>0$，把好动作的概率比推到 $1+\epsilon$ 以上不再得分；但若概率比跌到 $1-\epsilon$ 以下，损失仍继续变差，算法不会保护错误方向。 当 $\hat A_t<0$，把坏动作的概率比压到 $1-\epsilon$ 以下不再得分；但若坏动作反而变得更可能，惩罚仍继续加重。 所以 clip 是单侧的乐观收益上限。它不会把所有比率强行夹回区间，也不会保证一次优化后每个样本都满足区间约束。一个 batch 内不同样本共享参数，更新某个样本可能把另一个样本推得更远。 PPO 通常最小化总损失，因此策略项前面带负号，再加 critic 的均方误差和熵奖励： $$\mathcal{L}_{\mathrm{total}}=-L^{\mathrm{CLIP}}+c_v\mathbb{E}_t[(V_\phi(s_t)-\hat V_t^{\mathrm{target}})^2]-c_e\mathbb{E}_t[\mathcal{H}(\pi_\theta(\cdot\mid s_t))]$$ $c_v$ 控制价值损失权重，$c_e$ 控制探索强度，$\mathcal H$ 是策略熵。LLM 对齐还常加入相对参考模型 $\pi_{\mathrm{ref}}$ 的 KL 惩罚： $$r_t^{\mathrm{RLHF}}=r_t^{\mathrm{task}}-\beta\left(\log\pi_\theta(a_t\mid s_t)-\log\pi_{\mathrm{ref}}(a_t\mid s_t)\right)$$ 注意这里有两个不同的“旧模型”。$\pi_{\theta_{\mathrm{old}}}$ 是本轮采样快照，用来计算 PPO ratio；$\pi_{\mathrm{ref}}$ 通常是固定的 SFT 参考模型，用来限制长期漂移。把两者混成一个模型，会把更新稳定性和行为保持两个问题混在一起。 3.6 把公式还原成一轮训练 一轮 PPO 可以按下面的顺序执行。顺序很重要，因为每个张量对应的策略版本不同。 第一步，冻结当前 actor 为 $\pi_{\theta_{\mathrm{old}}}$，用它与环境交互，保存每一步的状态、动作、奖励、old_logprob、critic value、终止标记。语言模型场景还要保存 completion mask，明确哪些位置真的是策略采出的动作。 第二步，在 rollout 末端决定是否 bootstrap。若轨迹真正终止，next_value=0；若只是采样窗口用完，则由 critic 估计最后状态价值。然后从后向前计算 $\delta_t$ 和 GAE。此时优势应视作固定训练标签，actor 更新不能反向穿过 GAE 进入旧 value。 第三步，用 $\hat V_t^{\mathrm{target}}=\hat A_t+V_{\mathrm{old}}(s_t)$ 构造 critic 目标。这里加的是采样时保存的旧 value。如果一边更新 critic 一边用新 value 重算 target，目标本身会追着模型移动。 第四步，打乱 rollout，切成 minibatch。当前 actor 对已经采过的动作重新计算 new_logprob，然后用对数概率之差求 ratio。直接先算概率再做除法容易在大词表和低概率动作上出现下溢。 第五步，分别计算未裁剪和裁剪后的 policy surrogate，逐样本取 min，再按有效动作 mask 求平均。同时计算 value loss、entropy bonus，以及可选的 reference KL。 第六步，反向传播并做梯度范数裁剪。一个 rollout 通常会被重复使用若干 epoch；每轮都用同一份 old_logprob，但 new_logprob 随参数更新而变化。若 approximate KL 超阈值，就提前结束剩余 epoch。 第七步，丢弃这批 on-policy 数据，用更新后的 actor 重新采样。此时新策略成为下一轮的 old policy。PPO 能复用的是“一轮之内的有限次数”，并没有把数据永久变成离线训练集。 用张量形状检查这套流程也很有效。经典控制任务常见 reward/value/advantage/logprob 都是 [T, N]，其中 $T$ 是 rollout 长度，$N$ 是并行环境数；语言模型常见 [B, L]，其中 $B$ 是回答数，$L$ 是序列长度。只要其中一个量偷偷变成 [B] 并发生广播，训练可能不会报错，却会让整条回答的每个 token 共享一个本不该共享的系数。 3.7 手算一条两步轨迹 用一个最小数字例子把 GAE 和 clip 接起来。设两步奖励为 $[0,1]$，critic 预测为 $[0.4,0.7]$，第二步后真正终止，$\gamma=0.9$。最后一步的 TD 残差是 $\delta_1=1-0.7=0.3$；第一步是 $\delta_0=0+0.9\times0.7-0.4=0.23$。当 $\lambda=0.95$ 时： $$\hat A_0=0.23+0.9\times0.95\times0.3=0.4865$$ 对应的 value target 为 $0.4865+0.4=0.8865$，接近完整 Monte Carlo 回报 $0+0.9\times1=0.9$。若 $\lambda=0$，target 只有 $0.23+0.4=0.63$，它完全相信 critic 对下一状态的 0.7 估计。这几个数把“依赖 critic”和“纳入远期真实奖励”的差别直接展开了。 再设旧策略给第一步动作的概率是 0.25，新策略更新后变为 0.30，则 ratio 为 $0.30/0.25=1.2$。若该动作优势为 0.4865，正好到 $\epsilon=0.2$ 的上边界；继续把概率推到 0.35 时，未裁剪收益会按 ratio=1.4 计算，裁剪收益仍只按 1.2 计算，min 选择后者。若优势改为负数，1.4 一侧不会被保护，因为提高坏动作概率应该继续受到惩罚。 这个例子也说明两个超参数并不独立。$\lambda$ 改变优势的大小和噪声，进而改变样本多久撞上 clip；$\epsilon$ 决定同一批优势可以推动概率多远。只调其中一个而不观察 ratio 分布，往往解释不了训练曲线。 04. 代码实现 4.1 逐步算 GAE code/gae_minimal.py 只依赖 NumPy。它先对一条长度为 4 的固定轨迹从后向前递推，再模拟 20000 条随机轨迹，比较不同 $\lambda$ 下 value target 对真实 $V(s_0)$ 的偏差和方差。固定轨迹的真实输出是： rewards.shape=(4,), values.shape=(4,) lambda=0.00 delta=[0.1445 0.143 0.597 0.2 ] A=[0.1445 0.143 0.597 0.2 ] lambda=0.50 delta=[0.1445 0.143 0.597 0.2 ] A=[0.3858 0.4875 0.696 0.2 ] lambda=0.95 delta=[0.1445 0.143 0.597 0.2 ] A=[0.9734 0.8814 0.7851 0.2 ] lambda=1.00 delta=[0.1445 0.143 0.597 0.2 ] A=[1.0652 0.93 0.795 0.2 ] 同一行的 delta 不随 $\lambda$ 变化，因为它只由一步奖励和 critic 决定；$\lambda$ 改变的是未来残差往前传播的强度。最后一步已经终止，所以四组优势都等于 0.2。 随机实验故意给 critic 加入固定误差，结果是： episodes=20000, true_V0=4.2083, critic_V0=4.5083 lambda mean bias std mse 0.00 3.3133 -0.8950 0.7941 1.4315 0.50 3.8265 -0.3818 0.9201 0.9924 0.95 4.1767 -0.0316 1.8754 3.5181 1.00 4.2130 0.0048 2.2029 4.8529 $\lambda$ 从 0 增大到 1，偏差从 -0.8950 降到接近 0，标准差却从 0.7941 增至 2.2029。这个设定里 $\lambda=0.5$ 的均方误差最低；换一个 critic 误差和奖励噪声，最优点也会移动。这正是“偏差—方差折中”的可观测含义。 图 1：同一随机环境与 critic 误差下，$\lambda$ 改变了偏差和方差的配比。曲线由 make_figures.py 根据 20000 条固定随机种子轨迹生成。 4.2 看懂裁剪的符号 code/ppo_clip_minimal.py 先列出单样本裁剪表。下面四行足以解释 min： ratio A ratio*A clip(ratio)*A min clipped? 1.35 1.0 1.350 1.200 1.200 yes 0.65 1.0 0.650 0.800 0.650 no 1.35 -1.0 -1.350 -1.200 -1.350 no 0.65 -1.0 -0.650 -0.800 -0.800 yes 第一行是好动作已经涨太多，因此收益停在 1.2。第二行是好动作被错误地下调，目标仍取更差的 0.65。第三行是坏动作被错误地上调，惩罚保留 -1.35。第四行才是坏动作下降过多后停止继续奖励。clip 只截断“继续沿正确方向冲得更远”的激励。 图 2：蓝色阴影是 $[1-\epsilon,1+\epsilon]$。正优势在右侧形成平台，负优势在左侧形成平台，另一侧继续保留惩罚。 脚本还构造了一个两臂老虎机：旧策略给好、坏动作各 0.5 概率，同一批数据更新 40 轮。不裁剪时，好动作概率一路升到 0.9492，新旧策略 KL 达 0.8231；PPO-Clip 在第三轮越过 1.2 附近后梯度归零，停在 0.6097，KL 为 0.0246： mode epoch p(good) ratio_good grad KL(old||new) unclipped 3 0.6097 1.2193 0.4890 0.0246 unclipped 40 0.9492 1.8984 0.0990 0.8231 PPO-Clip 3 0.6097 1.2193 0.4890 0.0246 PPO-Clip 40 0.6097 1.2193 0.0000 0.0246 为什么停在 1.2193 而不是精确的 1.2？因为脚本用有限学习率做离散更新，第三步从区间内跨到了区间外，跨过以后才失去梯度。真实神经网络也会出现这种越界，所以还要监控 approximate KL，并在过大时提前停止 epoch。 05. 工业级实现对照 参考实现以 2026-09-08 的上游主分支为准。知识树原来记录的 huggingface/trl/trl/trainer/ppo_trainer.py 已经不在 TRL 当前主分支；当前 TRL 的 trainer 目录以 GRPO、DPO、RLOO 等实现 为主。因此本文把仍在维护、能直接核对 PPO 细节的 Stable-Baselines3 PPO.train 作为代码锚点，同时用 TRL GRPOTrainer 观察 LLM 对齐算法如何改写这些组件。 最小脚本与工业实现的差别主要有六处： rollout buffer：工业实现保存 observation、action、old log-prob、value、reward、termination mask，并在采样结束后统一算 GAE。old log-prob 必须冻结，不能在每个 epoch 重算。 优势标准化：常见实现按 minibatch 或整批做 $(A-\mu)/(\sigma+\varepsilon)$。它通常改善数值尺度，却会让一个样本的梯度依赖同批其他样本；batch 太小还可能产生不稳定统计量。 多轮 minibatch：Stable-Baselines3 的默认参数明确包含 n_epochs=10、gamma=0.99、gae_lambda=0.95、clip_range=0.2。这些是基准默认值，不应脱离任务直接复制。 价值函数处理：actor 与 critic 常共享 backbone，并用 vf_coef 合并损失；一些实现还裁剪 value 更新。value clipping 的尺度受奖励缩放影响，不能把 policy clip 的 $\epsilon$ 无脑复用。 额外护栏：除了 ratio clip，还会使用梯度范数裁剪、熵正则、approximate KL、target KL early stop、学习率退火和数值异常检查。Stable-Baselines3 的参数说明也明确指出，clip 本身不足以保证更新一定很小。 LLM 的 token mask：prompt token、padding 和 completion token 必须分开。策略损失通常只落在生成 token 上；序列末端奖励要变成 token 级回报，KL 往往逐 token 计算。错误的 mask 会把 padding 当动作，或者让 prompt 本身参与优化。 在 RLHF 里还要额外记录 reward/score、reward/non_score_reward、policy/approxkl、policy/clipfrac、loss/value、熵、响应长度和终止比例。只看总 reward 上升无法判断是任务能力提高、KL 惩罚变化，还是模型学会钻奖励与长度的空子。 5.1 指标应该怎样联读 clipfrac 必须先确认实现口径：常见日志统计 $|r_t-1|>\epsilon$，但真正进入目标平台还要满足“正优势且 ratio 过大”或“负优势且 ratio 过小”。越界比例不等于梯度被截断比例。它长期接近 0，可能说明学习率过小、epoch 太少、优势太弱，也可能只是策略已经接近局部平稳；它突然接近 1，说明新旧动作概率已大幅偏离，应结合优势符号与 KL 判断；错误方向的越界样本仍有纠正梯度。单独追求某个“漂亮”的 clipfrac 没有意义，应与 KL、entropy 和 reward 一起看。 approximate KL 快速增大而 reward 没有改善，通常先检查学习率、更新 epoch、优势尺度和 mask。如果 KL 很小、entropy 很快下降，则可能是分布变化集中在少数关键 token，平均 KL 把局部坍缩稀释了，需要再看 token 级分位数或最大值。 critic loss 下降也不一定是好消息。若 explained_variance 仍接近 0，critic 可能只学会了回报均值；若训练 loss 很低而新 rollout 上的 value error 很高，critic 在记忆同批噪声。此时 actor 收到的优势基线并不可靠。可以检查优势均值、标准差、与 return 的相关性，并把 value 网络的学习轮数与 actor 分开调节。 语言模型还有一个常见组合信号：任务 reward 上升、KL 惩罚绝对值变大、回答长度持续增加、人工质量不升。它常意味着奖励模型偏好长度或某种格式，策略在用分布漂移换分。解决顺序应是先按长度和题型切片评估奖励，再检查 EOS 奖励与 truncation，最后才调整 $\beta$。只增大 KL 系数会把症状压住，却不能修复有偏的反馈。 5.2 一个最小验收清单 正式放大训练前，可以用小 batch 做四个不依赖最终 reward 的验收。新旧模型完全一致时，所有有效位置的 ratio 应接近 1，approximate KL 应接近 0；把优势全设为 0 后，policy loss 对 actor 的梯度应为 0；交换优势正负号后，选中动作的更新方向应反转；只改变 padding 内容而保持 mask 不变时，损失应不变。这四项能抓住大部分 silent broadcasting、旧概率未冻结和 mask 泄漏问题。 06. 代价与边界 PPO 的优点是实现直接、可对 actor 和 critic 使用普通一阶优化器、同一批 rollout 可以复用多个 epoch。它的代价也很具体。 第一，PPO 仍是 on-policy 算法。ratio 能容忍有限的策略变化，却不能把很久以前的 replay data 变成可靠的当前策略数据。生成模型的 rollout 很贵，这也是许多对齐算法试图绕开在线 PPO 的原因。 第二，critic 会占显存、计算和调参预算。GAE 的质量由奖励噪声、value 误差、终止处理共同决定。critic 拟合得太慢，优势带偏；拟合得太快，又可能记住同一批高噪声回报。 第三，clip 是代理目标的启发式约束。它不等价于 TRPO 的显式 trust region，也不保证真实期望回报单调上升。大 batch、较小学习率、较少 epoch 和 target KL 能降低风险，仍需用新 rollout 验证。 第四，奖励尺度会渗透到整个系统。优势标准化能消除一部分尺度问题，critic loss、value clipping、梯度竞争和终止奖励仍受影响。LLM 的序列长度还会改变 token 级 KL 总量，导致长回答受到更大惩罚或获得更多优化机会。 以下情况不适合直接上 PPO：没有可信在线奖励或无法持续采样；动作空间可枚举且能直接做监督式偏好优化；离线数据远多于在线交互预算；安全约束必须严格满足而不能只靠软惩罚；环境极端稀疏奖励且 critic 没有可学习信号。此时应先解决反馈、数据或约束建模，而不是把希望寄托在 clip 上。 调试时也应先分清“估计坏了”还是“优化走远了”。前者常表现为优势高噪声、critic 解释方差差、不同随机种子更新方向不稳定，应检查奖励、bootstrap 和 GAE；后者常表现为前几个 minibatch 正常，随后 ratio、KL、clipfrac 一起升高，应减少学习率、epoch 或启用 KL early stop。两类问题都会造成 reward 抖动，但修复手段完全不同。 还要保留旧策略下的原始 rollout 指标。只比较更新后的 loss 无法回答“策略是否真的更好”；每轮更新后用新策略重新采样，并按任务、长度和难度切片比较回报，才能排除同一批数据上的代理目标过拟合。至少运行多个随机种子，因为一次 rollout 的偶然高分足以让小样本实验得出相反结论。 07. 经典论文脉络 REINFORCE：Williams 在 1992 年的 Simple Statistical Gradient-Following Algorithms for Connectionist Reinforcement Learning 给出基于采样回报和 log-derivative 的经典策略梯度形式，优点是通用，主要问题是方差高。 TRPO：Trust Region Policy Optimization（2015）用 KL 约束和二阶近似限制策略更新，为“旧数据上走多远”给出更严格的处理，但实现复杂。 GAE：High-Dimensional Continuous Control Using Generalized Advantage Estimation（2015）把 TD($\lambda$) 思想用于优势估计，明确控制 bias 与 variance。 PPO：Proximal Policy Optimization Algorithms（2017）用 clip 或 KL penalty 构造易实现的一阶替代目标，让同一批数据可做多轮 minibatch 更新。 InstructGPT：Training language models to follow instructions with human feedback（2022）展示了 SFT、奖励模型和 PPO 组合成语言模型 RLHF 流水线的代表性实践，也让参考模型 KL 成为对齐工程中的常见组件。 这条演进线并不是“新论文淘汰旧论文”。REINFORCE 给出梯度来源，GAE 改善估计，TRPO/PPO 限制更新，RLHF 再把状态、动作、奖励和约束映射到序列模型。后续算法通常只替换其中一到两块。 08. 常见误解 误解一：优势为正就说明奖励为正 优势是相对量。一个回答奖励为 -2，如果当前 prompt 下平均只能拿 -5，它的优势仍可能为正；一个回答奖励为 9，如果该状态平均是 9.5，它的优势反而为负。策略更新比较的是条件基线，不是绝对分数。 误解二：把优势乘常数不会影响训练 纯策略梯度方向不变，但真实 PPO 还有固定学习率、clip、value loss、熵和 KL 项。优势尺度改变后，各损失的相对权重、越过 clip 边界的速度都会改变。优势标准化是训练定义的一部分，不能只当日志美化。 误解三：ratio 超出区间后一定没有梯度 只有“超出区间且方向有利”的样本会被截断。好动作概率下降得太多、坏动作概率上升得太多时，目标仍保留梯度去纠正。04 节的四行表比背公式更可靠。 误解四：PPO clip 就是 KL trust region clip 在采样动作上限制替代收益，KL 衡量整个动作分布的变化。一个低概率 token 的概率可以相对变化很多而对平均 KL 贡献有限；许多小变化也可能累积出较大 KL。工程实现常同时监控 ratio、clip fraction 和 KL。 误解五：episode 截断等于终止 死亡、成功等真实终止意味着后续价值为零；时间上限只是观察窗口结束，底层状态可能仍有价值。两者都清零 bootstrap 会使 rollout 尾部的 value target 偏低。在固定最大生成长度的 LLM 训练中，这一点对应 EOS 与被长度上限截断的区别。 误解六：old policy 和 reference policy 是同一个概念 old policy 是一轮 PPO 的行为策略快照，几轮 minibatch 后就会更新；reference policy 是长期锚点，通常在 RL 开始时冻结。前者服务于重要性采样，后者服务于行为约束。 09. 动手验证 先在仓库根目录运行： python3 outputs/fundamentals_files/policy_gradient/code/gae_minimal.py python3 outputs/fundamentals_files/policy_gradient/code/ppo_clip_minimal.py 接着做三个小改动，每次只改一个变量并记录输出。 把 gae_minimal.py 的 critic_error 全部设为 0。你会看到 $\lambda=0$ 的偏差大幅下降；奖励噪声不变时，小 $\lambda$ 的低方差优势变得更有吸引力。 把奖励噪声标准差从 0.8 改为 0.1。完整 Monte Carlo target 的方差会明显下降，$\lambda=1$ 的代价随之减小。 把 ppo_clip_minimal.py 的学习率从 0.3 改为 0.03。PPO-Clip 会更接近 ratio=1.2 后才停住，说明 clip 边界不是参数投影，离散优化仍会跨界。 最后给脚本加入 target_kl=0.02 的提前停止条件：每轮更新后计算 KL(old||new)，超过阈值就停止。比较它与 ratio clip 的停止轮数。你会看到两种护栏观察的是不同量，也会理解工业实现为什么常把它们同时保留。 10. 延伸阅读 读完这篇可以继续看： 从 DPO 到 GRPO：去掉价值网络：比较偏好优化、组内相对优势与 PPO critic，追踪它们分别改了哪一项。 RLHF 工程实践：继续研究 token mask、参考模型 KL、奖励白化、长度偏置与分布式 rollout。 视频生成中的强化学习与奖励模型：把策略梯度和相对优势迁移到扩散与 flow matching 生成过程，观察奖励黑客如何出现。 回到开头那三个问题，现在可以逐项检查任何新对齐算法：它用什么分配信用，用什么基线或优势降低噪声，用什么机制约束策略变化。只要这三项说清楚，新名词就不会遮住算法真正改变的地方。 附录：完整代码 09 节用到的脚本全文如下（gae_minimal.py、make_figures.py、ppo_clip_minimal.py）。复制到本地存成同名文件，按各脚本开头的依赖说明准备环境后即可运行。 gae_minimal.py #!/usr/bin/env python3 """用 NumPy 展示 GAE 递推，以及 lambda 的偏差—方差折中。 依赖：numpy。运行：python3 gae_minimal.py """ import numpy as np def generalized_advantage_estimate(rewards, values, dones, next_value, gamma, lam): """返回 TD 残差、GAE 优势和 value target，输入均为 shape [T]。""" rewards = np.asarray(rewards, dtype=np.float64) values = np.asarray(values, dtype=np.float64) dones = np.asarray(dones, dtype=np.float64) advantages = np.zeros_like(rewards) deltas = np.zeros_like(rewards) gae = 0.0 for t in reversed(range(len(rewards))): value_next = next_value if t == len(rewards) - 1 else values[t + 1] not_done = 1.0 - dones[t] deltas[t] = rewards[t] + gamma * not_done * value_next - values[t] gae = deltas[t] + gamma * lam * not_done * gae advantages[t] = gae return deltas, advantages, advantages + values def fixed_trajectory_demo(): rewards = np.array([0.0, 0.0, 1.0, 0.5]) values = np.array([0.40, 0.55, 0.70, 0.30]) dones = np.array([0, 0, 0, 1]) gamma = 0.99 print("=== 固定轨迹 ===") print(f"rewards.shape={rewards.shape}, values.shape={values.shape}") for lam in (0.0, 0.5, 0.95, 1.0): deltas, advantages, targets = generalized_advantage_estimate( rewards, values, dones, next_value=0.0, gamma=gamma, lam=lam ) print( f"lambda={lam:>4.2f} delta={np.round(deltas, 4)} " f"A={np.round(advantages, 4)} target={np.round(targets, 4)}" ) def bias_variance_demo(seed=7, episodes=20000): """比较 GAE value target 对真实 V(s_0) 的偏差、标准差和均方误差。""" rng = np.random.default_rng(seed) gamma = 0.99 reward_means = np.linspace(0.2, 0.9, 8) horizon = len(reward_means) dones = np.zeros(horizon) dones[-1] = 1.0 true_values = np.zeros(horizon) for t in reversed(range(horizon)): future = 0.0 if t == horizon - 1 else true_values[t + 1] true_values[t] = reward_means[t] + gamma * future # 故意给 critic 加一组固定误差，使短视的 TD target 有偏。 critic_error = np.array([0.30, -0.90, 0.35, -0.25, 0.20, -0.15, 0.10, -0.05]) approx_values = true_values + critic_error rewards_batch = rng.normal(reward_means, 0.8, size=(episodes, horizon)) print("\n=== lambda 的偏差—方差折中（估计 V(s_0)）===") print(f"episodes={episodes}, true_V0={true_values[0]:.4f}, critic_V0={approx_values[0]:.4f}") print("lambda mean bias std mse") for lam in (0.0, 0.5, 0.95, 1.0): estimates = np.empty(episodes) for i, rewards in enumerate(rewards_batch): _, advantages, targets = generalized_advantage_estimate( rewards, approx_values, dones, next_value=0.0, gamma=gamma, lam=lam ) estimates[i] = targets[0] bias = estimates.mean() - true_values[0] mse = np.mean((estimates - true_values[0]) ** 2) print(f"{lam:>6.2f} {estimates.mean():>8.4f} {bias:>8.4f} {estimates.std():>8.4f} {mse:>8.4f}") if __name__ == "__main__": fixed_trajectory_demo() bias_variance_demo() make_figures.py #!/usr/bin/env python3 """生成本文的 GAE 偏差—方差图与 PPO 裁剪曲线。 依赖：numpy、matplotlib。运行：python3 make_figures.py 图片写入相邻的 figures/ 目录。 """ from pathlib import Path import matplotlib.pyplot as plt import numpy as np from gae_minimal import generalized_advantage_estimate from ppo_clip_minimal import clipped_surrogate OUT_DIR = Path(__file__).resolve().parent.parent / "figures" def gae_tradeoff_figure(seed=7, episodes=20000): rng = np.random.default_rng(seed) gamma = 0.99 means = np.linspace(0.2, 0.9, 8) dones = np.zeros(len(means)) dones[-1] = 1.0 true_values = np.zeros(len(means)) for t in reversed(range(len(means))): true_values[t] = means[t] + gamma * (0.0 if t == len(means) - 1 else true_values[t + 1]) approx_values = true_values + np.array([0.30, -0.90, 0.35, -0.25, 0.20, -0.15, 0.10, -0.05]) reward_batch = rng.normal(means, 0.8, size=(episodes, len(means))) lambdas = np.array([0.0, 0.25, 0.5, 0.75, 0.95, 1.0]) bias, std, mse = [], [], [] for lam in lambdas: estimates = [] for rewards in reward_batch: _, _, targets = generalized_advantage_estimate( rewards, approx_values, dones, 0.0, gamma, lam ) estimates.append(targets[0]) estimates = np.asarray(estimates) bias.append(estimates.mean() - true_values[0]) std.append(estimates.std()) mse.append(np.mean((estimates - true_values[0]) ** 2)) fig, ax = plt.subplots(figsize=(10, 5.3)) ax.plot(lambdas, np.abs(bias), "o-", linewidth=2.5, label="Absolute bias") ax.plot(lambdas, std, "s-", linewidth=2.5, label="Standard deviation") ax.plot(lambdas, mse, "^-", linewidth=2.5, label="Mean squared error") ax.set(xlabel="GAE lambda", ylabel="Error scale", title="GAE: less bias usually costs more variance") ax.grid(alpha=0.25) ax.legend(frameon=False) fig.tight_layout() path = OUT_DIR / "gae_bias_variance.png" fig.savefig(path, dpi=180, facecolor="white") plt.close(fig) return path def ppo_clip_figure(): ratio = np.linspace(0.4, 1.6, 500) fig, axes = plt.subplots(1, 2, figsize=(11, 4.8), sharex=True) for ax, advantage in zip(axes, (1.0, -1.0)): raw, clipped, objective = clipped_surrogate(ratio, np.full_like(ratio, advantage)) ax.plot(ratio, raw, linestyle="--", linewidth=2, label="Unclipped") ax.plot(ratio, objective, linewidth=3, label="PPO objective") ax.axvspan(0.8, 1.2, color="#2f6df6", alpha=0.08) ax.axvline(1.0, color="black", linewidth=1, alpha=0.5) ax.set_title(f"Advantage = {advantage:+.0f}") ax.set_xlabel("new / old probability ratio") ax.grid(alpha=0.2) axes[0].set_ylabel("Per-sample surrogate") axes[0].legend(frameon=False) fig.suptitle("PPO-Clip is one-sided and depends on the advantage sign", fontsize=14) fig.tight_layout() path = OUT_DIR / "ppo_clip_sign.png" fig.savefig(path, dpi=180, facecolor="white") plt.close(fig) return path if __name__ == "__main__": OUT_DIR.mkdir(parents=True, exist_ok=True) for output in (gae_tradeoff_figure(), ppo_clip_figure()): print(f"saved: {output} ({output.stat().st_size / 1024:.1f} KiB)") ppo_clip_minimal.py #!/usr/bin/env python3 """用 NumPy 拆解 PPO-Clip，并在两臂老虎机上复用同一批数据更新多轮。 依赖：numpy。运行：python3 ppo_clip_minimal.py """ import numpy as np def clipped_surrogate(ratio, advantage, epsilon=0.2): ratio = np.asarray(ratio, dtype=np.float64) advantage = np.asarray(advantage, dtype=np.float64) unclipped = ratio * advantage clipped = np.clip(ratio, 1.0 - epsilon, 1.0 + epsilon) * advantage objective = np.minimum(unclipped, clipped) return unclipped, clipped, objective def clipping_table(): ratios = np.array([1.35, 0.65, 1.35, 0.65, 1.10, 0.90]) advantages = np.array([1.0, 1.0, -1.0, -1.0, 0.5, -0.5]) raw, clipped, objective = clipped_surrogate(ratios, advantages) print("=== 单样本裁剪表 ===") print(" ratio A ratio*A clip(ratio)*A min clipped?") for r, a, u, c, o in zip(ratios, advantages, raw, clipped, objective): flag = "yes" if not np.isclose(u, o) else "no" print(f" {r:>4.2f} {a:>4.1f} {u:>7.3f} {c:>7.3f} {o:>6.3f} {flag}") def sigmoid(x): return 1.0 / (1.0 + np.exp(-x)) def bandit_objective(theta, clipped): """旧策略两动作概率各 0.5；动作 1 优势 +1，动作 0 优势 -1。""" p_good = sigmoid(theta) ratios = np.array([p_good / 0.5, (1.0 - p_good) / 0.5]) advantages = np.array([1.0, -1.0]) if clipped: return clipped_surrogate(ratios, advantages)[2].mean() return np.mean(ratios * advantages) def finite_difference_gradient(theta, clipped, h=1e-5): return (bandit_objective(theta + h, clipped) - bandit_objective(theta - h, clipped)) / (2 * h) def bernoulli_kl(old_p, new_p): return old_p * np.log(old_p / new_p) + (1.0 - old_p) * np.log((1.0 - old_p) / (1.0 - new_p)) def optimize_same_batch(clipped, steps=40, learning_rate=0.3): theta = 0.0 snapshots = [] watch = {0, 1, 2, 5, 10, 20, 39} for step in range(steps): grad = finite_difference_gradient(theta, clipped) theta += learning_rate * grad if step in watch: p_good = sigmoid(theta) ratio_good = p_good / 0.5 snapshots.append((step + 1, p_good, ratio_good, grad, bernoulli_kl(0.5, p_good))) return snapshots def repeated_update_demo(): print("\n=== 在同一批旧策略数据上更新 40 轮 ===") print("mode epoch p(good) ratio_good grad KL(old||new)") for mode, clipped in (("unclipped", False), ("PPO-Clip", True)): for epoch, p_good, ratio, grad, kl in optimize_same_batch(clipped): print(f"{mode:<10} {epoch:>5d} {p_good:>7.4f} {ratio:>7.4f} {grad:>7.4f} {kl:>7.4f}") if __name__ == "__main__": clipping_table() repeated_update_demo() 更多 AIGC 论文解读，关注微信公众号「人工智能炼丹君」 每日更新 · 论文精选 · 深度解读 · 技术脉络 微信搜索 人工智能炼丹君 或扫描下方二维码关注

AIGC 基本功|序列并行与 Ring Attention-SP 序列并行与 Ring Attention 所属方向：分布式训练 ｜ 难度：高阶 ｜ 前置知识：张量并行与流水线并行、注意力显存账本、FlashAttention 关键词：序列并行、Ring Attention、长序列、通信与计算重叠、Ulysses 01. 为什么需要它 一个视频或长文被编码成 131072 个 token，hidden size 为 8192，仅一个 BF16 的 $[L,d]$ 张量就占 2 GiB。Transformer 每层不只保存一个这样的张量，还要保存归一化输入、Q/K/V、MLP 中间量、dropout 状态等；标准 attention 若显式保存 $[L,L]$ 分数矩阵，单头 BF16 就是 32 GiB，根本无法进入反向。 数据并行沿 batch 切，每个 rank 仍要处理完整序列；张量并行沿 hidden/head 切，序列长度 $L$ 仍完整复制；流水线并行沿层切，某一层的长序列峰值仍在。长上下文把瓶颈推到第四个维度：必须沿 sequence 切 token。 但“每卡拿一段 token”并不自动成立。MLP 和 LayerNorm 对每个 token 独立，本地就能算；self-attention 中，每个 query 必须看到所有 key/value。若 rank 0 只拿前 1/8 的 K/V，它算出的就成了局部 attention，数学结果已经改变。序列并行的核心问题，是怎样在不复制全部长序列、不近似 attention 的前提下交换必要信息。 目前最常见的两类答案是 Ulysses 与 Ring Attention。Ulysses 先按 sequence 分片，再用 all-to-all 把布局变成按 attention head 分片，让每个 rank 在少量 head 上看到完整 sequence；算完后再 all-to-all 变回来。Ring Attention 则让每个 rank 保留本地 query，K/V block 沿环传递，经过 $P$ 轮后本地 query 恰好见过全序列。 即使尚未阅读FlashAttention 文章，因此先补最小背景：FlashAttention 并不近似 softmax，它把 Q/K/V 分块，并用在线 softmax 的最大值、分母和加权和递推，避免把完整 $L\times L$ 注意力矩阵写回 HBM。Ring Attention 把同一个分块思想扩展到多设备：块不只从显存搬到片上 SRAM，也沿网络搬到下一个 rank。 本文预算脚本给出：$L=131072,d=8192,P=8$ 时，一个 BF16 序列张量从每卡 2 GiB 降到 0.25 GiB，理想降低 8 倍；每层每 rank 绕环接收 K/V 的单向数据约 3.5 GiB。显存不是白省的，它被换成了通信与 $P$ 轮依赖。 02. 最小可用理解 三句话建立框架： 序列并行把 $[B,L,d]$ 沿 $L$ 切成 $P$ 份，使 LayerNorm、dropout、MLP 等序列形激活每卡近似降为 $1/P$。 Ulysses 用 all-to-all 在 sequence shard 与 head shard 间转置；Ring Attention 让 K/V block 绕环流动，本地 Q 用在线 softmax逐块累积精确结果。 两者没有把全注意力的 $O(L^2d)$ 算术复杂度变成线性，只是把计算与内存分散到多卡，并用通信重叠争取接近线性扩展。 如果只记一个区别：Megatron 风格的 sequence parallel 主要切 LayerNorm/dropout 等非 TP 区域的激活；Ulysses/Ring Attention 属于长上下文的 attention 并行，真正处理完整 sequence 上的全局注意力。很多系统把后者称为 context parallel，以免名称混淆。 03. 数学推导 3.1 标准 attention 的账本 令 $Q,K,V\in\mathbb{R}^{L\times d}$，为简洁省略 batch 与 head 维。scaled dot-product attention： $$O=\operatorname{softmax}\left(\frac{QK^\top}{\sqrt d}\right)V$$ 分数矩阵 $S=QK^\top/\sqrt d$ 有 $L^2$ 个元素，算力约为 $O(L^2d)$。若将序列均分到 $P$ 个 rank，每卡持有 $Q_r,K_r,V_r\in\mathbb{R}^{L/P\times d}$。逐 token 的 MLP 可以直接本地算，但 attention 的本地输出需要： $$O_r=\operatorname{softmax}\left(\frac{Q_r[K_0;K_1;\ldots;K_{P-1}]^\top}{\sqrt d}\right)[V_0;V_1;\ldots;V_{P-1}]$$ 也就是说，本地 query 仍必须与所有 rank 的 K/V 相遇。区别只在“把所有 K/V 一次聚齐”还是“逐块流过”。 3.2 在线 softmax：Ring 的数学底座 对一行 query，把 key/value 分成若干块。处理前 $t$ 个块后，保存行最大值 $m_t$、指数和 $\ell_t$ 与未归一化加权和 $u_t$。新块分数为 $s$，块内最大值为 $m_b=\max(s)$，更新全局最大值： $$m_{t+1}=\max(m_t,m_b)$$ 为了把旧累计量换到新的指数基准，定义 $\alpha=\exp(m_t-m_{t+1})$，则： $$\ell_{t+1}=\alpha\ell_t+\sum_j\exp(s_j-m_{t+1})$$ $$u_{t+1}=\alpha u_t+\sum_j\exp(s_j-m_{t+1})v_j$$ 处理完所有块后： $$o=\frac{u_T}{\ell_T}$$ 每次最大值变大时，旧分母与旧分子同时乘 $\alpha$，所以比值语义不变。算法只需保存 $m,\ell,u$ 和当前 score tile，不需保存全矩阵。它与减最大值的稳定 softmax 完全等价，而不是近似。 3.3 Ring Attention 的 $P$ 轮 第 $r$ 个 rank 固定本地 $Q_r$，初始持有 $K_r,V_r$。第 0 轮计算本地块，第 1 轮把 K/V 发给下一个 rank并接收上一个 rank 的块，如此循环。第 $t$ 轮 rank $r$ 处理来源 $(r-t)\bmod P$ 的 K/V；$P$ 轮后每个 $Q_r$ 看过所有块。 本地计算量近似： $$C_{\mathrm{rank}}=O\left(\frac{L}{P}\cdot L\cdot d\right)=O\left(\frac{L^2d}{P}\right)$$ 每 rank 需要接收 $P-1$ 次大小约 $2(L/P)d\,b$ 的 K/V，其中 $b$ 是每元素字节数： $$V_{\mathrm{ring}}=2\frac{P-1}{P}Ldb$$ 当 $L$ 与 $P$ 同比例增长、每卡本地 token 数保持不变时，单轮的 Q/KV block 大小、计算量和传输量近似不变，增长的是轮数和每卡总计算量。是否能覆盖通信取决于单轮 block 是否足够大、有效算力与链路带宽；弱扩展本身不会自动改善单轮的重叠条件。若本地 block 太小、网络慢或 kernel 没有异步双缓冲，通信会裸露出来。 训练反向不能只把前向环倒放那么简单：每个 rank 还要计算对本地 Q、流动 K/V 的梯度，并把属于原拥有者的 dK/dV 沿环累积或归还。checkpoint 与重计算还能减少保存量，但增加计算。 3.4 因果 mask 如何穿过环 自回归 attention 要求 query 位置 $i$ 只能看 key 位置 $j\le i$。分块后，根据全局位置判断： K/V block 完全位于 Q block 未来：整块跳过； 完全位于过去：无需 mask； 与 Q block 重叠：块内施加三角 mask。 跳过未来块可省计算，但不同 rank 的有效工作量会不均。某些实现使用 zigzag 或重新排列 token，使每卡同时持有前后位置，平衡 causal attention 工作量。若只看非因果公式估算吞吐，部署到 causal LLM 时可能偏差很大。 3.5 Ulysses 的 all-to-all 转置 设输入布局是 sequence-sharded： $$[B,L/P,H,D_h]$$ $H$ 是 head 数，$D_h$ 是每头维度。Ulysses all-to-all 把它转成： $$[B,L,H/P,D_h]$$ 每卡序列完整、head 只剩 $H/P$，于是能在本地执行普通 attention；输出再 all-to-all 回 sequence shard。它通常要求 $H$ 能被 $P$ 整除，若 GQA 只有很少 KV heads，约束会更紧。Ring 不要求按 head 数切，但要经历 $P$ 个顺序环步骤。 Ulysses 的 collective 容易利用成熟 all-to-all，但跨节点 all-to-all 对网络争用敏感；Ring 只与邻居点对点通信，更贴合环形拓扑，也更容易细粒度重叠。没有一种方法在所有硬件和 head 配置上恒优。 3.6 Megatron Sequence Parallel 不等于长上下文 attention 传统 TP 的 column/row parallel 边界往往让某些激活在 TP rank 上复制。Megatron sequence parallel 把 LayerNorm、dropout 等逐 token 操作沿 sequence 切分，并用 reduce-scatter 与 all-gather 衔接 TP 区域，使这些激活内存近似降为 $1/P$。它与 TP process group 绑定，主要减少非 TP 区域的重复激活。 但 attention 若仍按 head 做 TP，每个 head 仍可能看到完整 sequence。因此配置项 sequence_parallel=true 并不等于已经支持百万 token。需要进一步的 context parallel、Ulysses 或 Ring Attention 才能沿 attention 的上下文维扩展。 一个序列张量的每卡存储随 P 下降，而完整前向绕环接收的 K/V 字节增加并趋于上限。右图是接收量，不重复计发送量，尚未包含反向通信。 04. 代码实现 ring_attention_sim.py 只用标准库，实现完整 softmax attention 与分块在线递推。Q/K/V 均为 $[4,2]$，每个 K/V block 两行，相当于两轮 ring： Q=K=V shape=(4, 2), kv_block=2, ring_steps=2 full[0]=[1.713718770, 1.717687378] ring[0]=[1.713718770, 1.717687378] max_abs_diff=0.000e+00 本组小样本在打印精度内相同；算法等价性来自第 3.2 节的不变式，浮点实现一般仍存在求和顺序引起的舍入差异。脚本没有真正启动多进程；它把 K/V block 依次喂给同一递推器，模拟每个 rank 收到环上块后的本地数学。生产实现还要异步 send/recv、双缓冲、causal mask 与 backward。 sp_budget.py 对 $L=131072,d=8192,H=64,P=8$ 输出： L=131072, d=8192, heads=64, sp=8, dtype_bytes=2 one [L,d] tensor: replicated=2.00 GiB, sequence-sharded/rank=0.25 GiB ideal memory reduction for sequence-shaped activations=8.0x Ulysses head-divisible=True (64//8=8 heads/rank) Ring Attention rounds=8, local query tokens=16384 Ring K/V traffic per rank per layer (one direction, ideal)=3.50 GiB 2 GiB 只是一个张量，不是整层总显存；0.25 GiB 是理想静态 shard，也没算当前通信块、输出、梯度与工作区。3.5 GiB 对应 $2(P-1)Ldb/P$，是每层、每 rank、单方向的 K/V 接收量。实际网络还包含反向与协议开销。 05. 工业级实现对照 以 2026-09 为准，DeepSpeed 的 deepspeed/sequence/layer.py 中 DistributedAttention 用自定义 all-to-all 在 scatter 轴与 gather 轴间转换 Q/K/V，调用 local_attention 后再变回去。源码明确检查总 head 数必须能被 sequence parallel size 整除，并提供 stream 与 overlap handle 管理通信重叠。 这个工业实现比最小公式多处理 batch 维位置、不同 Q/K/V 布局、rotary position embedding、异步 stream、autograd 的反向 all-to-all 与进程组。布局索引错一位可能 shape 仍合法但语义错误，所以系统必须统一采用明确的张量维约定。 Megatron-LM 的模型并行配置中 sequence_parallel 用于并行 LayerNorm 与 dropout；context_parallel_size 则面向上下文切分。当前 Megatron Core 还支持多种 context parallel 通信方式。读配置时不要只看“SP”缩写，要追到具体张量在哪个轴分片以及 attention 内部采取哪种交换。 Ring Attention 的工业 kernel 通常建立两个 K/V buffer：当前块参与计算时，下一个块在独立通信 stream 接收；计算结束交换 buffer。要真正隐藏通信，必须满足： $$T_{\mathrm{attention\ block}}\ge T_{\mathrm{send/recv}}$$ 此外还要确保计算与 NCCL 没争抢同一资源到互相拖慢。仅在代码里使用 async op 不代表 timeline 上已经重叠，必须用 profiler 验证。 FlashAttention 提供单卡块级 IO 优化，Ring/Ulysses 提供多卡数据布局。两者不是竞争关系：每个 rank 的 local attention 往往正由 FlashAttention kernel 完成。前者减少 HBM 读写，后者减少单卡需要常驻的全局序列。 checkpoint 也必须保存并行元数据。若训练时改变 SP/CP 度，参数本身或许未沿 sequence 切，但 optimizer、随机数状态和位置相关缓存可能变化。恢复前应由框架的 distributed checkpoint 层重新映射，而不是直接让每个 rank 读取旧编号文件。 06. 代价与边界 序列并行省的是序列形激活，不一定省参数和优化器状态；它通常要与 ZeRO/FSDP、TP、PP 组合。完整世界规模可能写成 $W=d_p t_p p_p c_p$，其中 $c_p$ 是 context parallel 度。每多一根轴，都新增 process group、拓扑映射与整除约束。 Ring Attention 保留全注意力的二次计算量。卡数增加能把每卡计算降到 $1/P$，但集群总 FLOPs 仍为 $O(L^2d)$；当上下文翻倍、卡数不变时，总计算约四倍。它突破的是单卡内存和 wall-clock 可扩展性，不是算法复杂度。 通信也不是“免费”。Ring 的每 rank 字节量随 $L$ 线性增长且有 $P$ 轮延迟链；Ulysses 有前后 all-to-all，容易受跨节点网络与并发任务干扰。若 $L$ 尚短，通信与布局转换可能比 attention 本身更贵，普通 TP 更合适。 常见边界包括： Ulysses 的 head 或 KV head 不能被 SP 度整除，GQA/MQA 尤其容易受限； causal mask 造成各 block 计算不均，需要 load-balanced ring； position encoding 必须使用全局位置，不能把每 rank 的局部 token 从 0 重新编号； dropout 的随机数要在切分前后保持统计与重算一致，否则并行度改变会造成难解释的差异； variable-length packed sequence 需要携带边界，不能让不同样本互相 attention； 视频 token 在时间、空间上的排列会影响 causal 或局部结构，切块策略不能只按连续内存方便决定。 在低带宽以太网集群上，增加 context parallel 度可能只是把 OOM 变成通信瓶颈。应先用 activation checkpoint、FlashAttention、减小 microbatch 等单机方法，确认仍由单卡序列容量限制，再引入跨设备 SP。 6.1 从单张量扩展到整层激活账本 一个 $[L,d]$ BF16 张量占 $2Ld$ 字节，但训练层内可能同时存在 residual、norm 输入、Q/K/V、attention output 与 MLP 中间值。设所有与序列线性相关且必须保存到反向的张量合计为 $cLd$ 个元素，序列分片后的理想每卡保存： $$M_{\mathrm{linear/rank}}\approx\frac{cLdb}{P}$$ $c$ 由架构、融合 kernel 与 checkpoint 策略决定，不能用固定常数套所有模型。SwiGLU 会产生门控与 value 两个支路；视频 DiT 还可能有条件分支、cross-attention 与调制参数。最可靠方法是从 profiler 的 allocation timeline 识别跨 backward 存活的张量。 attention tile 与通信 buffer 是额外峰值。Ring 双缓冲至少需要当前和下一份 K/V block；FlashAttention kernel 还需要 score tile、softmax statistic 和 workspace。于是： $$M_{\mathrm{peak/rank}}=M_{\mathrm{linear/rank}}+M_{\mathrm{local\ attention}}+M_{\mathrm{double\ buffer}}+M_{\mathrm{runtime}}$$ 当 $P$ 继续增加时，第一项下降，但 buffer、runtime 和最小 kernel workspace 不会等比下降，显存收益逐渐偏离 $1/P$。用总 allocated 除卡数预测极限会过于乐观。 6.2 反向为何比前向更难 前向中，本地 Q 只需依次看过所有 K/V block。反向要得到 $dQ_r$、$dK_j$ 与 $dV_j$。当来源 $j$ 的 K/V 到达 rank $r$ 时，本地可计算它对 $dQ_r$ 的贡献，也会产生属于来源 rank $j$ 的 $dK_j,dV_j$ 部分。后两者必须跨所有 query shard 求和并归还拥有者。 因此每个块需要身份信息，环中既可能传 K/V，也可能传对应梯度累计。为了减少前向保存，backward 常重算局部 score 与概率，依赖前向保存的行最大值和归一化分母。若随机 dropout 参与 attention，重算还必须恢复相同随机 mask。 理论通信量必须同时列前向与反向。只用 $2(P-1)Ldb/P$ 报告 K/V 前向流量会低估训练网络负担；真实实现还包括 dK/dV、可能的 dQ 归约、控制同步和其他模块的 collective。性能分析应以 profiler 中每层总通信字节与裸露时间为准。 6.3 通信重叠需要满足哪些条件 把一个本地 query block 与一个 K/V block 的 attention 计算时间记为 $T_c$，邻居传输一个 K/V block 的时间记为 $T_n=\alpha+S/\beta$，其中 $\alpha$ 是链路延迟，$\beta$ 是有效带宽，$S$ 是消息字节。双缓冲后的理想每轮时间： $$T_{\mathrm{round}}\approx\max(T_c,T_n)$$ 只有 $T_c\ge T_n$，网络才大部分隐藏。减小 block 能降低临时显存，却增加轮数和延迟占比；增大 block 提高 GEMM 效率，却增大 buffer。block size 是计算、带宽、延迟与显存的共同旋钮。 实现还需避免隐式同步。例如在循环中把 GPU 标量取回 CPU、过早 wait 通信 handle、复用尚未完成接收的 buffer，都会把异步流水串行化。正确做法是独立 stream、event 依赖与 ping-pong buffer，并在 GPU timeline 上确认 send/recv 与 attention kernel 真正重叠。 环的物理顺序同样重要。逻辑 rank 相邻却跨交换机，会让每轮走慢链路；合理 ring 应优先沿 NVLink 域，再通过少量跨节点边连接。多维集群还可以分层：节点内用 Ulysses all-to-all，节点间用 Ring，或反过来根据硬件选择，这正是统一序列并行方法的动机。 6.4 长度、head 与并行轴的联合约束 设 context parallel 度为 $c_p$、tensor parallel 度为 $t_p$。若 TP 已把 head 切为 $H/t_p$，Ulysses 再沿 head 分 $c_p$，就需要本地可见 head 数继续可分： $$H\ \bmod\ (t_pc_p)=0$$ 具体约束取决于框架是在同一 head 轴叠加还是采用不同布局。GQA 中应将 $H_{KV}$ 单独代入；只有 8 个 KV head 的模型很难做大 head-parallel。某些系统复制 K/V 以放松整除，但通信和显存模型随之改变。 sequence length 通常也要求能被 $c_p$ 与 block size 整除。不整除时可以 padding 或不等长 shard。padding 实现简单但会浪费二次 attention 计算；不等长 shard 又使通信消息与计算不均。对长度差异大的数据，先按 token 数分桶往往比依赖动态不等长 ring 更稳定。 TP 与 CP 同时使用还会出现多种 collective。attention 投影可能在 TP 组 all-reduce，context attention 在 CP 组 send/recv，层外梯度还在 DP 组 reduce-scatter。若各组在不同 rank 上以不一致顺序发起 collective，可能死锁。框架必须给跨组通信建立确定顺序或独立 stream 依赖。 6.5 视频与多模态序列的特殊问题 视频 token 常按时间、空间高、空间宽展平。连续切 sequence 可能让每个 rank 持有若干完整帧，也可能持有同一帧的空间条带；两种布局对位置编码、局部 attention、数据增强和负载平衡影响不同。若模型有 temporal attention 与 spatial attention 分解结构，最优切分轴也可能随层变化。 全局位置必须从原始坐标生成。RoPE 若把每个 rank 的局部索引重新从 0 开始，不会报 shape 错，却让不同块位置混叠。二维或三维 RoPE 更要携带时间与空间坐标，而不是只传一个线性 offset。跨模态序列还要保持文本、图像、视频 segment 的 mask 与边界。 视频长度和分辨率经常动态变化。同一 global batch 内若 token 数差异大，padding 使最长样本决定所有 rank 工作量；packed sequence 可减少浪费，却要求 block 不跨样本做 attention。调度系统最好按总 token 数而非样本数组 batch，并把实际有效 token 纳入 loss normalization。 因果结构也不总是标准下三角。视频生成可能使用双向视觉 attention、文本到视频 cross-attention 或分块因果时间 mask。Ring 优化中“未来块可跳过”的判断必须来自真实 mask 结构，不能硬编码 LLM 假设。 6.6 正确性验收与故障定位 第一层验收是 shape：每个边界记录 global shape、local shape、分片轴、global offset、padding 与拥有者。第二层是数学：小尺寸 FP64 单卡作为 oracle，逐项比较 forward 和 dQ/dK/dV。第三层是随机性：开启 dropout 与 checkpoint 后，在固定 seed 下比较统计和可复现范围。 输出不一致时，可按轮 dump 每个 rank 处理的 K/V 来源编号。若少一个或重复一个，是 ring 调度；若所有来源齐全但数值错，检查在线 softmax 的旧累计重标定；若仅 causal 错，检查全局位置与块分类；若 forward 对而 backward 错，检查 dK/dV 归还与跨 query shard 求和。 OOM 时不要只减 block。先区分 persistent activation、通信双 buffer、kernel workspace 与碎片。若 activation 占主导，提高 CP 或 checkpoint 有效；若 buffer 占主导，减 block 才有效；若碎片主导，可调分配器或稳定 shape。错误手段可能降低吞吐却完全不改变峰值来源。 性能验收同时做 strong scaling 与 weak scaling。strong scaling 固定全局 $L$，增加卡数，观察端到端时间能否下降；容量弱扩展固定每卡 local token，令 $L\propto P$：精确全注意力的每卡计算 $L^2/P$ 此时随 $P$ 线性增长，不能期望每卡时间不变；应比较实测时间与这一增长基线，以及通信隐藏率。Ring Attention 可随设备数扩展可处理上下文，但总计算也随之增加。 最后把精度也纳入验收。在线 softmax 很稳定，但不同 block 顺序改变浮点加法；BF16、FP8 attention 更明显。应比较最终任务指标、loss 轨迹和梯度，而不是要求多卡输出 bitwise 等同。若误差远超单纯归约顺序，再排查 scale、mask 与布局。 6.7 Ulysses、Ring 与混合方案如何选 先看 head 约束。若本地 Q head 与 KV head 都能被目标 SP 度整除，且集群 all-to-all 带宽高，Ulysses 路径直接，local attention 可复用成熟 kernel。若 KV head 很少、SP 度很大，Ring 不沿 head 切，通常更容易扩展。再看拓扑：全连接高速交换更适合 all-to-all；邻接带宽强、跨组带宽弱时，环形点对点更自然。 再看本地块计算。长序列、大 head dimension 使每轮计算足以隐藏 K/V 传输，Ring 受益；较短序列或大量小 head 时，$P$ 轮延迟可能突出。Ulysses collective 次数少，但消息全局交换可能造成拥塞。决策依据应是每轮计算时间、有效链路带宽和 collective 实测，不是理论总字节一项。 混合方案把设备组织成二维网格。例如节点内 8 卡做 Ulysses，节点间若干组做 Ring；这样 head 整除只限制节点内度数，跨节点使用邻接通信。代价是布局更复杂、process group 更多，Q/K/V 需要在两个维度间保持一致顺序。只有单一方案确实受限时才值得引入。 还要把 TP 纳入考虑。如果 TP 已经切 head，Ulysses 可用 head 更少；有时降低 TP、提高 CP 更适合长序列，有时模型单层容量又要求 TP 不能降。最优并行度会随训练长度变化，短序列预训练与长序列继续训练未必使用同一布局。 6.8 推理场景与训练不同 prefill 处理整段 prompt，计算形态接近训练前向，序列并行能分摊长 prompt attention。decode 每步只有少量新 query，却要读取不断增长的 KV cache；此时瓶颈常是 KV 带宽与跨卡延迟，而非大块 GEMM。训练中优秀的 Ring block，在逐 token decode 里可能太细、同步轮数太多。 推理还涉及请求级并发。可以把不同请求分给数据并行副本，也可把一个超长请求沿 context 分片。前者吞吐高、单请求受单卡长度限制；后者支持超长上下文、占用多个设备。调度器应根据请求长度动态选择，而不是所有请求固定占一个 CP 组。 KV cache 的拥有关系必须稳定。若每 rank 保存一段 sequence，新增 token 的 K/V 放在哪一段、何时重平衡、beam search 如何复制，都要定义。Paged KV cache 与 context parallel 叠加时，逻辑页、物理 GPU 与全局位置形成三层映射，错误可能表现为偶发内容质量下降而非崩溃。 因此本文通信公式主要用于训练或 prefill，不能直接拿来预测 decode tokens/s。推理基准要分别报告 time-to-first-token、inter-token latency、并发吞吐、KV cache 容量和长短请求混合表现。 6.9 上线前检查清单 上线前确认：sequence、head、KV head 与 block 的整除或 padding 规则明确；每个 shard 保存全局位置；causal/packed mask 按全局索引生成；每轮 K/V 来源不重不漏；在线 softmax 保存并正确重标定最大值、分母与分子；backward 的 dK/dV 回到原拥有者；dropout 重算保持随机一致。 系统侧确认 ring 物理顺序符合拓扑，通信与计算 timeline 真重叠，buffer 生命周期无覆盖，所有 rank collective 顺序一致，checkpoint 记录分片元数据，变更 CP 度可以恢复。性能侧同时做 strong/weak scaling，并在真实长度分布而非单一最大长度上测试。 最后保留小尺寸 oracle。每次升级 attention kernel、通信库、RoPE 或 mask 实现，都自动生成随机 Q/K/V，与 FP64 单卡结果比较 forward 和全部梯度；再跑包含因果、变长、GQA、极端分数和非整除长度的边界集。长序列错误代价很高，小 oracle 是最便宜的保险。 6.10 最后的边界：更长不等于更有用 系统能够处理百万 token，只说明容量与运行时间可接受，不保证模型会利用远距离信息。长上下文实验还要检查位置外推、训练长度分布、检索准确率与有效注意范围。并行系统解决“算得出来”，数据与模型设计决定“学得会不会”。 因此扩展长度时同时保留短上下文质量基线，按距离分桶评估信息召回，并报告有效 token 而非 padding 后长度。否则可能用大量 GPU 计算模型并未使用的上下文。 07. 经典论文脉络 FlashAttention（arXiv:2205.14135）：用 IO-aware tiling 与在线 softmax 实现精确 attention，避免物化完整分数矩阵，是多卡分块 attention 的本地 kernel 基础。 Reducing Activation Recomputation in Large Transformer Models（arXiv:2205.05198）：提出 Megatron sequence parallel 与选择性重计算，减少 TP 区域之外的重复激活。 DeepSpeed Ulysses（arXiv:2309.14509）：用 all-to-all 在 sequence 与 head 布局间转换，使超长序列 attention 可扩展。 Ring Attention with Blockwise Transformers（arXiv:2310.01889）：让 K/V 块沿环流动并与分块计算重叠，把可处理上下文随设备数扩展。 USP（arXiv:2405.07719）：统一 Ulysses 与 Ring 两类序列并行，针对模型结构与网络拓扑组合两者。 这条演进线从“单卡不保存 $L^2$ 矩阵”，走到“多卡不复制 $L$ 激活”，再走到按 head 结构和网络拓扑组合通信方式。算法、kernel 与系统三层缺一不可。 08. 常见误解 误解一：序列并行把 attention 复杂度从二次降成线性。 精确全注意力的集群总计算仍为 $O(L^2d)$；它只把工作分摊，并让内存近似随本地 token 数增长。 误解二：每卡切一段 token 后直接做本地 attention。 那只允许 token 看同一分片，会改变模型。必须让 query 见到全局 K/V，或明确接受局部/稀疏近似。 误解三：Ring Attention 是近似注意力。 在线 softmax 递推在精确算术下与完整 softmax 等价；差别来自浮点归约顺序，不是丢弃连接。 误解四：开启 Megatron sequence_parallel 就解决长上下文。 它主要减少 LayerNorm、dropout 等激活复制；长上下文 attention 通常还需 context parallel。 误解五：异步 send/recv 就等于通信被隐藏。 只有计算时间覆盖传输、stream 真能并发且无资源争抢，timeline 上才会重叠。 误解六：Ulysses 的 SP 度只受 GPU 数限制。 head 数、KV head 数与 tensor layout 必须可切；GQA 模型常比标准 MHA 更早碰到上限。 09. 动手验证 先运行 ring_attention_sim.py，把 block_size 改成 1、2、4。预期 max_abs_diff 始终接近机器精度；再把 score 全部加 1000，稳定递推仍有限，而直接 exp(score) 会溢出。这同时验证精确性与减最大值的重要性。 为脚本加入 causal mask，使用 query/key 的全局索引屏蔽 $j>i$。预期分块结果与完整下三角 attention 一致。然后故意用局部索引 mask，会发现后续 rank 错误允许关注未来或错误屏蔽过去，这是真实分布式实现中很隐蔽的 bug。 运行 sp_budget.py，固定每卡 token 为 16384，让 $L$ 与 ranks 同比例从 1、2、4、8 增长。每卡序列张量保持 0.25 GiB，而 ring K/V 流量与全局 $L$ 增长。这个弱扩展实验说明“可放下更长序列”并不等于“通信不增长”，只是计算增长可能覆盖通信。 在真实集群比较 FlashAttention 单卡、Ulysses、Ring/context parallel：固定模型与 global sequence，记录峰值显存、attention kernel 时间、collective/点对点时间、重叠率与 MFU。分别在节点内和跨节点运行；预期最佳方法随 NVLink、InfiniBand、head 数和 local block 大小变化。 最后做数值验收：关闭 dropout，用同一 Q/K/V 比较单卡与多卡输出及 dQ/dK/dV。BF16 下用合理容差，同时单独检查每个 rank 拼接后的全局顺序。输出对而梯度错，通常说明 backward 中流动 K/V 梯度没有正确归还拥有者。 10. 延伸阅读 读懂本文的完整前置路径是： 数据并行与 ZeRO 显存切分：先区分状态分片与计算分片（已发布）。 张量并行与流水线并行：理解 process group、collective、microbatch 与拓扑（已发布）。 自注意力机制的计算与显存账本：理解 $L^2$ 分数矩阵从哪里来。 FlashAttention 为什么不需要存下注意力矩阵：深入单卡在线 softmax 与 IO 复杂度。 DDPM 训练目标与采样流程和 DiT：用 Transformer 替掉 UNet：理解视频 DiT 的时空 token 为何迅速推高注意力成本。 选 Ulysses 还是 Ring，不应从名字出发。先列出 $L,H,H_{KV},d$、网络拓扑、每卡内存与 causal 结构；再计算 local shape、通信字节、轮数和可重叠窗口；最后用 profiler 验证。长上下文系统真正的能力，是让数学等价、内存边界和物理网络三者同时对齐。 附录：完整代码 09 节用到的脚本全文如下（ring_attention_sim.py、sp_budget.py、make_figures.py）。复制到本地存成同名文件，按各脚本开头的依赖说明准备环境后即可运行。 ring_attention_sim.py #!/usr/bin/env python3 """Exact blockwise attention recurrence used by ring-style attention; stdlib only.""" import math def dot(a, b): return sum(x * y for x, y in zip(a, b)) def full_attention(q, k, v): outputs = [] scale = math.sqrt(len(q[0])) for query in q: scores = [dot(query, key) / scale for key in k] peak = max(scores) weights = [math.exp(score - peak) for score in scores] denom = sum(weights) outputs.append([sum(w * value[j] for w, value in zip(weights, v)) / denom for j in range(len(v[0]))]) return outputs def blockwise_attention(q, k, v, block_size): scale = math.sqrt(len(q[0])) peak = [-math.inf] * len(q) denom = [0.0] * len(q) numer = [[0.0] * len(v[0]) for _ in q] for start in range(0, len(k), block_size): kb, vb = k[start:start + block_size], v[start:start + block_size] for row, query in enumerate(q): scores = [dot(query, key) / scale for key in kb] new_peak = max(peak[row], max(scores)) correction = math.exp(peak[row] - new_peak) if peak[row] != -math.inf else 0.0 weights = [math.exp(score - new_peak) for score in scores] denom[row] = correction * denom[row] + sum(weights) numer[row] = [correction * old + sum(w * value[j] for w, value in zip(weights, vb)) for j, old in enumerate(numer[row])] peak[row] = new_peak return [[x / denom[row] for x in numer[row]] for row in range(len(q))] q = [[1.0, 0.0], [0.0, 1.0], [1.0, 1.0], [-1.0, 1.0]] k = [[1.0, 0.0], [0.0, 1.0], [1.0, 1.0], [1.0, -1.0]] v = [[1.0, 2.0], [2.0, 0.0], [0.0, 3.0], [4.0, 1.0]] full = full_attention(q, k, v) streamed = blockwise_attention(q, k, v, block_size=2) max_diff = max(abs(a - b) for ra, rb in zip(full, streamed) for a, b in zip(ra, rb)) print("Q=K=V shape=(4, 2), kv_block=2, ring_steps=2") print("full[0]=[" + ", ".join(f"{x:.9f}" for x in full[0]) + "]") print("ring[0]=[" + ", ".join(f"{x:.9f}" for x in streamed[0]) + "]") print(f"max_abs_diff={max_diff:.3e}") sp_budget.py #!/usr/bin/env python3 """Compare idealized per-rank sequence-memory and method constraints.""" import argparse parser = argparse.ArgumentParser() parser.add_argument("--tokens", type=int, default=131072) parser.add_argument("--hidden", type=int, default=8192) parser.add_argument("--heads", type=int, default=64) parser.add_argument("--ranks", type=int, default=8) args = parser.parse_args() if args.tokens % args.ranks: parser.error("tokens must be divisible by ranks") tensor_gib = args.tokens * args.hidden * 2 / 1024**3 local_gib = tensor_gib / args.ranks print(f"L={args.tokens}, d={args.hidden}, heads={args.heads}, sp={args.ranks}, dtype_bytes=2") print(f"one [L,d] tensor: replicated={tensor_gib:.2f} GiB, sequence-sharded/rank={local_gib:.2f} GiB") print(f"ideal memory reduction for sequence-shaped activations={args.ranks:.1f}x") print(f"Ulysses head-divisible={args.heads % args.ranks == 0} ({args.heads}//{args.ranks}={args.heads // args.ranks} heads/rank)") print(f"Ring Attention rounds={args.ranks}, local query tokens={args.tokens // args.ranks}") print(f"Ring K/V traffic per rank per layer (one direction, ideal)={2*tensor_gib*(args.ranks-1)/args.ranks:.2f} GiB") make_figures.py """Regenerate this article's deterministic teaching figure (numpy + matplotlib).""" from pathlib import Path import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt import numpy as np OUT=Path(__file__).resolve().parents[1]/"figures" OUT.mkdir(exist_ok=True) plt.rcParams.update({"font.sans-serif":["PingFang SC","Arial Unicode MS","DejaVu Sans"],"axes.unicode_minus":False}) p=np.array([1,2,4,8,16,32]); L=131072;d=8192;b=2;g=1024**3 fig,axes=plt.subplots(1,2,figsize=(10,4.4)) axes[0].plot(p,np.full(p.shape,L*d*b/g),"--",label="Replicated") axes[0].plot(p,L*d*b/p/g,"o-",label="Sequence shard") axes[0].set(xscale="log",xlabel="Ranks P",ylabel="One tensor per rank (GiB)");axes[0].legend() axes[1].plot(p,2*(p-1)/p*L*d*b/g,"o-") axes[1].set(xscale="log",xlabel="Ranks P",ylabel="Forward K/V received per rank (GiB)") for ax in axes:ax.grid(alpha=.25) fig.suptitle("Ring attention budget: L=131072, d=8192, BF16");fig.tight_layout() fig.savefig(OUT/'ring_budget.png',dpi=170) plt.close(fig) print(OUT/'ring_budget.png') 更多 AIGC 论文解读，关注微信公众号「人工智能炼丹君」 每日更新 · 论文精选 · 深度解读 · 技术脉络 微信搜索 人工智能炼丹君 或扫描下方二维码关注

AIGC 基本功|张量并行与流水线并行-TP-PP 张量并行与流水线并行 所属方向：分布式训练 ｜ 难度：高阶 ｜ 前置知识：数据并行与 ZeRO 显存切分 关键词：张量并行、流水线并行、Megatron-LM、GPipe、通信开销、并行策略 01. 为什么需要它 ZeRO-3 能把参数、梯度和优化器状态切到数据并行 rank 上，但每一层计算时仍要聚合这一层的参数。若单个 Transformer 层、词表投影或大矩阵乘本身就放不进一张卡，单纯增加数据并行度解决不了问题；若层能放下，但整网常驻状态或激活放不下，也可考虑沿网络深度切开；能否仅靠 ZeRO-3 和重计算解决，需要先算峰值。 这对应两把不同的刀。张量并行（Tensor Parallelism，TP）在一层内部切矩阵，让多张卡共同完成同一个 GEMM；流水线并行（Pipeline Parallelism，PP）沿层切模型，让不同设备各自保存连续或交错的一段层。前者解决“这一层太宽”，后者解决“模型太深”。数据并行沿 batch 复制模型，三者切的是三个不同维度。 先看一个浪费现场：4 个 pipeline stage 处理 1 个 microbatch。第 0 段前向时其余三段都空闲；数据流到第 3 段后，前几段又空闲。若把一个 global batch 切成 8 个 microbatch，让不同样本在各 stage 重叠，理想 forward sweep 从 $4\times8=32$ 个串行 stage-slot 压到 11 个时间槽，但仍有填充和排空的“气泡”。本文脚本算出 4 段、8 个 microbatch 的理想利用率只有 72.73%，气泡占 27.27%。 再看 TP：若一个 MLP 的第一层权重为 $[d,4d]$，可沿输出维切成 $p$ 份，每卡算 $[d,4d/p]$；第二层 $[4d,d]$ 再沿输入维与中间激活匹配切分。这样两次 GEMM 都不需要在中间把完整 $4d$ 激活拼回每张卡，只在必要边界做集合通信。这正是 Megatron-LM 的 column-parallel 与 row-parallel 配对。 难点不在“切成几份”，而在切后保持数学等价，并让通信落在高带宽拓扑上。TP 频繁同步，通常限制在 NVLink/NVSwitch 节点内；PP 在相邻 stage 前向传激活、反向传对应梯度，更适合跨节点。若反过来布置，同一套卡数会被网络延迟拖垮。 02. 最小可用理解 三句话建立框架： TP 切一层的隐藏维、输出通道或 attention head，每层都会通信，换来单层权重、激活和计算分摊。 PP 把层分给不同 stage，再把 batch 切成 microbatch 以重叠各段计算；microbatch 越多，气泡越小，但调度和激活管理更复杂。 生产系统通常组成 DP×TP×PP 的三维网格：TP 放在最快链路内，PP 穿过较慢边界，剩余设备做 DP 与 ZeRO/FSDP。 如果只记一句：TP 是“同一个样本的一层由多卡一起算”，PP 是“同一个样本依次经过多卡上的不同层”。它们都属于模型并行，但通信模式完全不同。 设世界规模 $W$，数据并行度为 $d_p$、张量并行度为 $t_p$、流水线段数为 $p_p$，若没有其他并行轴： $$W=d_p\,t_p\,p_p$$ 这不是越均匀越好。TP 度受矩阵维度、head 数、KV head 数和高速互联限制；PP 度受层数、切分均衡与 microbatch 数限制；最后再用 DP 提升全局吞吐。 03. 数学推导 3.1 Column Parallel：沿输出维切 线性层输入 $X\in\mathbb{R}^{n\times d_{\mathrm{in}}}$，权重 $A\in\mathbb{R}^{d_{\mathrm{in}}\times d_{\mathrm{out}}}$，输出为： $$Y=XA$$ 把 $A$ 沿列切为 $p$ 份： $$A=[A_0,A_1,\ldots,A_{p-1}],\qquad A_i\in\mathbb{R}^{d_{\mathrm{in}}\times d_{\mathrm{out}}/p}$$ 第 $i$ 个 rank 计算： $$Y_i=XA_i$$ 完整结果只是拼接： $$Y=[Y_0,Y_1,\ldots,Y_{p-1}]$$ 前向不必通信，只要输入 $X$ 在 TP 组内复制。若下一层能直接消费分片 $Y_i$，连 all-gather 都可省。反向时，每个 rank 产生一份输入梯度贡献 $dX_i=dY_iA_i^\top$，完整 $dX$ 要求和： $$dX=\sum_{i=0}^{p-1}dX_i$$ 因此 column-parallel 的关键通信落在反向输入梯度 all-reduce。 3.2 Row Parallel：沿输入维切 把输入和权重的匹配维切开： $$X=[X_0,X_1,\ldots,X_{p-1}],\qquad A=\begin{bmatrix}A_0\\A_1\\ \vdots\\A_{p-1}\end{bmatrix}$$ 每卡先算局部部分和 $Z_i=X_iA_i$，完整输出是： $$Y=XA=\sum_{i=0}^{p-1}X_iA_i=\sum_{i=0}^{p-1}Z_i$$ 所以 row-parallel 前向需要 reduce 或 all-reduce，反向则能从复制的 $dY$ 本地得到 $dX_i=dYA_i^\top$。Megatron 把 MLP 的第一层做 column parallel，激活函数逐元素地作用在本地 shard；第二层做 row parallel，最后只求和一次。attention 的 QKV 投影与输出投影也能形成类似配对。 这种设计的精髓不是“每个 Linear 都随便切”，而是让相邻层的分片布局衔接，从而避免在每个算子后都 all-gather。若一个自定义层偷偷需要完整 hidden，通信会突然增加。 3.3 TP 的通信账本 设传输张量有 $N$ 个元素、每元素 $b$ 字节，ring all-reduce 在理想带宽模型中每 rank 通信： $$V_{\mathrm{AR}}=2\frac{p-1}{p}Nb$$ $p$ 是 TP 组大小。Megatron 风格的一个 Transformer 层在前向与反向关键边界各发生集合通信；真实实现可用 reduce-scatter、all-gather 与 sequence parallel 改写，但总成本仍与 token 数、hidden size 和 TP 频率有关。TP 每层同步，所以延迟不能忽略：小矩阵、多层和大 TP 度会让消息切得太碎，算力利用率反而下降。 TP 对参数显存理想降为 $1/p$，但并非所有东西都跟着切。LayerNorm 参数很小，可能复制；embedding 与 LM head 有词表切分规则；激活是否分片取决于 sequence-parallel 配置；通信临时 buffer 还会增加峰值。因此应从真实 module state 和 profiler 测量，不可把全显存直接除以 TP 度。 3.4 PP 气泡从哪里来 设有 $p$ 个等耗时 stage，global batch 切成 $m$ 个 microbatch，每个 stage 处理一个 microbatch 的单向计算耗时为一格。第一个 microbatch 需要 $p$ 格到达尾部，之后每格完成一个。完成一个 forward sweep 共： $$T_{\mathrm{sweep}}=m+p-1$$ 其中每个 stage 真正工作 $m$ 格，所以理想效率与气泡比例： $$\eta_{\mathrm{pipe}}=\frac{m}{m+p-1},\qquad f_{\mathrm{bubble}}=\frac{p-1}{m+p-1}$$ 当 $p=4,m=8$，效率 $8/11=72.73\%$。要到 90% 以上，需要 $m\ge9(p-1)$。这也解释了为什么 stage 很多却只有几个 microbatch 时 PP 很差。 公式假设各段等耗时、通信完全隐藏、无数据依赖停顿。真实训练的气泡还包括 stage 不均衡、点对点传输、参数同步、数据加载与尾 batch。最后一段若有巨大词表投影，前面切得再平均也会被它卡住。 3.5 GPipe、1F1B 与交错调度 GPipe 先把所有 microbatch 前向跑完，再统一反向，调度简单但要保存更多未反向的激活。1F1B 在 warmup 后交替执行一次 forward 和一次 backward，使稳态内存更接近少量 microbatch；同步更新语义仍要求一个 global batch 的梯度累积完再 step。 交错 1F1B 让每个物理设备承载多个 virtual stage，把一个大气泡拆细。它可改善负载与气泡，却增加更多通信边界、依赖和调度复杂度。Pipeline 并行的优化对象不只是公式里的 bubble，还包括峰值激活、通信重叠和 kernel 连续性。 均衡同步流水线的理想利用率，按 m/(m+p−1) 计算。增加 microbatch 能摊薄填充与排空气泡；图中不含通信、stage 失衡与小 GEMM 效率变化。 04. 代码实现 tp_linear_sim.py 只用标准库，构造 $X:[2,4]$、$W_1:[4,6]$、$W_2:[6,3]$。先完整计算，再把第一层按列、第二层按行切到两个虚拟 rank： X shape=(2, 4), W1 shape=(4, 6), W2 shape=(6, 3), tp=2 column shards: W1=(4, 3) each, hidden=(2, 3) each row shards: W2=(3, 3) each, partial output=(2, 3) each full output=[[52, 42, 72], [132, 114, 180]] max hidden diff=0, max output diff=0 输出差为 0，验证“按列拼接、按行求和”与完整矩阵乘严格等价。真实浮点并行的归约顺序不同，通常会出现末位误差，不能要求 bitwise 一致，而应设置与 dtype 相称的容差。 pipeline_bubble.py 默认模拟 4 段、8 个 microbatch： stages=4, microbatches=8, ideal_slots_per_sweep=11 pipeline_efficiency=72.73%, bubble_fraction=27.27% forward schedule (slot -> active stage:microbatch) 00: S0:M0 01: S0:M1 S1:M0 02: S0:M2 S1:M1 S2:M0 03: S0:M3 S1:M2 S2:M1 S3:M0 04: S0:M4 S1:M3 S2:M2 S3:M1 05: S0:M5 S1:M4 S2:M3 S3:M2 06: S0:M6 S1:M5 S2:M4 S3:M3 07: S0:M7 S1:M6 S2:M5 S3:M4 08: S1:M7 S2:M6 S3:M5 09: S2:M7 S3:M6 10: S3:M7 最前面三格在填充，最后三格在排空，中间四个 stage 才全部忙碌。用参数改变 stages 与 microbatches，就能看到增加 microbatch 如何摊薄固定气泡。 05. 工业级实现对照 以 2026-09 为准，NVIDIA Megatron-LM 的 tensor_parallel/layers.py 仍提供 ColumnParallelLinear 与 RowParallelLinear。工业实现除切权重外，还管理参数初始化、bias、是否 gather 输出、异步梯度 all-reduce、梯度累积融合、sequence parallel、专家并行组和通信 buffer。最小脚本只证明线性代数，没有模拟 autograd 与 NCCL。 ColumnParallelLinear 的 gather_output 决定输出是否立刻汇总；RowParallelLinear 的 input_is_parallel 表示输入是否已经按最后一维切好。错误组合往往不是数值报错，而是多做一次 gather 或得到错位 shard。审计并行模型时，应给每个边界标清 global shape、local shape、分片轴、复制轴和预期 collective。 PP 代码位于 Megatron Core 的 pipeline_parallel 调度模块。配置中的 pipeline_model_parallel_size 切物理 stage，virtual_pipeline_model_parallel_size 启用交错。系统必须传递前向激活和反向梯度，处理 tied embedding、首尾 stage 特殊 loss、不同张量 shape、激活释放与通信重叠。 Megatron 2021 展示了 TP、PP 与 DP 的组合，并提出交错流水调度。工程上的常见布局是：同一节点内组成 TP 组，邻接节点组成 PP 链，跨副本组成 DP 组。原因是 TP 每层通信而 PP 只在 stage 边界传输；把频繁 collective 放在更快互联上更划算。 生产配置还要做 layer partition。按层数平均只在每层耗时相同才合理；MoE 层、cross-attention、视觉模块、embedding 与词表 head 的成本差异很大。应先 profile 单层前后向时间和激活尺寸，再按时间而不是按层数切 stage。 06. 代价与边界 TP 的主要代价是高频 collective。TP 度增加后，每卡 GEMM 变小，计算效率下降，通信延迟占比上升；跨节点 TP 尤其敏感。维度还必须能合理整除：attention head、KV head、MLP intermediate size、词表分片与低精度 tile 对齐都会形成约束。 PP 的主要代价是气泡、激活驻留与调度复杂度。microbatch 增多能降气泡，却让每次 GEMM 的 batch 更小，可能降低算力利用率；还会增加调度次数。梯度累积数也受 global batch、DP 度与 microbatch size 约束： $$B_{\mathrm{global}}=B_{\mathrm{micro}}\,m\,d_p$$ 为了把 $m$ 调大而偷偷改变 global batch，会连学习率与收敛语义一起改变。更稳妥的做法是减小 microbatch size、保持 global batch，再检查小 GEMM 是否仍高效。 PP stage 之间有顺序依赖，单个慢 stage 会拖住全链。设备故障、动态 shape、条件分支和 MoE 路由也比普通 DP 难处理。推理时 batch 与请求长度动态变化，训练得到的均衡切分未必仍均衡。 TP 与 ZeRO/FSDP 并非天然可任意叠加。两者可能切同一参数的不同轴，process group、参数初始化、checkpoint layout 和 optimizer state 都需协调。首先建立二维或三维 rank 映射，再启用框架明确支持的组合，避免自行套两层 wrapper。 什么时候不该用？模型单卡能放下且吞吐受数据不足或 CPU 限制时，TP/PP 只会增加同步；单层很小而层很多时，优先 PP 或 FSDP；单层巨大而层数不多时，TP 更直接；超长序列导致激活爆炸时，还要引入下一篇的序列/上下文并行。 6.1 从模型形状反推 TP 度 假设 decoder hidden size 为 $d$，MLP expansion 为 $4d$，attention head 数为 $H$，每头维度为 $d_h=d/H$。采用 $t_p$ 路 TP 时，至少希望 $H/t_p$ 与 $4d/t_p$ 为整数，且本地矩阵维度满足 Tensor Core tile 对齐。GQA 还要检查 KV head 数 $H_{KV}$；若 Q head 可整除而 KV head 不可整除，框架可能复制 K/V 或根本拒绝配置。 参数容量只是下界。以 MLP 两个权重为例，忽略 bias： $$P_{\mathrm{MLP}}=d(4d)+(4d)d=8d^2$$ 理想 TP 后每卡为 $8d^2/t_p$ 个参数。然而输入 $X$、残差、LayerNorm 与某些输出仍可能复制。是否启用 sequence parallel，会决定 $[B,L,d]$ 激活是每 TP rank 一份还是沿 token 切开。评估 TP 度时要分别列参数、可分片激活、复制激活和通信 buffer，不能只用总显存除 $t_p$。 性能上，本地 GEMM 的算术强度会随 $t_p$ 增大而下降。若 $d/t_p$ 太小，矩阵乘无法占满 SM，即使通信为零也会变慢。实际选型一般从“能放下的最小 TP 度”起步，再尝试少数相邻值；不是卡越多 TP 越大。 6.2 三维 rank 映射为什么决定速度 有 2 个节点、每节点 8 卡，总计 16 卡。假设 TP=8、PP=2、DP=1，合理映射是每个节点内部构成一个 TP 组，两个节点作为相邻 PP stage。这样每层 TP collective 走 NVLink/NVSwitch，只有 stage 边界激活跨节点。 若 rank 编号错误，使每个 TP 组横跨两节点，那么每个 Transformer 层的 all-reduce 都经过网络；PP 反而在节点内。数学结果完全正确，吞吐却可能大幅下降。这类问题从配置数字看不出来，必须导出每个 process group 的物理 GPU、主机名、PCIe/NVLink 路径和 NIC 亲和性。 当 DP>1 时，DP 组应从相同 TP/PP 坐标上取不同副本。例如把 rank 写成坐标 $(r_d,r_p,r_t)$，模型同一 shard 的梯度只在 $r_d$ 维同步；层内 tensor collective 只改变 $r_t$；pipeline 点对点只沿 $r_p$ 邻接。把坐标语义明确下来，checkpoint 分片和故障定位才不会依赖偶然 rank 编号。 多 NIC 节点还要关注通信并发。TP all-reduce、PP send/recv、DP reduce-scatter 可能同时争同一链路。单独 benchmark 每个 collective 很快，不代表组合后仍快；应在真实 schedule 上看每条 stream 和 NIC 的时间线。 6.3 Pipeline 内存不是简单除以段数 PP 将参数按层分段，参数内存近似降为 $1/p_p$，但激活取决于 schedule。GPipe 前向完所有 $m$ 个 microbatch 才反向，每 stage 可能保留 $O(m)$ 份边界与层内激活。1F1B warmup 后尽早反向，可显著减少同时存活的 microbatch 数；不同 stage 的 warmup 长度又不同，峰值不完全一致。 activation checkpoint 将层内保存换成反向重算，但 stage 边界张量通常仍要保留或重新通信。若 PP 与 ZeRO-3 叠加，重算还可能再次触发参数 all-gather；如果 prefetch 与 release 时机不匹配，理论显存节省会被临时完整参数覆盖。 一个更实用的峰值模型是： $$M_s=M_{\mathrm{params},s}+n_{\mathrm{live},s}M_{\mathrm{act/micro},s}+M_{\mathrm{comm},s}+M_{\mathrm{workspace},s}$$ $s$ 是 stage，$n_{\mathrm{live},s}$ 是该调度下同时存活的 microbatch 数。选切分点的目标应是最小化 $\max_s M_s$ 并平衡每段时间，而不是让每段层数相同。首段 embedding、尾段 vocab projection 和 loss 往往需要单独计量。 6.4 Schedule 的一致性与权重版本 同步 GPipe 或 1F1B 中，一个 global batch 的所有 microbatch 应使用同一版参数，梯度累积后再统一 optimizer step。因此它与非流水同步训练在数学上等价，只改变操作顺序。随机 dropout 若要严格比较，还需保证不同 schedule 消耗一致的随机数流。 异步 pipeline 为减少 flush 可能允许某些 microbatch 使用旧权重，产生 weight staleness。PipeDream 的权重暂存与调度就是为管理这个问题。吞吐更高不等于优化轨迹相同；对需要可复现或大规模预训练的任务，同步 schedule 通常更容易验收。 梯度累积的 loss normalization 也常出错。若每个 microbatch loss 已取 mean，再把 $m$ 份梯度直接相加，最终梯度比全局 mean 大 $m$ 倍；框架可能在 loss、backward 或 optimizer wrapper 某处除 $m$。迁移实现时必须确认缩放发生在哪里，特别是最后一个不完整 microbatch。 6.5 如何系统调优而不是枚举所有组合 第一步做容量约束：根据参数、优化器、激活和临时 buffer，排除会 OOM 的 TP/PP。第二步做整除约束：检查层数、head、KV head、MLP width、词表与 microbatch。第三步按拓扑放组：TP 留在最快域，PP 穿过节点，DP 使用余下副本。第四步才短跑候选方案。 短跑至少记录每个 stage 的 forward/backward 时间、TP collective、PP send/recv、DP collective、bubble、MFU 和峰值显存。若 stage 时间方差大，先重切层；若 collective 裸露，检查异步与 bucket；若 GEMM 利用率低，减少 TP 或增大 microbatch；若 bubble 大，增加 microbatch 或 virtual stage。 还要避免只优化稳态。训练中 checkpoint、评估、数据切换、动态 loss scale 与长短样本混合会改变节拍。视频模型的序列长度可能随分辨率和帧数变化，固定层切分在不同 batch 上会失衡。可以按长度分桶、限制每批 token 数，或使用能处理动态 shape 的 schedule，但都要重新验证 global batch 语义。 6.6 Checkpoint 与并行度变更 TP checkpoint 的一个权重可能沿行或列分片，PP checkpoint 又只在拥有该层的 stage 上出现。保存时应记录全局 shape、分片轴、offset、replica group 与 tied-weight 关系。仅按 rank 存文件却没有布局元数据，会显著增加从 TP=8 改成 TP=4 的恢复难度；若完整掌握原切分约定仍可转换，但必须验证。 成熟 distributed checkpoint 会把逻辑参数名与物理 shard 解耦，加载时重新规划切片。验证转换不能只看“文件读完”，应抽样 all-gather 后与原始全参数比较，并跑一个确定性 forward。optimizer state 也必须按相同参数布局重分片，特别是 Adam 的 m、v 与 FP32 主权重。 PP 的 tied embedding 是典型边界：输入 embedding 在首段，输出投影在尾段，但两者可能共享参数。系统要么在两个 stage 间同步梯度，要么采用明确的复制与更新协议。忽略它会让模型能跑、loss 也下降，却不再是原架构。 最后为每个组合保留机器可读配置：world size、各轴度数、rank 坐标映射、global/micro batch、累积数、schedule、virtual stage、precision 与 checkpoint schema。没有这份清单，性能回归时很难判断是代码变化还是并行布局变化。 6.7 一个具体的 64 卡选型例子 假设 8 个节点、每节点 8 卡，模型单层在 4 卡上能放下，96 层在单节点放不下，目标 global batch 又允许 4 个数据副本。可先试 TP=4、PP=4、DP=4，乘积正好 64。每个节点容纳两个 TP 组，相邻两个节点组成一条四段 pipeline；相同 TP/PP 坐标跨四个副本形成 DP 组。 为什么不直接 TP=8、PP=2、DP=4？它减少 PP 气泡，但本地 GEMM 变小、每层 collective 参与卡数翻倍。为什么不 TP=2、PP=8、DP=4？单层可能放不下，而且 PP 段更多，要求更多 microbatch 才能摊薄气泡。三个方案都满足乘积约束，只有容量、profile 和拓扑能决定赢家。 设 PP=4、microbatch 数 $m=16$，理想气泡为 $3/19=15.79\%$；若为配合 DP=4，global batch 满足 $B_{\mathrm{global}}=B_{\mathrm{micro}}\times16\times4$。当目标 global batch 固定时，microbatch size 可能被压得太小。此时交错 PP 可在不继续增加 $m$ 的情况下缩小气泡，但要多传 stage 边界。 选定后再检查每卡峰值。若尾段词表 head 使 stage 3 明显更慢，可把少量 Transformer 层从尾段移到前段；若 stage 0 embedding 占显存而计算很少，切分目标应同时满足容量与时间，而不是追求参数量绝对相等。 6.8 故障现象与定位路径 若所有 GPU 利用率呈周期性锯齿且空白集中在迭代首尾，优先看 PP bubble；若每层 GEMM 后都有长 NCCL 条带，优先看 TP 通信或 rank 跨节点；若只有某一个 stage 长期满载、其他 stage 等待，是层切分失衡；若迭代末尾集中等待，是 DP 梯度同步没有充分 overlap。 出现 hang 时，先核对各 rank collective 调用序列。条件分支导致某些 rank 少调用一次 all-reduce，或 PP 两端 send/recv shape 不一致，都会永久等待。为通信操作记录 group、sequence number、peer、shape 与 dtype，比只看 Python stack 更有用。设置超时只能让错误更快暴露，不能修复顺序。 数值不一致时，从一个微型模型开始，关闭 dropout 和 fused kernel，分别测试 TP、PP，再测试组合。TP 常见错误是切分轴、bias 重复相加、输出误 gather；PP 常见错误是 loss 缩放、跨段激活 requires-grad、共享参数同步。一次只启用一根并行轴，能把搜索空间从三维降为一维。 OOM 若只发生在第一步，往往是 optimizer state 首次建立或通信 bucket 惰性分配；若在若干步后发生，可能是 graph retention、动态 shape 缓存或碎片；若只在某个 PP stage，先看该段真实 activation 与临时 buffer。不要看到 OOM 就统一减 microbatch，这可能掩盖泄漏却降低所有卡效率。 6.9 上线验收清单 发布配置前确认：世界规模等于各并行轴乘积；每个 rank 坐标唯一；TP group 位于预期高速域；head、KV head、MLP 与词表可整除；各 stage 时间和峰值接近；global batch 算术正确；loss normalization 不依赖 microbatch 数；tied weight 有同步协议；checkpoint 能跨目标并行度恢复。 性能门禁同时保存 tokens/s、MFU、bubble、每类 collective 的裸露时间、峰值显存和最慢 stage。只保存总迭代时间无法解释回归。集群拓扑或通信库升级后，即使模型代码未变，也要重跑基准，因为并行策略本质上是对物理系统的映射。 最后预留降级方案：某节点高速链路异常时能否减少 TP、增加 PP；某个 checkpoint 是否可重分片；global batch 是否仍保持；恢复后数值是否连续。真正稳健的并行配置，不只是峰值最快，也要能在硬件变化和断点恢复时保持语义清楚。 6.10 理论通信量为什么不等于训练时间 同样传 1 GiB，一个大 all-reduce 与数百个小 collective 的耗时不同。常用模型 $T=\alpha n+\beta V$ 中，$n$ 是消息次数，$V$ 是字节量，$\alpha$ 表示每次启动延迟，$\beta$ 表示每字节时间。TP 在每层频繁通信，尤其受 $\alpha$ 影响；PP 消息次数少，但单次边界激活可能很大。 通信是否裸露还取决于依赖。梯度 all-reduce 可与更早层 backward 重叠，PP send 可与下一份本地计算重叠；位于关键路径上的 collective 即使字节少也会直接延长 step。性能报告应区分总 NCCL 时间与 exposed communication time。 bucket 太小会增加消息次数，太大又推迟通信启动，减少 overlap。最优 bucket 与层大小、网络延迟和反向节奏有关。框架默认值是通用折中，不一定适合视频 DiT 的大激活或 MoE 的不均匀参数。 最后还要观察尾延迟。一个 rank 因热降频、ECC、数据抖动或网络拥塞变慢，collective 会让整个组等待。平均 GPU 时间看似健康，最慢 rank 才决定训练。按 rank 记录分位数，并把异常节点与拓扑关联，是大规模并行调优的基本动作。 理论模型最终要由 trace 校准。为每次实验保存并行配置、网络拓扑、逐 stage 时间线与通信矩阵，下一次才能定位回归。若只留一行 tokens/s，既无法判断瓶颈是 GEMM、气泡还是慢 rank，也无法安全迁移到另一代 GPU。并行策略不是模型的附属启动参数，而是计算图和硬件共同组成的一部分，理应像模型结构一样接受版本管理和回归测试。 07. 经典论文脉络 GPipe（arXiv:1811.06965）：用 microbatch 流水化跨加速器的模型分段，并以同步 mini-batch 语义训练巨型网络。 PipeDream（arXiv:1806.03377）：探索 1F1B 与异步流水，揭示吞吐、权重版本和一致性之间的权衡。 Megatron-LM（arXiv:1909.08053）：提出 Transformer 内高效的 tensor model parallel 切法，用少量 collective 支撑数十亿参数。 Efficient Large-Scale Language Model Training（arXiv:2104.04473）：系统组合 TP、PP、DP，并以交错流水减少气泡，扩展到千卡与万亿参数。 论文演进说明：单一并行轴只能解决一种容量瓶颈，规模继续增长后，真正的问题变成如何把多个轴映射到硬件拓扑并共同调度。 08. 常见误解 误解一：TP 就是把每层平均切开。 切分轴决定通信。Megatron 的 column-row 配对是为了让中间分片直接衔接，不是机械地切一半权重。 误解二：PP 段数越多，显存越省且速度越快。 参数容量会下降，但固定气泡随 $p-1$ 增长；若 microbatch 不够，更多 stage 反而更闲。 误解三：microbatch 越多越好。 它降低理想气泡，却减小 GEMM、增加调度，并可能扩大激活队列。要联动测 MFU 与显存。 误解四：1F1B 是异步优化。 常用同步 1F1B 只是改变前后向顺序，仍在一个 global batch 梯度完成后统一更新；PipeDream 式异步才涉及权重陈旧。 误解五：TP 能把所有显存除以 TP 度。 小参数、复制激活、通信 buffer 与 runtime 不会理想均分。必须看逐项账本。 误解六：三维并行度乘起来等于卡数就配置正确。 整除只是必要条件。拓扑、矩阵 tile、head 数、stage 均衡、global batch 都可能让配置不可用。 09. 动手验证 先运行 tp_linear_sim.py，把 TP 从 2 改成能整除维度的 3，重新切 $W_1$ 的输出与 $W_2$ 的输入。预期完整输出仍一致。然后故意让第二层 shard 顺序交换，结果会静默错误；这说明 distributed tensor layout 必须携带明确语义。 运行 pipeline_bubble.py，固定 stages=8，分别取 microbatches=1、8、32、72。理论效率依次为 $1/8$、$8/15$、$32/39$、$72/79$。脚本当前只实现均衡 stage 的闭式公式；扩展成事件调度模拟后，再给每个 stage 设置不同时间，以最慢 stage 为节拍比较，观察均分层数为何不等于均分时间。 在真实集群做 TP=1/2/4/8 对照，固定 global batch 和模型，记录每层 GEMM 时间、collective 时间、MFU、峰值显存。预期早期 TP 能解除容量或提高总吞吐，超过某点后小 GEMM 和通信使收益反转。用拓扑工具确认 TP rank 是否真的落在 NVLink 域内。 对 PP 做相同实验：固定 PP 度，逐步增加 microbatch 数；再固定 microbatch，增加 virtual stage。记录 bubble、激活峰值与端到端 tokens/s。不要只看 framework 打印的理论 bubble，GPU timeline 中的空白才是真实气泡。 最后做一次等价性验收：单卡、TP、PP、TP+PP 使用相同初始化和样本，关闭 dropout，比较一个 step 的 loss 和若干参数梯度。低精度与归约顺序会造成小差异，但若误差随层爆炸，优先检查 shard 顺序、bias、loss 归一化与 tied weight。 10. 延伸阅读 读完这篇可以继续看： 数据并行与 ZeRO 显存切分：DP 沿 batch 复制计算、ZeRO 沿 DP 组切状态，是三维并行的第一根轴（已发布）。 序列并行与 Ring Attention：当瓶颈从模型宽度转向超长 token，沿 sequence 切激活与 attention（本文系列下一篇）。 混合精度与数值稳定性：低精度改变 TP/PP 通信字节与归约误差，配置并行前要先明确累加精度（本文系列上一篇）。 选型时先问容量瓶颈在哪里：单层太宽选 TP，层总数太深选 PP，副本吞吐选 DP，超长激活选 SP。随后才是把这些轴映射到实际互联。公式给出上限，profile 决定最终配置。 附录：完整代码 09 节用到的脚本全文如下（tp_linear_sim.py、pipeline_bubble.py、make_figures.py）。复制到本地存成同名文件，按各脚本开头的依赖说明准备环境后即可运行。 tp_linear_sim.py #!/usr/bin/env python3 """Verify column- and row-parallel linear algebra without any framework.""" def matmul(a, b): return [[sum(x * y for x, y in zip(row, col)) for col in zip(*b)] for row in a] def split_columns(a, parts): width = len(a[0]) // parts return [[row[i * width:(i + 1) * width] for row in a] for i in range(parts)] def split_rows(a, parts): height = len(a) // parts return [a[i * height:(i + 1) * height] for i in range(parts)] def add(a, b): return [[x + y for x, y in zip(ra, rb)] for ra, rb in zip(a, b)] x = [[1, 2, 3, 4], [5, 6, 7, 8]] # [tokens=2, hidden=4] w1 = [[1, 0, 2, 0, 3, 0], [0, 1, 0, 2, 0, 3], [1, 1, 1, 1, 1, 1], [2, 1, 0, 1, 2, 1]] # [4, 6] w2 = [[1, 0, 1], [0, 1, 1], [1, 1, 0], [2, 0, 1], [0, 2, 1], [1, 0, 2]] # [6, 3] full_hidden = matmul(x, w1) full_output = matmul(full_hidden, w2) # Column parallel W1: every rank computes a slice of output features. w1_shards = split_columns(w1, 2) hidden_shards = [matmul(x, shard) for shard in w1_shards] column_joined = [left + right for left, right in zip(*hidden_shards)] # Row parallel W2: input features and W2 rows use matching shards, then sum. w2_shards = split_rows(w2, 2) partials = [matmul(h, w) for h, w in zip(hidden_shards, w2_shards)] row_reduced = add(partials[0], partials[1]) max_hidden_diff = max(abs(a - b) for ra, rb in zip(full_hidden, column_joined) for a, b in zip(ra, rb)) max_output_diff = max(abs(a - b) for ra, rb in zip(full_output, row_reduced) for a, b in zip(ra, rb)) print("X shape=(2, 4), W1 shape=(4, 6), W2 shape=(6, 3), tp=2") print("column shards: W1=(4, 3) each, hidden=(2, 3) each") print("row shards: W2=(3, 3) each, partial output=(2, 3) each") print(f"full output={full_output}") print(f"max hidden diff={max_hidden_diff}, max output diff={max_output_diff}") pipeline_bubble.py #!/usr/bin/env python3 """Idealized GPipe bubble calculator; stdlib only.""" import argparse parser = argparse.ArgumentParser() parser.add_argument("--stages", type=int, default=4) parser.add_argument("--microbatches", type=int, default=8) args = parser.parse_args() if args.stages < 1 or args.microbatches < 1: parser.error("stages and microbatches must be positive") p, m = args.stages, args.microbatches slots = m + p - 1 efficiency = m / slots print(f"stages={p}, microbatches={m}, ideal_slots_per_sweep={slots}") print(f"pipeline_efficiency={efficiency:.2%}, bubble_fraction={1-efficiency:.2%}") print("forward schedule (slot -> active stage:microbatch)") for t in range(slots): active = [f"S{s}:M{t-s}" for s in range(p) if 0 <= t - s < m] print(f"{t:02d}: " + " ".join(active)) make_figures.py """Regenerate this article's deterministic teaching figure (numpy + matplotlib).""" from pathlib import Path import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt import numpy as np OUT=Path(__file__).resolve().parents[1]/"figures" OUT.mkdir(exist_ok=True) plt.rcParams.update({"font.sans-serif":["PingFang SC","Arial Unicode MS","DejaVu Sans"],"axes.unicode_minus":False}) m=np.arange(1,129) fig,ax=plt.subplots(figsize=(8,4.8)) for p in [2,4,8,16]:ax.plot(m,m/(m+p-1),label=f"{p} stages") ax.set(xlabel="Microbatches per global batch",ylabel="Ideal utilization",ylim=(0,1.02),title="Balanced synchronous pipeline: m / (m + p - 1)") ax.legend();ax.grid(alpha=.25);fig.tight_layout() fig.savefig(OUT/'pipeline_utilization.png',dpi=170) plt.close(fig) print(OUT/'pipeline_utilization.png') 更多 AIGC 论文解读，关注微信公众号「人工智能炼丹君」 每日更新 · 论文精选 · 深度解读 · 技术脉络 微信搜索 人工智能炼丹君 或扫描下方二维码关注

AIGC 基本功|混合精度与数值稳定性-AMP 混合精度与数值稳定性 所属方向：分布式训练 ｜ 难度：进阶 ｜ 前置知识：无 关键词：混合精度、FP16、BF16、FP8、loss scaling、溢出、数值稳定 01. 为什么需要它 同一份训练代码，把 BF16 改成 FP16，吞吐可能更高，也可能几十步后 loss 直接变成 NaN。另一个常见现场是：模型看起来正常收敛，但某些参数长期没有变化；检查梯度才发现，一些绝对值不超过 $2^{-25}\approx2.98\times10^{-8}$ 的梯度在 round-to-nearest-even 写回 FP16 时已经变成了 0。 这不是“半精度不准”一句话能解释的。训练里的风险至少有三种：数太大，超过格式的最大有限值而溢出成 inf；数太小，舍入后落到 0（或硬件将次正规数 flush-to-zero）；数虽然在范围内，却因为有效位太少，在加法里被大数吞掉。三者的修复手段不同：loss scaling 能救小梯度，却救不了激活溢出；换 BF16 能扩展动态范围，却不会自动改善尾数精度；把归约留在 FP32 能减小累加误差，却不代表所有输入和权重都要回到 FP32。 为什么还要承担这些麻烦？因为 Transformer 的大头通常是矩阵乘。现代加速器对低精度矩阵乘提供更高吞吐，权重和激活每元素从 4 字节降到 2 字节也能减小显存与带宽压力。Mixed Precision Training 给出的核心配方是：低精度做大部分前后向，保留 FP32 主权重进行更新，并对 loss 缩放以防小梯度消失；其模型显存可接近减半，同时保持精度。 “混合”的重点不是选一个统一 dtype，而是给不同数值角色分工。大 GEMM 适合低精度输入，softmax、归一化、loss 和长归约通常需要更大的范围或更高精度，优化器状态又有自己的要求。AMP 的价值正是把逐算子的 dtype 路由表和梯度缩放流程标准化。 一个数字能说明 FP16 与 BF16 的性格差异：FP16 最大有限值只有 65504，而 BF16 与 FP32 一样有 8 位指数，最大值约 $3.39\times10^{38}$。反过来，FP16 有 10 位小数尾数，1 附近的间隔约 $9.77\times10^{-4}$；BF16 只有 7 位尾数，间隔约 $7.81\times10^{-3}$。BF16 更不容易炸，FP16 在可表示范围内更细；“更稳定”与“更精确”不是同一维度。 02. 最小可用理解 三句话先建立框架： AMP 让适合 Tensor Core 的矩阵乘使用 FP16/BF16，让 softmax、归约等敏感算子保留 FP32，而不是粗暴地把整个模型全部转成 half。 FP16 训练通常用 loss scale $S$ 把反向梯度整体放大，更新前再除回去；发现 inf/NaN 时跳过这一步并缩小 $S$。 BF16 动态范围接近 FP32，通常不需要 GradScaler，但尾数更短；FP8 则必须进一步管理张量尺度、绝对最大值历史和累加精度。 如果只记一个检查顺序：先问异常发生在前向还是反向，再区分 overflow、underflow 与 rounding，最后才决定换 dtype、调 scale，还是把局部算子提升到 FP32。 AMP 中常见的四种角色也要分开：参数存储 dtype、算子输入 dtype、乘法累加 dtype、优化器状态 dtype。日志打印“BF16 training”并不能证明四者都是 BF16。很多 GEMM 是 BF16 输入、FP32 累加；Adam 的一阶矩和二阶矩仍为 FP32；某些框架还会保存 FP32 主权重。 03. 数学推导 3.1 浮点数到底牺牲了什么 对非零正规数，一个二进制浮点数可写成（次正规数没有隐含的前导 1）： $$x=(-1)^s(1.f)_2\,2^{e-\mathrm{bias}}$$ $s$ 是符号位，$e$ 是指数域，$f$ 是小数域。指数位数决定动态范围，尾数位数决定相邻可表示数的间距。FP16 是 1 位符号、5 位指数、10 位小数；BF16 是 1、8、7；FP32 是 1、8、23。于是： FP16：最小正规数 $2^{-14}=6.1035\times10^{-5}$，最小次正规数 $2^{-24}=5.9605\times10^{-8}$，最大有限值 65504； BF16：最小正规数 $2^{-126}\approx1.1755\times10^{-38}$，最大有限值约 $3.3895\times10^{38}$； FP32：动态范围与 BF16 同阶，但 23 位小数显著降低舍入误差。 机器 epsilon 是 1 与下一个可表示数的间隔；下面列的是 epsilon。在 round-to-nearest 模式下，通常定义的 unit roundoff 为这些值的一半： $$\epsilon_{\mathrm{FP16}}=2^{-10},\qquad \epsilon_{\mathrm{BF16}}=2^{-7},\qquad \epsilon_{\mathrm{FP32}}=2^{-23}$$ 因此 $1+10^{-4}$ 写入 FP16 或 BF16 都可能仍是 1。更危险的是参数更新：若权重 $w=1$，学习率乘梯度只有 $10^{-5}$，直接在低精度权重上做 $w-\eta g$，变化会被舍掉。FP32 主权重就是在高精度副本上累计细小更新，再把结果舍入给低精度前后向。 在 round-to-nearest-even 且保留次正规数时，绝对值不超过 $2^{-25}$ 的 FP16 输入舍入为零；介于这个阈值和 $2^{-24}$ 之间的数可能舍入为最小次正规数，而非一律为零。支持 flush-to-zero 的执行路径还需另查硬件与算子规则。 3.2 为什么 loss scaling 不改梯度 设损失为 $L(\theta)$，参数为 $\theta$，真实梯度为 $g=\nabla_\theta L$。把 loss 乘常数 $S$ 后反向： $$g_s=\nabla_\theta(SL)=S\nabla_\theta L=Sg$$ 只要在优化器读取梯度前除以 $S$，就恢复原梯度： $$g=\frac{g_s}{S}$$ 关键是量化顺序。若 $g=10^{-8}$，先写入 FP16 会变 0，之后再乘任何数都救不回来；若先由链式法则得到 $Sg=1.024\times10^{-5}$，它能被 FP16 表示，写回后再用 FP32 除以 1024，就能保留接近 $10^{-8}$ 的非零值。 静态 scale 要人工选 $S$。太小救不了下溢，太大又会让大梯度超过 65504。动态 GradScaler 维护随训练变化的 $S_t$： $$S_{t+1}=\begin{cases}\beta S_t,&g_s\text{ 含 inf/NaN}\\ \gamma S_t,&\text{连续 }K\text{ 步有限}\\ S_t,&\text{其他情况}\end{cases}$$ $\beta<1$ 是回退因子，$\gamma>1$ 是增长因子，$K$ 是增长间隔。出现非有限梯度时必须跳过 optimizer step，否则坏更新会污染权重和 Adam 状态。PyTorch 文档还提醒：scale 不保证始终大于 1，BF16 预训练模型强转 FP16 时可能因数值过大而一路回退。 3.3 为什么 softmax 和归约容易出事 直接计算 softmax： $$p_i=\frac{e^{x_i}}{\sum_j e^{x_j}}$$ 若 $x_i$ 很大，指数会溢出。稳定写法先减最大值 $m=\max_jx_j$： $$p_i=\frac{e^{x_i-m}}{\sum_j e^{x_j-m}}$$ 这样最大指数为 1，不改变结果，却压下溢出风险。LayerNorm 方差也不宜用 $E[x^2]-E[x]^2$ 的低精度朴素形式：两个接近的大数相减会灾难性消减。长向量求和误差还会随项数累积，因此 AMP 通常让归约在 FP32 完成。 还要区分“算子输出 dtype”和“内部累加 dtype”。矩阵乘输入、输出可以是 BF16，但硬件乘积常进入 FP32 accumulator，最终再舍入回 BF16。若自定义 kernel 把 accumulator 也降成 16 位，它不再与常规 AMP 拥有同样数值性质。 3.4 FP8 为什么不只是再少 8 位 FP8 常见 E4M3 与 E5M2。FP8 Formats for Deep Learning 用 E4M3 提供较多有效位，用 E5M2 提供较大范围。原始张量很难恰好落进狭窄范围，通常给每个张量或块维护尺度 $a$： $$q=Q_{\mathrm{FP8}}(x/a),\qquad x_{\mathrm{deq}}=a\,q$$ 尺度可由绝对最大值 amax 与格式上限估计。若每步都同步全局 amax，又会引入开销；工业实现会用历史窗口、延迟缩放或块缩放。FP8 的正确性同时依赖格式、粒度、尺度更新、异常值分布和高精度累加，不能机械替换 dtype 字符串。 直接量化的相对误差。小数值端 FP16 可能舍入为零（相对误差 1），超过最大有限值后溢出；BF16 的范围更宽，但正规数的尾数精度更低。精确可表示点的误差为零，图中作绘图下限处理。 04. 代码实现 两个脚本都只依赖 Python 标准库。float_formats.py 用 IEEE 二进制打包模拟 FP16，并先按舍入位加偏置，再截去 FP32 低 16 位，实现 BF16 round-to-nearest-even。真实输出： format exp frac min_normal max_finite epsilon_at_1 FP16 5 10 6.1035e-05 6.5504e+04 9.7656e-04 BF16 8 7 1.1755e-38 3.3895e+38 7.8125e-03 FP32 8 23 1.1755e-38 3.4028e+38 1.1921e-07 value -> FP16 | BF16 1.0001 -> 1 | 1 1e-05 -> 1.00136e-05 | 1.00136e-05 100000 -> inf | 99840 100000 在 FP16 变成 inf，在 BF16 仍有限；1.0001 在两种 16 位格式里都舍入成 1。动态范围与精度的区别由此可见。 loss_scaling_sim.py 把五个小梯度直接量化，再先乘 1024、量化、最后除回去： shape=(5,), loss_scale=1024 gradient direct_fp16 scaled_fp16/unscaled 1.000e-08 0.000e+00 1.001e-08 3.000e-08 5.960e-08 2.998e-08 1.000e-07 1.192e-07 1.000e-07 1.000e-06 1.013e-06 1.000e-06 1.000e-05 1.001e-05 9.999e-06 nonzero: direct=4/5, scaled=5/5 $10^{-8}$ 直接写 FP16 已为 0，缩放路径却保留为 $1.001\times10^{-8}$。这不是凭空增加精度：量化误差仍在，只是把数搬进可表示区间。 真实 PyTorch CUDA 训练的核心顺序： scaler = torch.amp.GradScaler("cuda") for inputs, targets in loader: optimizer.zero_grad(set_to_none=True) with torch.autocast("cuda", dtype=torch.float16): loss = loss_fn(model(inputs), targets) scaler.scale(loss).backward() scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) scaler.step(optimizer) scaler.update() 梯度裁剪必须在 unscale 之后，若对放大后的梯度仍用阈值 C 裁剪，反缩放后的有效阈值会变成 C/S。若用 BF16，通常保留 autocast 而不用 GradScaler；是否省略仍应由实际梯度分布验证。 05. 工业级实现对照 以 2026-09 的实现为准，PyTorch 的 torch/amp/grad_scaler.py 中 GradScaler.scale 把 loss 乘当前尺度；unscale 按 device 与 dtype 分组检查并反缩放梯度；step 只在没有 inf/NaN 时调用优化器；update 再根据各 optimizer 的 found-inf 状态调整尺度。它还处理稀疏梯度、多设备、多优化器、checkpoint state dict 和惰性初始化，这些都是最小模拟没有覆盖的工程边界。 官方 AMP 文档 当前推荐统一使用 torch.autocast 与 torch.amp.GradScaler；旧的 torch.cuda.amp 接口已标记弃用。autocast 只包前向和 loss，backward 放在上下文外。不要启用 autocast 后再全局 model.half，否则会绕开逐算子的安全策略。 算子策略不是“白名单里全降精度”。线性层、卷积通常进入 lower-precision；softmax、部分 loss 和归约倾向 FP32；多输入算子可能提升到最宽输入类型。显式 dtype、原地算子或 out 版本可能不参与 autocast，自定义代码要检查实际 dtype 流。 FP8 方面，NVIDIA Transformer Engine 提供 E4M3/E5M2 recipe、amax 历史与 delayed scaling，并在支持硬件上让 Transformer 层进入 FP8 路径。它仍以 BF16/FP16 保存部分张量并用高精度累加；分布式训练还可能跨 rank 归约 amax。“开启 FP8”是引入量化运行时，不是把参数永久存成单一 8 位格式。 生产监控至少记录当前 scale、增长和回退次数、跳过的 step 数、梯度范数、参数与激活 amax、NaN 首次出现的层、关键算子 dtype。只看总 loss 会把静默下溢藏很久。 06. 代价与边界 混合精度通常节省权重、梯度和激活带宽，但不保证总显存恰好减半。Adam 的 FP32 主权重与两个 moment 仍可能占 12 字节/参数；部分算子保存 FP32 中间量；通信 buffer、cast buffer、workspace 与碎片也增加峰值。 吞吐提升也有条件。矩阵尺寸太小、CPU 或数据加载受限、频繁转换、未使用低精度加速单元，都会让 AMP 几乎不加速。小 batch 下，kernel launch 与 cast 成本甚至抵消收益。正确基准应同时报告 tokens/s、峰值显存、收敛曲线和最终指标。 以下情况尤其要谨慎： 模型源自 BF16 预训练且激活常超过 65504，FP16 loss scaling 只能处理梯度，救不了前向溢出； 长归约、概率、指数、对数、方差或很小正则项的自定义算子没有 autocast 策略； 梯度累积时各 microbatch 使用不同 scale，或反缩放前裁剪梯度； 多 optimizer 共用计算图，却在所有梯度就绪前更新 scale； FP8 中异常值主导 amax，张量级缩放让普通值量化过粗，需要块缩放或高精度旁路。 调试时先用 FP32 建立可复现基线，再开 BF16，然后才是 FP16/FP8。每步只改一个变量，并比较前若干 step 的 loss、梯度范数和权重更新。数值问题最怕同时调整学习率、并行度、batch 与 dtype。 6.1 一次更新里究竟有哪些精度 以 AdamW 为例，一次更新至少涉及前向激活、反向梯度、参数、副本和两个动量。低精度 GEMM 只覆盖其中一部分。设低精度参数为 $\theta_{16}$，FP32 主参数为 $\theta_{32}$，反缩放后的梯度为 $g_{32}$，则更新链条可以写成： $$\theta_{32}^{t+1}=\operatorname{AdamW}(\theta_{32}^{t},g_{32}^{t},m_t,v_t),\qquad \theta_{16}^{t+1}=Q_{16}(\theta_{32}^{t+1})$$ $m_t,v_t$ 是 FP32 一阶、二阶矩，$Q_{16}$ 表示舍入到前后向 dtype。若框架没有主参数副本，而直接更新 BF16 参数，许多小于当前参数 ULP 的更新会消失。是否保留主权重是优化器实现细节，不能仅从模型参数 dtype 推断。 梯度累积又多一层顺序。假设一个 optimizer step 包含 $A$ 个 microbatch，它们必须使用同一 scale $S$，先把已缩放梯度累加： $$g_s=\sum_{a=1}^{A}Sg_a=S\sum_{a=1}^{A}g_a$$ 等所有 microbatch 完成后只反缩放一次。若中途 update scale，累加 buffer 中不同 microbatch 带有不同倍数，最后无法用一个 $S$ 恢复。若某个 microbatch 产生 inf，这一整个 optimizer step 都应跳过，而不是只丢掉坏 microbatch，否则有效 batch 和采样权重已变化。 6.2 分布式训练里的非有限值共识 数据并行中，每个 rank 只看到本地样本。rank 3 出现 inf 而其他 rank 有限时，所有副本仍必须对“是否更新”达成一致；否则 rank 3 跳步、其余 rank 更新，参数立即分叉。工业 GradScaler 会把 found-inf 状态跨相关设备和优化器汇总，分布式封装还需确保所有 DP rank 做相同决策。 梯度裁剪也有全局语义。ZeRO/FSDP 把梯度分片后，每卡只能算局部平方和： $$\|g\|_2=\sqrt{\sum_{r=0}^{P-1}\sum_{i\in\mathcal{S}_r}g_i^2}$$ 必须 all-reduce 局部平方和才能得到全局范数。正确顺序通常是：完成梯度同步，反缩放，检查非有限值，计算全局 norm，裁剪，执行 optimizer step，最后更新 scale。任一步调换都可能让“裁剪阈值 1.0”失去原含义。 混合精度还会改变 collective 的数值误差。BF16 梯度 all-reduce 比 FP32 少一半字节，但不同归约树改变加法顺序；卡数越多、梯度尺度跨度越大，末位差异越明显。若训练对归约误差敏感，可使用 FP32 梯度通信或分块高精度累加，但要接受带宽代价。排查时应分开比较“本地梯度生成精度”与“跨 rank 归约精度”。 6.3 一张可执行的 dtype 选型表 选 FP16 还是 BF16，不应只看硬件宣传峰值。可按四步判断： 第一，确认硬件原生路径。某些设备对 BF16 与 FP16 吞吐相同，某些旧设备没有 BF16 Tensor Core；软件模拟 BF16 可能更慢。第二，检查模型来源。BF16 预训练 checkpoint 的激活范围未必适合 FP16，直接转换容易前向溢出。第三，检查数值敏感区域。softmax、归一化、概率 loss、长归约优先保留 FP32。第四，跑短程 A/B，比较吞吐、峰值、跳步率和收敛，而非只确认“能启动”。 可以把训练方案分成四档： FP32 基线：最慢但诊断最清晰，用于建立 loss 与梯度参照； BF16 AMP：大范围、无需常规 loss scaling，是现代大模型训练的常见起点； FP16 AMP：尾数稍细，但指数窄，必须认真监控动态 scaling 与前向溢出； FP8 混合精度：只把适合的 GEMM 张量降到 FP8，保留 16/32 位旁路，并维护尺度元数据。 所谓“纯 BF16”或“纯 FP8”常是营销简写。softmax、norm statistic、optimizer state、某些 residual 累加和 master gradient 仍可能更高精度。真正有意义的配置描述，应列出 input、weight、output、accumulator 和 optimizer 五个维度。 6.4 从第一个 NaN 反推根因 排查顺序应尽量靠近异常源。先在每个模块前后记录有限值比例、绝对最大值、绝对非零最小值与 dtype；找到第一个从有限变非有限的边界。若 attention score 在 softmax 前已经 inf，检查 Q/K 范数、缩放因子与位置编码；若 loss 有限而 backward 首先 inf，检查导数奇点、自定义 backward 和 loss scale；若梯度有限但 step 后参数异常，检查 optimizer state、权重衰减和反缩放顺序。 静默下溢比 NaN 更难发现。建议记录零梯度比例，并按参数组看更新比率： $$r_{\mathrm{update}}=\frac{\|\Delta\theta\|_2}{\|\theta\|_2+\varepsilon}$$ 若某些层长期为 0，而 FP32 基线非零，说明低精度存储或 scale 窗口吞掉了更新。若所有层的比例突然增大，则更像学习率、loss normalization 或跳步恢复后的 scale 问题。 不要看到 NaN 就立即降低学习率。学习率过大确实可能导致发散，但 dtype overflow、错误 mask 产生全负无穷、空 batch 的除零、坏数据和通信错误都可能表现为同一个 NaN。先定位首次非有限张量，才能选择对应修复。 6.5 Checkpoint、编译与自定义算子的坑 恢复训练时必须保存 GradScaler 状态。若只恢复模型和 optimizer，却把 scale 重置为很大初值，前几步可能反复溢出并被跳过；重置太小则产生额外下溢。保存点最好位于完整 optimizer step 之后，避免记录到一半累积的混合状态。 activation checkpoint 会在 backward 重跑前向。重算必须进入与原前向一致的 autocast 上下文，否则保存路径是 BF16、重算路径变 FP32 或 FP16，梯度对不上。torch.compile、CUDA Graph 和 fused optimizer 还可能改变 autocast 边界或引入 CPU-GPU 同步，升级版本后要重新做数值与性能回归。 自定义 autograd Function 不能假定输入永远 FP32。前向若内部要求 FP32，应显式关闭 autocast 并转换输入；反向要返回与调用约定兼容的梯度。写自定义 CUDA kernel 时，明确 accumulator 类型、饱和或 inf 行为、次正规数处理以及随机舍入。一个算子“支持 half”只代表能执行，不代表适合训练。 最后，评估数值稳定不能只跑几十步。动态 scale 的增长周期可能是数千步，某些罕见 batch 才触发极端值。至少保留一次覆盖学习率峰值、warmup 结束和验证阶段的长程对照，并把跳步数作为训练产物记录；否则断点续训后出现指标差异，很难追溯是数据还是精度状态造成。 6.6 性能基准该怎样设计 AMP 基准至少要有三次 warmup 与多次稳定迭代，计时前后做设备同步；否则异步 kernel 会让 CPU 提交时间冒充 GPU 执行时间。显存要同时报告 allocated 与 reserved 峰值，前者是活跃张量，后者包含分配器缓存。对比方案必须使用同一 batch、相同梯度累积与相同 checkpoint 策略。 吞吐最好报告有效 token/s，而不是 batch/s。变长序列下，一个 batch 的 padding 比例不同，batch/s 会误导。质量侧至少比较训练 loss、验证指标、梯度范数和被跳过 step 数。若 AMP 每秒更快却需要更多 step 才到相同验证指标，最终 time-to-quality 未必更优。 对 FP8 还应增加量化覆盖率：多少 GEMM 真正走 FP8，多少因 shape、算子或 recipe 回退到 BF16；记录每层 amax、饱和比例和 scale 更新。仅看配置显示 FP8 无法证明加速路径被命中。kernel trace 能确认实际指令与 cast 开销。 建立门禁时，不必要求低精度与 FP32 每步完全相同。可先规定前 100 步 loss 相对误差、梯度 cosine、最终指标容差与最大跳步率，再用多个随机种子评估。数值差异是浮点并行的正常现象，持续偏向、层级爆炸或指标显著退化才是故障信号。 6.7 一份上线前检查清单 上线前逐项回答：硬件是否原生支持目标 dtype；矩阵尺寸是否对齐加速 tile；敏感算子是否保持 FP32；FP16 是否启用动态 scaling；裁剪是否在 unscale 后；所有 DP rank 是否共享跳步决策；梯度累积期间 scale 是否不变；checkpoint 是否保存 scaler；自定义算子是否声明 autocast 与 accumulator；监控是否能定位第一处非有限值。 然后做三个故障注入。人为把 scale 调得极大，确认系统发现 inf、跳过更新并回退；给输入加入一个幅值异常样本，确认前向监控能定位层；从 checkpoint 恢复，确认 scale、optimizer 与随机数状态连续。没有做过故障注入的告警，往往只在真正长跑失败后才发现无效。 最后保留 FP32 或 BF16 安全开关。生产训练发生异常时，能在不改数据顺序和并行布局的条件下提升精度复现，定位效率远高于临时改一堆超参。安全路径不一定长期运行，但必须定期测试，避免代码演进后早已失效。 这些流程看似比设置一个 autocast 开关繁琐，却能把“偶尔 NaN”“换卡就掉点”“断点后不收敛”变成有指标、有复现实验、有回退方案的普通工程问题。 6.8 如何判断变化来自精度而不是随机性 一次训练 A 比 B 的 loss 高，并不能证明精度方案更差。数据顺序、dropout、并行归约顺序和非确定 kernel 都会制造波动。公平实验要固定数据索引与初始化，尽可能使用确定算法，并运行多个随机种子。先比较同一步同一层的输出、梯度与更新，再比较长程最终指标。 可用 FP32 输出 $y_{32}$ 作为局部参照，计算相对误差与余弦相似度： $$e_{\mathrm{rel}}=\frac{\|y_{\mathrm{low}}-y_{32}\|_2}{\|y_{32}\|_2+\varepsilon}$$ 单层误差略大未必影响训练，但若误差沿深度单调放大，往往说明 residual 累加、归一化或某个敏感算子精度不足。把统计按层绘制比只比较最终 logits 更容易定位。 还应区分可复现性与正确性。集合通信改变加法顺序后，bitwise 结果可能不同，但两条训练曲线仍落在相同统计分布。反之，两次运行逐位一致也可能稳定地实现了错误缩放。验收既要有小规模数学 oracle，也要有多种子任务指标。 混合精度上线后，持续监控跳步比例和 scale 分布。数据配方变化、序列变长、加入新 loss 或更换初始化，都可能改变数值范围；一次验证通过不意味着未来配置永久安全。把 precision 当成模型配置的一部分进行版本化，才能复现每次训练。 一个实用原则是把所有隐式转换显式化到观测层：训练启动时抽样打印关键模块的参数、输入、输出与归约 dtype，同时保存硬件、驱动、框架版本。低精度 kernel 与 autocast 策略会随版本更新，旧实验结论不能无条件外推。版本升级后的第一件事应是重跑短程数值基线与吞吐基线，而不是直接续跑昂贵训练。这样才能知道速度变化来自 kernel，精度变化来自策略，还是数据本身发生了变化。 还应把这些元数据写进 checkpoint 清单，使恢复任务能够拒绝不兼容的精度配置，而不是在数小时后以 NaN 形式暴露。可诊断性本身就是混合精度系统的一项能力。 07. 经典论文脉络 Mixed Precision Training（arXiv:1710.03740）：提出低精度前后向、FP32 主权重与 loss scaling，奠定现代 AMP。 A Study of BFLOAT16 for Deep Learning Training（arXiv:1905.12322）：说明 BF16 用 FP32 相同指数范围换取更短尾数，使训练更少依赖 loss scaling。 FP8 Formats for Deep Learning（arXiv:2209.05433）：提出 E4M3/E5M2 两种互补编码，并在大模型训练中验证 8 位浮点路径。 FP8-LM（arXiv:2310.18313）：处理 FP8 大语言模型训练的精度、通信与优化器环节，说明端到端 FP8 不只是替换 GEMM dtype。 共同主题是：硬件格式越窄，软件越需要知道张量的数值角色。格式提供可能性，尺度管理、累加方式和回退路径才决定能否稳定训练。 08. 常见误解 误解一：BF16 比 FP16 精度更高。 BF16 指数范围更大，更少 overflow；但尾数只有 7 位，1 附近分辨率更粗。应说通常更稳定，不是处处更精确。 误解二：loss scaling 能修复所有 NaN。 它主要防 FP16 小梯度下溢。前向激活溢出、除零、softmax 不稳定、学习率过大，都不会被它根治。 误解三：scale 越大越好，而且一定大于 1。 scale 太大会溢出。PyTorch 允许动态 scale 降到 1 以下；应观察 found-inf 和跳步频率。 误解四：用了 autocast 就无需关心 dtype。 自定义算子、显式 dtype、原地运算和策略表外 op 可能保持输入 dtype。关键归约与 loss 仍需审计。 误解五：AMP 会把训练显存直接砍半。 激活可能下降，FP32 优化器状态仍在。账本还包括梯度、通信 buffer、workspace 与碎片。 误解六：FP8 是 BF16 的无痛升级。 FP8 依赖尺度、amax 历史、格式选择和硬件 kernel；错误配置可能不 NaN，却悄悄降低收敛质量。 09. 动手验证 先运行文末脚本，修改 loss_scaling_sim.py 中 scale 为 1、128、1024、65536。预期 scale 变大时更多微小梯度被保留；继续增大并加入大梯度后会触发 FP16 inf。这正是动态 scale 在 underflow 与 overflow 之间寻找窗口的原因。 在支持 BF16/FP16 的 GPU 上做四组短跑：FP32、BF16 autocast、FP16 autocast 不带 scaler、FP16 autocast 带 scaler。固定随机种子和 batch，记录 200 步吞吐、峰值显存、loss、全局梯度范数、跳步数。预期 FP16 无 scaler 最容易出现零梯度；BF16 与 FP16 加 scaler 更接近 FP32，但结论取决于模型分布。 定位首个异常层时，给模块注册 hook，输出有限值比例、绝对最大值和 dtype。若前向先出现 inf，优先检查指数、归一化、attention score 与输入范围；若反向先异常，再检查 scale、梯度裁剪顺序和自定义 backward。 最后验证累加精度：构造一个大数和大量小数的向量，用 FP32、模拟 BF16 顺序累加、分块 FP32 累加分别求和。预期低精度顺序加法吞掉更多小项；这解释为什么 GEMM 输入低精度不等于 accumulator 也该低精度。 10. 延伸阅读 读完这篇可以继续看： 量化：从 INT8 到 FP4：区分训练混合精度与推理量化的目标、校准和误差模型（还没写）。 数据并行与 ZeRO 显存切分：混合精度改变每项状态字节数，ZeRO 决定状态在哪些 rank 上复制或分片（已发布）。 张量并行与流水线并行：低精度减少通信字节，但集合通信次数、拓扑与同步点仍由并行策略决定（本文系列下一篇）。 可靠的 AMP 心智模型不是“半精度开关”，而是一张数值预算表：每个张量需要多大范围、多少有效位，在何处累加、何时舍入、出现异常如何回退。把这张表画清楚，NaN 就从玄学变成可以定位的工程问题。 附录：完整代码 09 节用到的脚本全文如下（float_formats.py、loss_scaling_sim.py、make_figures.py）。复制到本地存成同名文件，按各脚本开头的依赖说明准备环境后即可运行。 float_formats.py #!/usr/bin/env python3 """Only the Python standard library is required.""" import math import struct def fp16(x: float) -> float: try: return struct.unpack("e", struct.pack("e", x))[0] except OverflowError: return math.copysign(math.inf, x) def bf16(x: float) -> float: """Round an IEEE FP32 value to BF16, ties-to-even, then widen for printing.""" bits = struct.unpack(">I", struct.pack(">f", x))[0] if (bits & 0x7F800000) != 0x7F800000: bits += 0x7FFF + ((bits >> 16) & 1) return struct.unpack(">f", struct.pack(">I", bits & 0xFFFF0000))[0] FORMATS = ( ("FP16", 5, 10, 2.0**-14, (2 - 2.0**-10) * 2.0**15, 2.0**-10), ("BF16", 8, 7, 2.0**-126, (2 - 2.0**-7) * 2.0**127, 2.0**-7), ("FP32", 8, 23, 2.0**-126, (2 - 2.0**-23) * 2.0**127, 2.0**-23), ) print("format exp frac min_normal max_finite epsilon_at_1") for name, exp, frac, low, high, eps in FORMATS: print(f"{name:5s} {exp:3d} {frac:4d} {low:.4e} {high:.4e} {eps:.4e}") values = (1.0001, 0.00001, 100000.0) print("\nvalue -> FP16 | BF16") for value in values: print(f"{value:g} -> {fp16(value):g} | {bf16(value):g}") loss_scaling_sim.py #!/usr/bin/env python3 """Show how FP16 loss scaling rescues tiny gradients; stdlib only.""" import math import struct def fp16(x: float) -> float: try: return struct.unpack("e", struct.pack("e", x))[0] except OverflowError: return math.copysign(math.inf, x) gradients = [1e-8, 3e-8, 1e-7, 1e-6, 1e-5] scale = 1024.0 direct = [fp16(g) for g in gradients] scaled_then_unscaled = [fp16(g * scale) / scale for g in gradients] print(f"shape=({len(gradients)},), loss_scale={scale:g}") print("gradient direct_fp16 scaled_fp16/unscaled") for g, raw, rescued in zip(gradients, direct, scaled_then_unscaled): print(f"{g:11.3e} {raw:11.3e} {rescued:11.3e}") print(f"nonzero: direct={sum(x != 0 for x in direct)}/{len(direct)}, scaled={sum(x != 0 for x in scaled_then_unscaled)}/{len(direct)}") make_figures.py """Regenerate this article's deterministic teaching figure (numpy + matplotlib).""" from pathlib import Path import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt import numpy as np OUT=Path(__file__).resolve().parents[1]/"figures" OUT.mkdir(exist_ok=True) plt.rcParams.update({"font.sans-serif":["PingFang SC","Arial Unicode MS","DejaVu Sans"],"axes.unicode_minus":False}) from float_formats import fp16,bf16 x=np.geomspace(1e-9,1e5,1500) y16=np.array([fp16(float(v)) for v in x]);yb=np.array([bf16(float(v)) for v in x]) fig,ax=plt.subplots(figsize=(8,4.8)) for name,y in [("FP16",y16),("BF16",yb)]: err=np.abs(y-x)/x;err[~np.isfinite(err)]=np.nan ax.loglog(x,np.maximum(err,1e-12),label=name,alpha=.75) ax.axvline(65504,color="grey",ls=":",label="FP16 max finite") ax.set(xlabel="Positive input value",ylabel="Relative rounding error",title="Quantization range and precision (nearest-even)",ylim=(1e-8,2)) ax.legend();ax.grid(alpha=.25);fig.tight_layout() fig.savefig(OUT/'precision_error.png',dpi=170) plt.close(fig) print(OUT/'precision_error.png') 更多 AIGC 论文解读，关注微信公众号「人工智能炼丹君」 每日更新 · 论文精选 · 深度解读 · 技术脉络 微信搜索 人工智能炼丹君 或扫描下方二维码关注

AIGC 基本功|数据并行与 ZeRO 显存切分-ZeRO 数据并行与 ZeRO 显存切分 所属方向：分布式训练 ｜ 难度：进阶 ｜ 前置知识：无 关键词：数据并行、DDP、ZeRO、优化器状态切分、FSDP、显存占用估算 01. 为什么需要它 先算一笔会让很多人意外的账：训练一个 75 亿参数模型，模型权重明明只有 15 GB，为什么放进 80 GB 显存的 GPU 还会爆？ 假设用混合精度 Adam 训练。每个参数除了 2 字节的 FP16/BF16 权重，还要留下 2 字节梯度、4 字节 FP32 主权重、4 字节一阶动量和 4 字节二阶动量。合计不是 2 字节，而是 16 字节/参数： $$7.5\times10^9\times16\ \mathrm{bytes}=120\ \mathrm{GB}$$ 这 120 GB 还没有算激活、临时通信缓冲区、CUDA context 和显存碎片。换句话说，模型甚至没开始处理一帧视频，单是“为了训练而保存的状态”就已经放不下。 自然的反应是：“那我上 64 张卡。”可如果只用普通数据并行，每张卡都会保留完整的参数、完整的梯度和完整的 Adam 状态。64 张卡只是把不同 mini-batch 同时算得更快，每张卡看到的仍然是同一份 120 GB。总显存从 80 GB 变成 5120 GB，但单卡瓶颈一点没松，模型照样启动不了。 这就是 ZeRO 要解决的矛盾：数据并行已经拥有整个集群的总显存，却因为每张卡都复制同样的训练状态，只能使用其中一张卡的容量。ZeRO 不改变模型的数学计算，而是按顺序切掉三种冗余副本： ZeRO-1 切优化器状态； ZeRO-2 再切梯度； ZeRO-3 连参数也切，只在某一层即将计算时临时拼回来。 把刚才的 75 亿参数模型放到 64 个数据并行 rank 上，三阶段的模型状态显存分别是 31.41 GB、16.64 GB 和 1.88 GB。最后一个数字正好是 $120/64$。这些数字不是宣传页上的模糊“最高节省多少”，而是本文后面会逐项推出来、再用代码复现的结果。 对视频生成尤其要补一句：视频 DiT 的 token 数很长，激活往往比模型状态还大。ZeRO 能拆掉模型状态冗余，但不会自动消灭激活。先分清是哪一类显存爆了，再决定用 ZeRO、激活重计算还是序列并行，比盲目把配置改成 stage 3 更重要。 02. 最小可用理解 三句话先建立框架： 数据并行让每个 rank 保存同一个模型、读取不同数据，反向后把局部梯度求平均，因而数学上等价于在合并后的大 batch 上做一次同步更新。 既然所有 rank 最终拿到相同梯度，它们各自保存一整套 Adam 状态并重复做同一份更新就是冗余；ZeRO 把状态分片，每个 rank 只负责其中 $1/N$。 切得越彻底，常驻显存越少，但参数就越需要“用前聚合、用后释放”，因此省下的显存会转化成通信、调度和峰值控制问题。 如果只记一件事：DDP 是复制计算，ZeRO 是切分状态；ZeRO 没有改变梯度本身。 这里的 rank 可以先理解为一个独立训练进程，通常一个 rank 独占一张 GPU。$N$ 表示数据并行进程数，world_size 就是 $N$。每个 rank 处理本地 batch，但每一步结束后所有 rank 必须得到一致参数，否则下一步就不再是同一个模型。 03. 数学推导 3.1 为什么平均梯度等价于一个大 batch 设全局 batch 有 $B$ 个样本，均匀切给 $N$ 个 rank，每个 rank 得到 $b=B/N$ 个样本。模型参数记作 $w$，第 $i$ 个样本的损失是 $\ell(x_i;w)$。 第 $r$ 个 rank 的局部平均损失为： $$L_r(w)=\frac{1}{b}\sum_{i\in\mathcal{B}_r}\ell(x_i;w)$$ 其中 $\mathcal{B}_r$ 是 rank $r$ 拿到的样本集合，$b$ 是本地 batch size。对参数求导，得到局部梯度： $$g_r=\nabla_w L_r(w)=\frac{1}{b}\sum_{i\in\mathcal{B}_r}\nabla_w\ell(x_i;w)$$ 同步数据并行对所有局部梯度取平均： $$g=\frac{1}{N}\sum_{r=1}^{N}g_r=\frac{1}{Nb}\sum_{r=1}^{N}\sum_{i\in\mathcal{B}_r}\nabla_w\ell(x_i;w)=\frac{1}{B}\sum_{i=1}^{B}\nabla_w\ell(x_i;w)$$ 右边正是把所有样本放在一张卡上算全局平均损失所得到的梯度。等价成立依赖三个前提：损失能按样本分解且样本计算不因分组改变（如本地 BatchNorm 统计会破坏这一条件）；各 rank 从同一个参数 $w$ 出发；样本权重一致；更新前完成同步。若最后一个 batch 在各 rank 上不等长，仍然机械地“先本地平均再按 rank 平均”，样本就会被赋予错误权重。这也是分布式数据加载器必须谨慎处理尾 batch 的原因。 DDP 的核心任务由此非常明确：它不替你切输入，而是在 autograd 产生梯度后，把各 rank 的 $g_r$ 聚合成相同的 $g$。 3.2 先把 16 字节拆明白 设模型有 $P$ 个可训练参数，采用低精度前后向和 FP32 Adam 更新。每个参数对应： 低精度参数：2 字节； 低精度梯度：2 字节； FP32 主参数：4 字节； Adam 一阶矩 $m$：4 字节； Adam 二阶矩 $v$：4 字节。 后面三项统称优化器状态，共 12 字节。普通 DDP 每个 rank 的模型状态显存是： $$M_{\mathrm{DDP}}=(2+2+12)P=16P$$ 注意这是一个状态下界模型，没有包含激活等剩余显存。它的价值不是预测 nvidia-smi 到个位数，而是先回答“哪一类状态值得切”。 实际套公式时，$P$ 应取需要优化器更新的参数量，不是配置文件里笼统写出的模型总参数量。冻结的视觉编码器通常仍要保存推理权重，却不产生梯度和 Adam 状态；共享 embedding 在参数统计中也只能算一次。优化器同样会改常数：SGD 无动量时几乎没有 $m,v$，8-bit Adam 会压缩这两份状态，某些纯 BF16 配置也不保留 FP32 主参数。所以正确动作不是背住“训练恒等于 16 字节/参数”，而是先列出当前训练栈真正持有的张量，再把每一项的元素数乘 dtype 字节数。本文使用 16 字节，是为了与 ZeRO 原论文的混合精度 Adam 假设严格对齐。 3.3 ZeRO 三阶段到底切了什么 ZeRO-1：只切优化器状态。 参数和梯度仍然每卡各有一份，12 字节的 FP32 主参数与 Adam 两个矩按 $N$ 份切开： $$M_{\mathrm{Z1}}=4P+\frac{12P}{N}$$ 当 $N$ 很大，第二项趋近于 0，但第一项中的 2 字节参数和 2 字节梯度仍然复制，所以 ZeRO-1 的极限是 $4P$，相对 $16P$ 最多节省 4 倍。 ZeRO-2：再切梯度。 常驻的完整副本只剩 2 字节低精度参数；梯度与优化器状态合计 14 字节一起分片： $$M_{\mathrm{Z2}}=2P+\frac{14P}{N}$$ 当 $N$ 很大，极限是 $2P$，相对普通 DDP 最多节省 8 倍。这里“切梯度”不是说某个 rank 永远见不到别人的梯度，而是 reduce-scatter 聚合后，每个 rank 只保留自己负责更新的那一片结果，其余梯度用完即释放。 ZeRO-3：参数也切。 参数、梯度和优化器状态全部只保留 $1/N$： $$M_{\mathrm{Z3}}=\frac{16P}{N}$$ 这时单卡模型状态随卡数线性下降，理论上终于可以使用整个数据并行组的聚合显存。代价是某层计算前必须 all-gather 出完整参数，计算后再 reshard。这里的“完整”通常只针对一个 FSDP unit 或若干层，而不是一次把全模型永远拼回显存；包装粒度如果选错，峰值仍然可能爆掉。 代入 $P=7.5\times10^9$、$N=64$： DDP：$16P=120.00$ GB； ZeRO-1：$4P+12P/64=31.41$ GB； ZeRO-2：$2P+14P/64=16.64$ GB； ZeRO-3：$16P/64=1.875$ GB。 这与 ZeRO 原论文表 1中的 120、31.4、16.6、1.88 GB 对上了。论文使用十进制 GB；如果监控工具显示 GiB，同一字节数会小约 7.4%，比较数字前先统一单位。 3.4 显存省了，通信为什么没有立刻爆炸 设梯度张量大小为 $G$ 字节。高带宽 ring all-reduce 通常拆成 reduce-scatter 和 all-gather。忽略延迟、拓扑和协议常数后，每个 rank 在两个阶段分别移动： $$V_{\mathrm{RS}}=\frac{N-1}{N}G,\qquad V_{\mathrm{AG}}=\frac{N-1}{N}G$$ 因此普通 DDP 的总通信量近似为： $$V_{\mathrm{DDP}}=2\frac{N-1}{N}G\approx2G$$ ZeRO-1/2 把原来的 all-reduce 改排成“梯度 reduce-scatter + 更新后参数 all-gather”。在本文参数与梯度均为两字节的假设下，两者大小同为 $G$，理想总字节量仍约为 $2G$。若用 FP32 梯度而参数为 BF16，两项大小不同，不能直接沿用该字节比。区别在于数据留在哪里、何时释放，而不是凭空少传了一半。 ZeRO-3 还要在前向和反向各聚合一次参数，再做梯度 reduce-scatter。按 ZeRO 论文用参数元素数 $P$ 计量，普通数据并行约移动 $2P$ 个元素，Stage 3 约移动 $3P$ 个元素，所以是基线的 1.5 倍。这个 1.5 倍假设参数/梯度通信 dtype 相同并采用所述聚合与释放调度，是理想带宽模型，不是训练时间必然乘 1.5：通信能否和计算重叠、跨没跨节点、消息是否太碎，都会改变真实结果。 3.5 别把模型状态公式当成整机显存公式 实际峰值更接近： $$M_{\mathrm{peak}}=M_{\mathrm{model\ states}}+M_{\mathrm{activations}}+M_{\mathrm{temporary}}+M_{\mathrm{runtime}}+M_{\mathrm{fragmentation}}$$ 第一项就是前面推导的参数、梯度和优化器状态。激活由 batch、序列长度、空间分辨率、层数和 checkpoint 策略决定；临时项包含 all-gather bucket、梯度 bucket 和算子工作区；runtime 包含 CUDA context、通信库等；碎片则取决于张量生命周期和分配器状态。 ZeRO 原论文给过一个很有提醒意义的例子：15 亿参数 GPT-2 的模型状态至少 24 GB；序列长度 1024、batch 32 时，激活约 60 GB，即使用激活重计算降到约 8 GB，临时 FP32 扁平缓冲还可能再占 6 GB。只看 $16P$ 判断“32 GB 正好能放下”会直接翻车。 模型状态随数据并行卡数的变化。固定 FP16 参数/梯度与 FP32 主权重及 Adam 状态，纵轴为十进制 GB；不含激活、临时聚合缓冲和运行时。 04. 代码实现 本节有两个完整脚本，均只依赖 Python 标准库。第一个复现论文显存表；第二个把 DDP all-reduce 与 ZeRO-2 的 reduce-scatter + all-gather 拆开，验证它们做出同一次 Adam 更新。完整代码在文末附录。 4.1 用公式生成显存账本 zero_memory_ledger.py 把每个阶段写成“每卡复制多少字节 + 分片多少字节”： STAGES = ( Stage("DDP", 16.0, 0.0, 1.0), Stage("ZeRO-1", 4.0, 12.0, 1.0), Stage("ZeRO-2", 2.0, 14.0, 1.0), Stage("ZeRO-3", 0.0, 16.0, 1.5), ) def bytes_per_parameter(self, world_size: int) -> float: return self.replicated_bytes + self.sharded_bytes / world_size replicated_bytes 是每个 rank 必须完整保留的部分，sharded_bytes 是可以除以 $N$ 的部分。默认参数就是论文的 7.5B/64 卡案例。实际运行： model=7.5B, world_size=64 assumption: fp16 params 2B + fp16 grads 2B + fp32 master/m/v 12B = 16 bytes/parameter stage bytes/param model-state GB vs DDP comm/step GB DDP 16.0000 120.00 1.00x 30.00 ZeRO-1 4.1875 31.41 3.82x 30.00 ZeRO-2 2.2188 16.64 7.21x 30.00 ZeRO-3 0.2500 1.88 64.00x 45.00 note: activation, temporary buffers and fragmentation are excluded 为什么 64 卡的 ZeRO-1 只有 3.82 倍而不是宣传里的 4 倍？因为 4 倍是 $N\to\infty$ 的上限，有限卡数下 $12P/N$ 还没有消失。ZeRO-2 的 7.21 倍同理。ZeRO-3 没有复制项，所以恰好获得 64 倍。 再换成 1.5B/8 卡： model=1.5B, world_size=8 assumption: fp16 params 2B + fp16 grads 2B + fp32 master/m/v 12B = 16 bytes/parameter stage bytes/param model-state GB vs DDP comm/step GB DDP 16.0000 24.00 1.00x 6.00 ZeRO-1 5.5000 8.25 2.91x 6.00 ZeRO-2 3.7500 5.62 4.27x 6.00 ZeRO-3 2.0000 3.00 8.00x 9.00 note: activation, temporary buffers and fragmentation are excluded 这组数字揭示一个选型习惯：模型状态只差几 GB 时，ZeRO-1/2 往往已经够用；不必为了追求最低常驻显存直接承担 Stage 3 的参数聚合。 4.2 ZeRO 为什么没有改掉优化结果 zero_update_simulator.py 构造 4 个 rank、8 个参数。每个 rank 看见不同数据，因此产生不同局部梯度。DDP 路径先得到完整平均梯度，再完整执行 Adam；ZeRO-2 路径只把平均梯度对应的两元素分片交给各 rank，每个 rank 维护自己的 $m,v$，更新后再把四个参数片聚合起来。 核心区别只有分片发生的位置： reduced_full_grad = average_columns(local_grads) ddp_params, ddp_m, ddp_v = adam_first_step(PARAMS, reduced_full_grad) param_shards = shard(PARAMS, WORLD_SIZE) grad_shards = shard(reduced_full_grad, WORLD_SIZE) updated_shards = [] for params_for_rank, grads_for_rank in zip(param_shards, grad_shards): updated, local_m, local_v = adam_first_step(params_for_rank, grads_for_rank) updated_shards.append(updated) zero_params = [value for part in updated_shards for value in part] 真实运行输出： world_size=4, parameter_count=8, shard_width=2 local gradient matrix shape=(4, 8) reduce-scatter result shape=(4, 2) averaged gradient: [0.02, 0.045, 0.07, 0.095, 0.12, 0.145, 0.17, 0.195] DDP updated params: [0.099, -0.201, 0.299, -0.401, 0.499, -0.601, 0.699, -0.801] ZeRO-2 gathered params: [0.099, -0.201, 0.299, -0.401, 0.499, -0.601, 0.699, -0.801] max_abs_diff=0.000000000000 Adam m/v scalars per rank: DDP=16, ZeRO-2=4, reduction=4.00x max_abs_diff=0 是本文最重要的代码结果：同一份平均梯度、同一优化器规则下，把参数更新分给不同 rank 并不会改变更新后的模型。每卡 Adam 的 $m,v$ 元素数则从 16 降到 4，正好是 4 倍。真实系统还需要 bucket、异步通信、混合精度和异常恢复，但数学骨架就是这几十行。 05. 工业级实现对照 最小实现为了看懂“切什么”，生产实现要解决的则是“什么时候切、什么时候聚合、怎么让网络传输藏在计算后面”。下面以 2026-09-05 可见的官方实现为准。 5.1 PyTorch DDP：复制参数，反向时同步梯度 PyTorch 的 DistributedDataParallel 采用一进程一卡。它不会自动切分输入；应用通常用 DistributedSampler 保证不同 rank 读取不同样本。模型参数在每卡完整复制，autograd hook 在反向过程中把就绪梯度装进 bucket，并尽早发起 all-reduce，以便通信和后续层反向计算重叠。 这解释了两个常见现象：第一，DDP 通常比单进程 DataParallel 快，因为没有主卡收集和 Python 线程瓶颈；第二，DDP 加卡能缩短时间，却不会降低模型状态的单卡显存。它首先是一种吞吐扩展方案，不是大模型装载方案。 5.2 ZeroRedundancyOptimizer：最小改动的 Stage 1 思路 PyTorch 的 ZeroRedundancyOptimizer 可以和 DDP 组合。每个 rank 只为大约 $1/N$ 的参数维护本地优化器状态，更新自己负责的参数后广播结果，让所有 DDP 副本重新一致。这对应 ZeRO-1 的核心思想；参数仍完整复制，所以不能把它当成 FSDP 的替代品。 官方文档还标记该 API 为 experimental，并提示启用 overlap_with_ddp=True 时，最初若干迭代可能因为梯度 bucket 尚未稳定而不做参数更新。工程上不能只看“少了多少 GB”，还要读清 API 的更新时序、checkpoint 聚合和版本状态。 5.3 DeepSpeed ZeRO：三个阶段直接写进配置 DeepSpeed 官方 ZeRO 教程把三个阶段定义得很直接：Stage 1 切优化器状态，Stage 2 再切梯度，Stage 3 再切参数。一个典型配置是： { "zero_optimization": { "stage": 2, "contiguous_gradients": true, "overlap_comm": true, "reduce_bucket_size": 500000000 } } contiguous_gradients 针对碎片与通信连续性，overlap_comm 尝试把通信藏进计算，reduce_bucket_size 在“消息大到能吃满带宽”和“临时 bucket 不要撑爆显存”之间取舍。这些开关没有脱离前面的公式，只是在控制公式之外的临时项和时间轴。 如果 GPU 显存仍不够，ZeRO-Offload 可以把优化器状态与计算移到 CPU，ZeRO-Infinity 进一步使用 CPU 和 NVMe。但 官方 ZeRO-Offload 教程 展示的 10B 单 V100 案例并不意味着 offload 免费：PCIe 传输、CPU Adam 吞吐、NUMA 与磁盘带宽会变成新的瓶颈。 5.4 PyTorch FSDP：Stage 2/3 的框架化实现 知识树里的代码锚点 torch/distributed/fsdp/fully_sharded_data_parallel.py#FullyShardedDataParallel 在当前 PyTorch main 仍存在。FSDP1 的 ShardingStrategy 可以这样理解： NO_SHARD：参数、梯度、优化器状态都复制，行为接近 DDP； SHARD_GRAD_OP：梯度和优化器状态分片，参数在计算窗口内保持完整，接近 ZeRO-2； FULL_SHARD：参数、梯度和优化器状态全分片，接近 ZeRO-3； HYBRID_SHARD：节点内全分片、节点间复制，减少低带宽跨节点通信。 不过当前 PyTorch FSDP2 教程 已经明确建议迁移到 fully_shard：FSDP2 以 DTensor 做逐参数分片，不再依赖 FSDP1 的扁平参数；reshard_after_forward=True 对应 FULL_SHARD，设为 False 则更像 SHARD_GRAD_OP。文章或配置里只写“用了 FSDP”已经不够，必须同时说明代际和策略。 5.5 一层参数在 FSDP 中的生命周期 以 full shard 为例，一个 FSDP unit 在一次训练步里大致经历： 常驻状态只有本 rank 的参数分片； pre-forward hook 发起 all-gather，临时物化完整参数； 执行该 unit 的前向； 释放完整参数，恢复分片； 反向前再次 all-gather 参数； 计算梯度后执行 reduce-scatter，每个 rank 只留下本地梯度片； 本地优化器只更新本 rank 的参数与状态分片。 真正决定峰值的不是“最终只存 $1/N$”，而是同一时刻有多少 unit 正在 all-gather、预取队列有多深、最大 unit 有多大。如果把整个模型只包成一个巨型 unit，那么计算前仍要暂时物化全模型；如果把每个很小的算子都单独包起来，又会制造大量小消息，延迟和调度开销反而吞掉吞吐。Transformer 常按 block 包装，就是在这两端之间折中。 FSDP2 的公开契约仍是同一个时间逻辑：前向/反向前由 hook unshard，之后 reshard，梯度用 reduce-scatter 汇聚。实现细节从 flat parameter 变成 DTensor，不代表 ZeRO 的“按需物化”思想变了。 5.6 一个实用的选择顺序 先测量，再逐级加复杂度： 模型状态能放下，只想提吞吐：先用 DDP； Adam 状态是主要缺口：DDP + ZeRO-1/ZeroRedundancyOptimizer； 梯度也造成明显压力：ZeRO-2 或 FSDP 的 shard-grad 策略； 单卡连参数副本都放不下：ZeRO-3/FSDP full shard； 模型状态已经很小但长视频仍 OOM：处理激活重计算、micro-batch、FlashAttention 或序列并行，而不是继续折腾 ZeRO stage； 网络太慢：优先让分片组留在节点内，再用张量/流水线或混合分片扩到节点间。 06. 代价与边界 ZeRO 不减少总计算量。 同一个 batch 的前向、反向和 Adam 更新并没有少。它让每个 rank 少存状态，并通过通信在需要时恢复视图。若模型原本就能舒适放下，Stage 3 很可能只是增加通信和 hook 调度，吞吐反而下降。 通信字节相同，不代表时间相同。 ZeRO-2 与 DDP 在论文模型里都是约 $2G$，但一次大 all-reduce 和许多 layer-wise reduce-scatter/all-gather 的延迟特征不同。NVLink 节点内、InfiniBand 节点间、普通以太网的最优 bucket 大小不会一样；跨节点带宽不足时，1.5 倍 Stage 3 通信尤其明显。 平均显存很低，峰值仍可能 OOM。 参数 all-gather、预取、梯度归约、算子 workspace 可能同时在场。只看稳定阶段的 memory_allocated 会漏掉峰值；要记录 max_memory_allocated，并逐步调 wrap 粒度、prefetch 和 bucket。 ZeRO 解决不了激活随序列长度增长。 视频 DiT 中，时间、宽、高一起扩张，token 数可能成倍增加，注意力和 MLP 保存的激活随之增长。参数切到 1 GB 后仍然爆显存，并不说明 ZeRO 失效，而是瓶颈已经从模型状态转移到激活。 checkpoint 变复杂。 每个 rank 手里只有一片状态，保存时要选择 full、sharded 或 local state dict。把完整 checkpoint 聚合到 rank 0 可能让 CPU 内存瞬间成为瓶颈；只保存分片又要求恢复时正确处理 world size 和布局。训练能跑并不等于容灾链路可用，上线前至少做一次“保存—退出—换步数恢复”的演练。 全局操作必须认识分片。 梯度范数裁剪、参数检查、EMA、冻结部分参数、权重共享都不能默认“当前 rank 能看到完整张量”。例如 FSDP 提供自己的 clip_grad_norm_，就是因为全局范数需要跨分片归约。绕开框架直接遍历本地 .grad，算到的只是局部值。 offload 是拿带宽换容量。 CPU 内存比显存大，NVMe 又比 CPU 内存大，但层级越远，带宽越低、延迟越高。小模型或计算密度不够的模型会被数据搬运压垮；只有“不 offload 根本跑不了”或计算足以覆盖传输时，这个交换才划算。 大 batch 会改变优化问题。 增加数据并行度时，如果每卡 batch 不变，全局 batch 会随 $N$ 增长。ZeRO 保证同一全局 batch 下更新等价，却不保证扩大 batch 后收敛曲线不变。学习率、warmup、梯度累计和数据采样仍要一起调整。 什么时候不该用 ZeRO-3：模型与激活在单卡尚有充足余量、网络较慢、模型由大量极小模块组成且难以形成高效 bucket，或者你更看重最低延迟和调试简单性。此时 DDP 或 ZeRO-1/2 往往是更好的工程答案。 07. 经典论文脉络 这条路线不是“突然发明一种分布式训练”，而是一步步把数据并行里的冗余拆掉： Horovod（arXiv:1802.05799，2018）把 ring all-reduce 做成易接入的训练抽象，奠定了现代同步数据并行“各算各的、梯度集体归约”的工程基线；它解决吞吐扩展，但每个 worker 仍保存完整模型状态。 ZeRO（arXiv:1910.02054，2019）指出数据并行浪费的不是总显存，而是参数、梯度和优化器状态的重复副本；三阶段切分在保持数据并行计算粒度的同时，把模型状态显存从 $16P$ 推到 $16P/N$。 ZeRO-Offload（arXiv:2101.06840，2021）把优化器状态与计算搬到 CPU，并针对 CPU Adam 优化，让单张 32 GB V100 训练 10B 模型成为论文展示案例；留下的新问题是主机带宽和异构调度。 ZeRO-Infinity（arXiv:2104.07857，2021）继续把 CPU 与 NVMe 纳入内存层级，用带宽感知的分区和预取突破 GPU 内存墙；容量继续扩大，但 I/O 调度成为系统的核心。 PyTorch FSDP（arXiv:2304.11277，2023）总结了把 fully sharded data parallelism 纳入 PyTorch eager 训练栈的实践，让 ZeRO-3 类思想不再只属于单一外部训练引擎，并推动后续 FSDP2 的逐参数分片。 主线可以压缩成一句话：all-reduce 证明“计算可以复制、结果可以同步”，ZeRO 进一步问“既然结果会同步，状态为什么还要复制”。 后续 Offload、Infinity、FSDP 都是在回答状态放在哪、何时出现、以什么粒度搬运。 08. 常见误解 “数据并行会把模型切到多张卡上。” 不会。普通 DDP 是每卡一个完整模型，只切数据。能把吞吐从一张卡扩到八张，不代表能装下超过单卡容量的参数。 “ZeRO-3 后每卡永远只有 $1/N$ 参数。” 常驻状态是 $1/N$，计算某个 FSDP unit 前仍需临时 all-gather 完整参数。忽略这个瞬时窗口，就会得到理论显存能放、实际 forward 前仍 OOM 的配置。 “ZeRO-2 比 DDP 少传一半梯度，所以一定更快。” ZeRO-2 用 reduce-scatter 只留下梯度片，但更新后还要 all-gather 参数片。按论文的理想字节模型，总通信量与 DDP all-reduce 相同；快慢来自重叠、bucket、拓扑与实现，而非简单少一半。 “用了 ZeRO 就不用 activation checkpointing。” 两者切的是不同账本。ZeRO 处理模型状态冗余；activation checkpointing 用额外重算减少前向激活保存。长视频训练经常需要两者同时用。 “参数是 BF16，所以 Adam 也只占 2 字节。” 常见混合精度训练会保留 FP32 主参数与 FP32 的 $m,v$，优化器部分仍是 12 字节/参数。不同优化器、8-bit optimizer 或纯 BF16 更新会改变常数，因此使用 $16P$ 前必须先列出真实 dtype 与状态。 “64 卡 ZeRO-1 就一定省 4 倍。” 4 倍是 $N$ 很大时的渐近上限。64 卡的精确值是 $16/(4+12/64)=3.82$ 倍；8 卡只有 2.91 倍。宣传中的“up to”不能替代自己的显存账本。 “只要 loss 一样，所有数据并行更新都严格一致。” 浮点归约顺序会改变末位误差，随机数、dropout、数据尾 batch 和非确定性 kernel 也会影响复现。数学上等价不等于 bitwise identical；本文模拟得到 0 误差，是因为使用同一确定性顺序与双精度标量。 09. 动手验证 下面五个实验都能在没有 GPU 的机器上完成。两个脚本全文在附录，复制保存后直接用 Python 3 运行。 实验一：复现原论文表格。 python zero_memory_ledger.py 预期看到 7.5B/64 卡下 DDP、ZeRO-1/2/3 分别为 120.00、31.41、16.64、1.88 GB。如果结果差约 7%，检查自己是不是把 GiB 和十进制 GB 混用了。 实验二：观察有限卡数离理论上限有多远。 python zero_memory_ledger.py --params-b 1.5 --world-size 8 预期 ZeRO-1 只省 2.91 倍，ZeRO-2 只省 4.27 倍，ZeRO-3 才精确省 8 倍。再把 world-size 改成 2、4、16、64，观察 Stage 1 向 4 倍、Stage 2 向 8 倍收敛。 实验三：验证更新等价性。 python zero_update_simulator.py 预期 DDP 与 ZeRO-2 的八个新参数完全相同，max_abs_diff=0；每 rank 的 Adam $m,v$ 标量数从 16 降为 4。把 local_gradient() 的公式改掉，只要两条路径仍使用同一平均梯度，最终参数就应继续一致。 实验四：故意制造尾 batch 权重错误。 在模拟器里让最后一个 rank 只代表一个样本、其他 rank 各代表两个样本，然后比较“先对每个 rank 求平均再除以 4”和“按总样本数加权平均”。预期两个梯度不同。这能解释为什么真实训练要使用 DistributedSampler、drop_last 或正确的样本权重，而不能只相信 all-reduce。 实验五：给自己的模型补全账本。 从训练日志记录参数量、每种状态 dtype、峰值激活与临时 buffer。先用脚本算模型状态理论值，再和框架峰值相减。如果差值随序列长度或分辨率快速增长，瓶颈是激活；如果几乎不随输入变化，才优先继续切模型状态。预期这一步比直接试三个 ZeRO stage 更快找到真正的 OOM 原因。 10. 延伸阅读 读完这篇可以沿分布式训练知识树继续： 混合精度与数值稳定性：本文的 16 字节常数建立在混合精度 Adam 上；理解 FP16、BF16、FP8 的动态范围后，才能判断哪些状态可以继续降精度。 张量并行与流水线并行：当单个算子或单层连临时 all-gather 都放不下时，需要切计算本身；下一步要算 TP 通信量和 PP 气泡率。 序列并行与 Ring Attention：模型状态已经切完，长视频仍被激活卡住时，才轮到沿序列维度切注意力计算。 本文的核心资料均来自原始论文与官方实现：ZeRO 的显存和通信公式以 ZeRO 原论文为准；三阶段配置参考 DeepSpeed ZeRO 官方教程；DDP、优化器分片和 FSDP 策略分别核对了 PyTorch DDP 文档、ZeroRedundancyOptimizer 文档 与 FSDP 文档。代码路径和 API 会随上游重构，工业实现部分注明的日期就是版本边界。 附录：完整代码 09 节用到的脚本全文如下（zero_memory_ledger.py、zero_update_simulator.py、make_figures.py）。复制到本地存成同名文件，按各脚本开头的依赖说明准备环境后即可运行。 zero_memory_ledger.py #!/usr/bin/env python3 """复现 ZeRO 论文中的混合精度 Adam 模型状态显存账本。 这里只计算参数、梯度和优化器状态，不包含激活、临时通信缓冲区、 CUDA context 与内存碎片。单位同时使用十进制 GB，便于和论文表格对照。 """ from __future__ import annotations import argparse from dataclasses import dataclass @dataclass(frozen=True) class Stage: name: str replicated_bytes: float sharded_bytes: float communication_multiple: float def bytes_per_parameter(self, world_size: int) -> float: return self.replicated_bytes + self.sharded_bytes / world_size STAGES = ( Stage("DDP", 16.0, 0.0, 1.0), Stage("ZeRO-1", 4.0, 12.0, 1.0), Stage("ZeRO-2", 2.0, 14.0, 1.0), Stage("ZeRO-3", 0.0, 16.0, 1.5), ) def gb(byte_count: float) -> float: """转为十进制 GB；ZeRO 原论文的 120 GB 使用这个口径。""" return byte_count / 1_000_000_000 def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--params-b", type=float, default=7.5, help="参数量，单位十亿；默认复现论文的 7.5B 示例") parser.add_argument("--world-size", type=int, default=64, help="数据并行进程数") args = parser.parse_args() if args.params_b <= 0 or args.world_size <= 0: raise ValueError("params-b 与 world-size 必须为正数") params = args.params_b * 1_000_000_000 gradient_bytes = params * 2 # fp16/bf16 梯度 baseline_comm = 2 * gradient_bytes # ring all-reduce: reduce-scatter + all-gather print(f"model={args.params_b:g}B, world_size={args.world_size}") print("assumption: fp16 params 2B + fp16 grads 2B + " "fp32 master/m/v 12B = 16 bytes/parameter") print("stage bytes/param model-state GB vs DDP comm/step GB") ddp_bytes = STAGES[0].bytes_per_parameter(args.world_size) for stage in STAGES: bpp = stage.bytes_per_parameter(args.world_size) state_gb = gb(params * bpp) saving = ddp_bytes / bpp comm_gb = gb(baseline_comm * stage.communication_multiple) print(f"{stage.name:<8} {bpp:>10.4f} {state_gb:>17.2f} " f"{saving:>9.2f}x {comm_gb:>15.2f}") print("note: activation, temporary buffers and fragmentation are excluded") if __name__ == "__main__": main() zero_update_simulator.py #!/usr/bin/env python3 """用单进程模拟 DDP 与 ZeRO-2 的一次 Adam 更新。 脚本不依赖 GPU、PyTorch 或分布式运行时。四个“rank”先各自产生一份局部 梯度，再比较两条路径：DDP 对完整梯度做 all-reduce；ZeRO-2 对梯度做 reduce-scatter、每个 rank 只更新自己的参数片，最后 all-gather 参数。 """ from __future__ import annotations import math WORLD_SIZE = 4 PARAMS = [0.10, -0.20, 0.30, -0.40, 0.50, -0.60, 0.70, -0.80] LEARNING_RATE = 1e-3 BETA1 = 0.9 BETA2 = 0.999 EPSILON = 1e-8 def local_gradient(rank: int, width: int) -> list[float]: """构造确定性的局部梯度，模拟每个 rank 看见不同数据。""" return [((rank + 1) * (index + 2) - 3) / 100.0 for index in range(width)] def average_columns(rows: list[list[float]]) -> list[float]: return [sum(column) / len(rows) for column in zip(*rows)] def shard(vector: list[float], world_size: int) -> list[list[float]]: if len(vector) % world_size: raise ValueError("为了让示例清楚，参数量必须能被 world_size 整除") width = len(vector) // world_size return [vector[rank * width:(rank + 1) * width] for rank in range(world_size)] def adam_first_step(params: list[float], grads: list[float]) -> tuple[list[float], list[float], list[float]]: """执行 Adam 的第 1 步，并返回新参数、m、v。""" new_params: list[float] = [] m_state: list[float] = [] v_state: list[float] = [] for value, grad in zip(params, grads): m = (1.0 - BETA1) * grad v = (1.0 - BETA2) * grad * grad m_hat = m / (1.0 - BETA1) v_hat = v / (1.0 - BETA2) updated = value - LEARNING_RATE * m_hat / (math.sqrt(v_hat) + EPSILON) new_params.append(updated) m_state.append(m) v_state.append(v) return new_params, m_state, v_state def main() -> None: local_grads = [local_gradient(rank, len(PARAMS)) for rank in range(WORLD_SIZE)] # DDP：每个 rank 经 all-reduce 得到同一份完整平均梯度，并完整更新参数。 reduced_full_grad = average_columns(local_grads) ddp_params, ddp_m, ddp_v = adam_first_step(PARAMS, reduced_full_grad) # ZeRO-2：reduce-scatter 的结果等价于先平均，再只保留所属分片。 param_shards = shard(PARAMS, WORLD_SIZE) grad_shards = shard(reduced_full_grad, WORLD_SIZE) updated_shards: list[list[float]] = [] state_sizes: list[int] = [] for params_for_rank, grads_for_rank in zip(param_shards, grad_shards): updated, local_m, local_v = adam_first_step(params_for_rank, grads_for_rank) updated_shards.append(updated) state_sizes.append(len(local_m) + len(local_v)) # all-gather 后每个 rank 都能拿到同一份新参数；这里拼接一次代表该结果。 zero_params = [value for part in updated_shards for value in part] max_abs_diff = max(abs(a - b) for a, b in zip(ddp_params, zero_params)) shard_width = len(PARAMS) // WORLD_SIZE print(f"world_size={WORLD_SIZE}, parameter_count={len(PARAMS)}, " f"shard_width={shard_width}") print(f"local gradient matrix shape=({WORLD_SIZE}, {len(PARAMS)})") print(f"reduce-scatter result shape=({WORLD_SIZE}, {shard_width})") print("averaged gradient:", [round(x, 4) for x in reduced_full_grad]) print("DDP updated params:", [round(x, 6) for x in ddp_params]) print("ZeRO-2 gathered params:", [round(x, 6) for x in zero_params]) print(f"max_abs_diff={max_abs_diff:.12f}") print(f"Adam m/v scalars per rank: DDP={len(ddp_m) + len(ddp_v)}, " f"ZeRO-2={state_sizes[0]}, reduction={WORLD_SIZE:.2f}x") if __name__ == "__main__": main() make_figures.py """Regenerate this article's deterministic teaching figure (numpy + matplotlib).""" from pathlib import Path import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt import numpy as np OUT=Path(__file__).resolve().parents[1]/"figures" OUT.mkdir(exist_ok=True) plt.rcParams.update({"font.sans-serif":["PingFang SC","Arial Unicode MS","DejaVu Sans"],"axes.unicode_minus":False}) n=np.array([1,2,4,8,16,32,64]); p=7.5e9 fig,ax=plt.subplots(figsize=(8,4.8)) for name,y in [("DDP",np.full(n.shape,16.)),("ZeRO-1",4+12/n),("ZeRO-2",2+14/n),("ZeRO-3",16/n)]: ax.plot(n,y*p/1e9,"o-",label=name) ax.set(xscale="log",yscale="log",xlabel="Data-parallel ranks",ylabel="Model states per rank (GB)",title="7.5B model: FP16 weights/grads + FP32 master/m/v") ax.legend();ax.grid(alpha=.25);fig.tight_layout() fig.savefig(OUT/'zero_states.png',dpi=170) plt.close(fig) print(OUT/'zero_states.png') 更多 AIGC 论文解读，关注微信公众号「人工智能炼丹君」 每日更新 · 论文精选 · 深度解读 · 技术脉络 微信搜索 人工智能炼丹君 或扫描下方二维码关注

AIGC 基本功|视频生成中的强化学习与奖励模型-VideoRL 视频生成中的强化学习与奖励模型 所属方向：对齐与强化学习 ｜ 难度：前沿专题 ｜ 前置知识：DiffusionRL、FlowMatching 关键词：视频强化学习、奖励模型、Flow-GRPO、时序一致性奖励、人类偏好、reward hacking 01. 为什么需要它 先看一个用于说明奖励投机的假设场景；下面的“两千步”不是本文实际运行的视频实验记录。 你有一个文生视频模型，想用 RL 把它对齐到人类偏好。你手上有个打分器，在图像上验证过效果不错——清晰度、细节、构图都能打分。你把它接上视频，跑 GRPO，训了两千步。奖励曲线漂亮地往上走。 然后你去看生成结果：画面几乎不动了。 人物站在原地，衣角不飘，背景纹理纹丝不动。prompt 明明写着 "a person walking across the street"，模型给你一张会呼吸的照片。 这不是 bug，是模型算得比你清楚。静止画面有一系列白拿的好处： 形状永远不会崩——因为它不变 帧间没有闪烁——因为帧之间完全一样 细节可以做到极精细——不用分算力去渲染运动模糊 而它放弃的只有「动态性」这一项。如果静止带来的加分超过动态性与指令遵循等全部损失，那么静止可能比运动拿到更高奖励，RL 会毫不犹豫地找到它。 视频奖励增加了运动与时序一致性维度。在有限模型能力下，运动幅度与形变、闪烁等伪影可能存在经验权衡；这不是“运动必然降低画质”的物理定律。图像的清晰度、构图、语义也可能冲突，不能假设图像奖励天然相容。 后面几节会把这笔账算成具体数字——用 VisionReward 真实发布的 29 个权重。 02. 最小可用理解 三句话： 采一组。对同一个 prompt 采样 $G$ 条视频（比如 8 条），用奖励模型给每条打一个分。 组内比。把这组分数减去组内均值、除以组内标准差，得到每条的优势（advantage）——比同组平均好就是正，差就是负。 按优势推。优势为正的往上推概率，为负的往下压。不需要训练价值网络（critic），组内均值就当基线用了。 这概括了 GRPO 的组内优势机制；完整目标还包含后文的概率比、裁剪和可选 KL。它省掉 critic 的代价是必须一次采一组，好处是少训一个网络、少一堆调不动的超参。 如果只记一件事：RL 不会给你「更好的视频」，只会给你「奖励更高的视频」。这两者的差距全部藏在奖励模型里。本文一半篇幅在讲奖励模型，不是跑题。 03. 数学推导 3.1 从策略梯度到组内基线 目标很直白：最大化生成样本的期望奖励。 $$J(\theta) = \mathbb{E}_{x \sim \pi_\theta} [r(x)]$$ 其中 $\pi_\theta$ 是生成模型（策略），$x$ 是一条采样出来的视频，$r(x)$ 是奖励模型给的分。对参数求梯度，用对数导数技巧： $$\nabla_\theta J = \mathbb{E}_{x \sim \pi_\theta} [r(x) \nabla_\theta \log \pi_\theta(x)]$$ 这个式子能用，但方差大得没法训。原因是 $r(x)$ 的绝对值直接乘在梯度上：假如所有样本的奖励都在 8.0 到 8.2 之间，那么每个样本都在说「往我这边推」，只是力度略有差别。真正有用的信息是「谁比谁好」，而不是「绝对分多高」。 标准解法是减一个基线 $b$： $$\nabla_\theta J = \mathbb{E}[(r(x) - b) \nabla_\theta \log \pi_\theta(x)]$$ 只要 $b$ 与 $x$ 无关，这个替换是无偏的（因为 $\mathbb{E}[b \nabla_\theta \log \pi_\theta] = b \nabla_\theta \mathbb{E}[1] = 0$）。PPO 那一路用一个学出来的价值网络当 $b$，GRPO 的选择更省事——同一个 prompt 采一组，用组内均值当基线： $$A_i = \frac{r_i - \mathrm{mean}(r_1, \dots, r_G)}{\mathrm{std}(r_1, \dots, r_G)+\varepsilon}$$ 这里不能直接套用上面的无偏证明：组均值含有 $r_i$，依赖当前样本。固定 prompt、独立同分布采样且不除标准差时， $$E[(r_i-\bar r)\nabla\log\pi_\theta(x_i)]=(1-1/G)\nabla J.$$ 使用不含自身的 leave-one-out 均值可消除该缩放；除以随机组标准差还会改变样本及 prompt 的相对权重。$\varepsilon>0$ 防止同分组除零，同分组的奖励优势为零。标准化是有用的工程选择，不能据此宣称完整 GRPO 无偏；奖励均值高本身也不等于期望梯度大。 这张图要看什么：左图是同一个 prompt 采出的 8 条视频的奖励，全挤在 8.0 到 8.2 之间——如果直接把 $r_i$ 乘在梯度上，8 条样本会一起喊「往我这边推」，有用的信息全被绝对值淹没（箭头方向完全一致）。右图减掉组内均值、除以组内标准差之后，同样的 8 个数字变成了有正有负的优势：绿的往上推、红的往下压。这就是 GRPO 不用 critic 也能训的原因——基线是从同组样本里白捡的（图由 code/make_figures.py 生成，下同）。 GRPO 出自 DeepSeekMath（arXiv:2402.03300），原本是给语言模型的数学推理用的，后来被搬到视觉生成上。 3.1.1 完整目标里还有两个护栏 上面那个式子是骨架。真实的 GRPO 目标还包着两层保护，后面 04 节的最小实现会把它们省掉，这里先说清省的是什么。 第一层是重要性采样比的截断（沿用 PPO 的做法）： $$\mathcal{L} = -\mathbb{E}\left[ \min\left( \rho_i A_i, \ \mathrm{clip}(\rho_i, 1-\epsilon, 1+\epsilon) A_i \right) \right]$$ 式子里的 $\rho_i$ 是新旧策略的概率比，$\rho_i = \pi_\theta(x_i) / \pi_{\theta_{\text{old}}}(x_i)$。采样用的是旧策略 $\pi_{\theta_{\text{old}}}$，更新的是新策略 $\pi_\theta$，$\rho_i$ 修正这个偏差。clip 截断的是朝有利方向继续外推的收益，不保证每个 ratio 都在区间内，也不是 KL 的硬约束。神经网络共享参数，仍可能越界；需结合 KL 和梯度日志监控。视觉去噪实现通常按每个随机转移计算 ratio，不能把这里的整段视频记号误当成可直接算出的边缘密度。 第二层是对参考模型的 KL 惩罚： $$\mathcal{L}_{\text{total}} = \mathcal{L} + \beta \, \mathrm{KL}\left[ \pi_\theta \, \| \, \pi_{\text{ref}} \right]$$ $\pi_{\text{ref}}$ 通常是 RL 开始前的初始模型。这一项把策略拴在原始分布附近，是对付 reward hacking 最基础的护栏——因为大多数 hack（比如 01 节那个静止画面）都要求策略跑到离原始分布很远的地方去。$\beta$ 调的就是「允许它跑多远」。 代价是 $\beta$ 很难调：太大则学不动，太小则护栏形同虚设，而合适的值随任务和奖励模型变化。这也是为什么 04 节要先把护栏拆掉看清裸的优化行为——理解 hacking 怎么发生，比直接套护栏更重要。 3.2 flow matching：怎样得到随机转移的似然 Flow matching 的 ODE 为 $dx_t/dt=v_\theta(x_t,t)$。本文用 $t=1$ 表示噪声、$t=0$ 表示数据：给定初始噪声，逐步转移是确定性的条件 delta，不能直接代入高斯转移 log-prob 做 PPO/GRPO。 这不等于输出没有概率分布。随机初值经正则、可逆的 ODE 流仍可产生具有边缘密度的输出，CNF 可用散度积分计算密度；可微奖励还可以沿 ODE 用路径导数优化，确定性策略梯度也有自己的适用条件。受限的是本文这条“随机逐步转移似然 + score function”的路线。 Flow-GRPO 第 4.2 节 为它构造反向时间 SDE： $$dx_t=\left[v_\theta(x_t,t)-\frac{\sigma_t^2}{2}\nabla\log p_t(x_t)\right]dt+\sigma_t\,d\bar W_t,\qquad dt<0.$$ 由于时间从 1 向 0 走，score 修正前是减号；离散噪声标准差为 $\sigma_t\sqrt{-dt}$。若改用从噪声出发递增的时间坐标，漂移和 score 符号必须一起变换。准确 score、匹配的速度场与连续时间求解条件下，SDE 和 ODE 共享边缘分布；近似网络、有限步长及 RL 更新会带来偏差，不能保证实现后精确不变。 这使被训练的非退化转移具有可计算的高斯 log-prob。终步若噪声为零，仍要排除或另行处理，不能对零方差求对数密度。Flow-GRPO 还用 Denoising Reduction 减少训练采样步数、保留完整推理步数；这是论文所测任务上的经验结果，需在视频任务重新验证。 3.3 奖励从哪来：把打分变成一串是非题 现在缺的是 $r(x)$。视频的「好」是多维的，而人类偏好数据通常只有「A 比 B 好」这种成对比较，信息量很稀疏。 VisionReward（arXiv:2412.21059）的做法是把打分拆成一串可解释的是非题，再线性加权： $$r(x) = \frac{1}{K} \sum_{k=1}^{K} w_k \cdot a_k(x), \quad a_k(x) \in \{+1, -1\}$$ 其中 $a_k$ 是一个视觉语言模型对第 $k$ 个问题的回答（yes 记 $+1$，no 记 $-1$），$w_k$ 是该问题的权重。视频版一共 $K = 29$ 个问题，从「是否满足 prompt 的全部要求」到「细节是否精细」。 这个设计有两个好处。可解释：分低的时候能看出是哪一维拖的，不是一个黑盒数字。可控：想让模型更重视某一维，直接改 $w_k$ 就行。 代价是它把「好」压成了一个线性组合。而线性组合最容易被套利——下一节就用它真实发布的权重，把套利空间算出来。 04. 代码实现 4.1 先把奖励函数照抄一遍 VisionReward 的打分逻辑短到可以完整贴出来。这是它仓库里 inference-video.py 的核心两行（2026-08 的 main 分支）： answers = np.array([1 if answer == 'yes' else -1 for answer in answers]) return np.mean(answers * weight).item() 有个细节值得停一下：no 记的是 $-1$，不是 $0$。这让答错变成实打实的扣分，而不只是「没拿到加分」。同一个维度上，yes 和 no 之间差了 $2 w_k$ 而不是 $w_k$——算套利空间时这个因子 2 不能漏。 29 个权重是仓库里 VisionReward_Video/weight.json 直接给出的。先看看它们长什么样（code/reward_hacking.py）： 29 个维度，权重总和 6.2772，均值 0.2165 最高 1.1418（指令-未完全失败） 最低 0.0085（细节-不粗糙） 最高／最低 = 134 倍 指令遵循三档合计 2.3486，占总权重 37.4% 构图＋色彩合计 0.0924，占总权重 1.5% 两个数字值得盯着看。最高和最低差 134 倍——这不是一个「各维度都重要」的均衡奖励，而是有着强烈倾向的。指令遵循三档占了 37.4%，构图加色彩只占 1.5%：在真实的人类偏好数据里，指令三档的权重合计是构图与色彩权重合计的约 25 倍。但这不是人类把指令遵循看得比整体美感重 25 倍的因果结论：其他维度也反映画质，且题目相关性、回答频率与权重尺度会影响解释。 这个结构反直觉但合理——观众要的是他要的东西，不是一张漂亮但无关的图。 这张图要看什么：左图把 29 个权重从低到高排开——最高的「指令-未完全失败」1.1418，最低的「细节-不粗糙」0.0085，相差 134 倍（红色是指令遵循三档，灰色是构图加色彩）。右图按语义归堆后更刺眼：指令遵循合计 2.3486、一家吃掉 37.4%，构图加色彩只有 0.0924、占 1.5%。指令三档的权重合计是构图与色彩这两组的二十多倍——记住这个比例，4.2 和 4.4 节的套利与翻转全是从它推出来的。 4.2 静态套利的账 回到 01 节那个静止画面。它能白拿哪些分、要放弃哪些分，现在可以精确算： 静止能白拿：全程形状保持＋不混乱＋画质稳定 = 0.7166 静止要放弃：镜头高度动态＋不微弱 = 0.1634 净收益 +0.5532 （yes/no 相差 2 个单位，实际影响 ×2 = +1.1064） 上述 +1.1064 是未除以 $K$ 的加权和变化；score() 实际变化是 $1.1064/29\approx0.03815$。权重总和 6.2772 对应分数上限 0.2165，分数全范围宽度约 0.4329，因此这笔差额约占全范围 8.8%。这只是所选维度的局部账，不能省略指令等其他维度后宣称静止全局最优。 那它会不会因为违背 prompt 而被罚回来？我构造了两个视频画像对比（prompt 要求「人物走过街道」）： 视频 A（几乎静止的精致画面，没照做 prompt） +0.1220 视频 B（照做了 prompt，但有形变和抖动） +0.1144 奖励模型更偏好：A（静态） 差值 0.0076 静态那个赢了。 拆开看差异来自哪： -1.9088 指令-全部满足 +0.8586 细节-非常精细 +0.5372 画质-非常稳定 +0.5272 全程形状-不混乱 -0.5048 指令-大部分满足 +0.3688 全程形状-完美保持 静态视频在指令遵循上亏了 2.41（这是权重最高的维度！），但靠画质那几项一点点堆回来，最后反超 0.0076。 需要说清楚：这两个 yes/no 画像是我按各维度语义人工假设的，不是真实模型的输出，所以这个结论说明的是「权重结构允许这样的套利」，不是「VisionReward 实际会这么判」。差值只有 0.0076 也说明它非常接近临界——是否会被优化器实际找到，还取决于策略可达性、采样噪声与正则约束，不能由这一对假设画像推断必然收敛。 4.3 让 GRPO 自己去找套利 手工构造画像终究不如让算法自己找。下面是完整的单参数 GRPO（code/grpo_minimal.py），策略只有一个参数：运动强度 $m$。 维度压到 1 维是刻意的——才能看清 GRPO 究竟把策略推向哪里。玩具概率按人为假设设定：$m$ 越大，指令遵循和动态性越容易拿 yes，形状保持、画质稳定、细节精细越容易拿 no。 def grpo(steps=400, group=32, lr=0.35, sigma=0.15, theta0=0.0, seed=0, weight=None): rng = np.random.default_rng(seed) theta = theta0 for step in range(steps): # 策略：m = clip(sigmoid(theta) + noise)，高斯探索 mu = 1 / (1 + np.exp(-theta)) eps = rng.normal(0, sigma, group) ms = np.clip(mu + eps, 0.0, 1.0) rs = np.array([reward(m, rng, weight) for m in ms]) # 组内相对优势：减均值除标准差（3.1 节那个式子） adv = (rs - rs.mean()) / (rs.std() + 1e-8) # ∂mu/∂theta = mu(1-mu) grad = float(np.mean(adv * eps / sigma ** 2) * mu * (1 - mu)) theta += lr * grad return 1 / (1 + np.exp(-theta)) m_final = grpo() print(f"GRPO 收敛到 m = {m_final:.4f}") # GRPO 收敛到 m = 0.9657 十几行，没有 critic，也刻意去掉了 3.1.1 节那两层护栏（clip 和 KL）——保留的是组标准化 REINFORCE 示意，不是完整 PPO/GRPO 更新。代码中的高斯 score 对应裁剪前潜在动作 $a=\mu+\epsilon$，环境再执行 $m=\mathrm{clip}(a,0,1)$；不能把它当成裁剪后混合分布的普通高斯 log-prob。用完整的 29 个权重跑： step 运动强度 m 组内平均奖励 0 0.5000 -0.0238 80 0.9227 +0.0395 160 0.9420 +0.0346 320 0.9649 +0.0452 399 0.9660 +0.0370 GRPO 收敛到 m = 0.9657 → 策略选择：动态（照做了 prompt） 没有 hack。 策略从中性的 0.5 一路推到 0.966，选择了照做 prompt。换三个不同初值（0.27 / 0.50 / 0.73）出发，都收敛到 0.977 附近，说明这不是初值凑巧。 原因就是 4.1 节那个 37.4%：指令遵循的权重足够大，压住了画质维度的全部诱惑。这是玩具里的结果，不能验证 Flow-GRPO 的视频表现。Flow-GRPO 在其图像实验中他们报告「very little reward hacking occurred」，奖励涨上去了而画质和多样性没有明显退化。 4.4 关键对照：把指令遵循拿掉 那 01 节那个失败场景怎么发生的？答案在奖励模型上，不在算法上。 01 节的设定是「手上有个在图像上验证过的打分器」——图像打分器很可能根本不看时序，或者对指令遵循的权重远没有这么高。把三个指令遵循维度的权重置零，模拟这种情况，同一个 GRPO 再跑一遍： step 运动强度 m 组内平均奖励 0 0.5000 -0.0016 80 0.0305 +0.0269 160 0.0178 +0.0363 320 0.0099 +0.0396 399 0.0084 +0.0324 收敛到 m = 0.0084（完整权重下是 0.9657） 从 0.966 翻转到 0.008。 策略选择了几乎完全静止。注意中间那列奖励——它一路从 $-0.0016$ 涨到 $+0.04$，总体上也比初始值高，但逐步有明显随机波动。这就是 01 节说的：奖励曲线不能告诉你模型学到了什么。 奖励曲面看得更清楚——同一组物理约束，两种权重下的形状完全相反： m 完整权重 去掉指令遵循 0.0 -0.0393 +0.0393 0.3 -0.0298 +0.0228 0.5 +0.0006 -0.0006 0.8 +0.0446 -0.0322 1.0 +0.0379 -0.0394 → 完整权重的最高点在 m = 0.8（+0.0446） → 去掉指令遵循后最高点在 m = 0.0（+0.0393） 完整权重下曲面在 $m = 0.8$ 处见顶；去掉指令遵循后，在这一玩具设定下整体偏向低运动强度；表格是每点 4000 次随机采样的估计，不能仅凭稀疏网格确立精确单调性或全局最优。 这张图要看什么：左图是同一组物理约束下的奖励曲面——完整权重（蓝实线）在 $m=0.8$ 附近见顶，去掉指令遵循三档后（红虚线）整条曲线翻转成单调递减，最高点直接落到 $m=0$。右图是同一个 GRPO、同一个初值跑 400 步：蓝线冲到 0.966（照做 prompt），红线掉到 0.008（躺平）。而两条浅色细线是各自的组内平均奖励——二者总体奖励改善，但都有采样噪声，不能靠奖励曲线判断是否符合需求。奖励曲线不能告诉你模型学到了什么。 结论落在这里：reward hacking 不是算法出了毛病，而是它精确地最大化了你真正写下来的那个目标。 同一个 GRPO、同一个物理约束，只改奖励的权重结构，行为就翻转了。查 hacking 该去查奖励函数，而不是调 RL 的超参。 还要区分两种目标：图中的曲面是在固定运动强度 $m$ 上的采样均值，策略优化的却是高斯探索后再裁剪的期望奖励。$\mu=0.966$ 是裁剪前策略均值参数，不是实际运动强度的期望；它与网格中 $m=0.8$ 的峰值不能直接比较。因此这些数字不能证明“策略梯度过冲”，更不能归因为探索噪声必然造成方向性偏置。 05. 工业级实现对照 参考实现（以 2026-08 的 main 分支为准，上游会重构）： THUDM/VisionReward → inference-video.py：score、compare_two_videos THUDM/VisionReward → VisionReward_Video/weight.json：29 个线性权重 上面的最小实现把奖励当成一个现成函数，生产上它是个 VLM，差别都在这。 5.1 打分要跑 29 次前向 score() 的实际流程是：抽帧、对每个问题拼一次 prompt、让 VLM 输出 yes/no。29 个问题就是 29 次 VLM 前向。 这直接决定了 RL 训练的成本结构。GRPO 每步要采一组（比如 8 条视频），每条视频要 29 次 VLM 前向，也就是每个训练步 232 次前向——而这还没算生成那 8 条视频本身的去噪开销。奖励计算不是附带开销，它经常是训练循环里最贵的一环。这也是 Flow-GRPO 要做 Denoising Reduction 的现实压力来源。 5.2 两套问题清单，别拿错 仓库里有两个文件长得很像： VisionReward_video_qa.txt：64 个问题，用于细粒度查询（问模型某一项如何） VisionReward_video_qa_select.txt：29 个问题，用于打分 inference-video.py 顶部读的是 _select 那个： QUESTIONS_PATH = "VisionReward_Video/VisionReward_video_qa_select.txt" WEIGHT_PATH = "VisionReward_Video/weight.json" 29 条问题对应 29 个权重，一一对齐。拿 64 个那份去乘权重会直接维度不匹配——这算好事，至少会报错。 5.3 比较两个视频不是比分数 想判断 A 和 B 哪个好，直觉是各算一次 score() 再比大小。它没这么做——compare_two_videos() 的结尾是： answers1 = np.array([1 if answer == 'yes' else -1 for answer in answers1]) answers2 = np.array([1 if answer == 'yes' else -1 for answer in answers2]) diff = answers1 - answers2 return np.sum(diff * weight).item() > 0 先求两个答案向量的逐维差 diff，再加权求和判正负。 只要两次打分使用同一组答案与权重，这与比较两个 score() 完全等价（仅差正的常数 $K$），并不能阻止不同维度相互抵消。VisionReward 论文的“多维一致性”另指偏好优化时筛选在各语义维度都一致占优的样本对；不是把减法移到加权和里面。 5.4 那些数字：这套设计到底有多少提升 VisionReward 报告的两个结果： 偏好预测准确率比 VideoScore 高 17.2% 用它做奖励的文生视频模型，pairwise win rate 比用 VideoScore 高 31.6% 这两个百分比对应不同指标与实验，不能相除推断“奖励改进被 RL 放大了”。VisionReward 的偏好优化还包括 DPO 等流程；本文引用的是论文报告结果，不是视频 GRPO 的对照实测。 06. 代价与边界 杠杆是双向的。 持续优化可能放大奖励偏差，但幅度不是由 5.4 节的两个异质指标决定的。奖励模型里一个不起眼的偏见——比如偏爱暖色调、偏爱中心构图——在评测里可能只是一两个百分点，但 RL 会把它当成目标，训几千步之后就是全部输出都发黄。奖励模型的偏见不会被平均掉，会被放大。 优化的是奖励，不是质量。 这句话看着像废话，但它有个实践推论：奖励曲线上升不能作为训练成功的证据。4.4 节那次 hack 的奖励曲线（$-0.0016 \to +0.04$）和 4.3 节那次成功训练长得一模一样。 那怎么才能发现？三个手段，按性价比排： 盯住奖励模型看不见的指标。 挑几个不参与奖励计算的量化指标定期记录，比如帧间光流的平均幅度、逐帧 LPIPS 差异。静态 hack 在这类指标上一眼就露馅（光流幅度趋近 0），而奖励模型完全不看它们。这个办法便宜、能自动化，应该默认开着。 看奖励的维度分解，不只看总分。 VisionReward 这种线性结构的好处正在这里——把 29 维的得分分别记下来。总分涨、但某几维在持续下跌，就是套利正在发生。4.4 节第 5 段那张表就是这么读的。 固定一批 prompt 定期人眼看。 最贵但不可替代。前两个手段只能发现你预料到的失效模式，人眼能发现你没预料到的。 至于用第二个奖励模型交叉验证——听起来对称，实际有个陷阱：如果两个奖励模型是用同源的偏好数据训的，它们很可能共享同样的盲区，交叉验证会一起点头。要用就得选训练数据和架构都不同的。 成本很实在。 按 5.1 节的账，每个训练步是「$G$ 条视频的完整采样＋ $29G$ 次 VLM 前向」。视频采样本身就比图像贵一个量级（多了时间维度），再乘上组大小，这是 RL 在视频上落地慢的主要原因。Denoising Reduction 这类工程手段不是优化项，是可行性前提。 线性加权是双刃剑。 它带来可解释性，也带来可套利性——线性函数没有交互项，意味着「这一维差到极点」不会拖累其他维度的加分。真实的人类判断不是这样的：一个完全静止的「视频」在人眼里直接不合格，不管它多精细。线性模型表达不了这种否决关系，所以 04 节那个套利在数学上成立。 什么时候不该用 RL。 如果你的问题能用更直接的手段解决，就别上 RL。想让输出更清晰——去修数据和 VAE；想让它听懂 prompt——先确认是不是文本编码器或者标注质量的问题。RL 适合的是「说得清好坏、但写不出损失函数」的目标，比如整体美感、运动自然度这类只能靠比较来表达的偏好。上 RL 之前先问一句：我的奖励模型真的比我的损失函数更懂这件事吗？ 07. 经典论文脉络 这条线分两支：怎么优化和拿什么当奖励。两支交替推进。 奖励这一支： ImageReward（arXiv:2304.05977，2023-04）第一个规模化的文生图人类偏好模型，确立了「收集成对偏好数据训一个打分器」这个范式。它是个黑盒标量，可解释性问题从这里就留下了。 VideoScore（arXiv:2406.15252，2024-06）把细粒度人类反馈搬到视频上，成为视频奖励模型的常用基线——也就是 VisionReward 那两个数字（17.2% / 31.6%）超越的对象。 VisionReward（arXiv:2412.21059，2024-12）用分层是非题 + 线性加权换来可解释性，图像和视频统一处理。代价是线性结构本身可被套利（06 节）。 优化这一支： DDPO（arXiv:2305.13301，2023-05）把去噪过程当成多步决策，第一次让 policy gradient 在扩散模型上跑通。 DPOK（arXiv:2305.16381，2023-05）同期工作，补上了 KL 正则——防止策略为了刷奖励跑离原始分布太远。这是对付 hacking 最基本的护栏。 Diffusion-DPO（arXiv:2311.12908，2023-11）绕开 RL：直接用成对偏好数据做监督式优化，省掉在线采样。简单稳定，代价是只能利用已有的偏好数据，没有探索。 GRPO（arXiv:2402.03300，2024-02）出自 DeepSeekMath，本是给语言模型数学推理用的。用组内相对优势替掉 critic，把 RLHF 的工程复杂度砍掉一大块。 Flow-GRPO（arXiv:2505.05470，2025-05）第一个把在线 policy gradient 接进 flow matching 的工作，靠 ODE→SDE 转换取得随机逐步转移的可计算似然（3.2 节）。在文生图上把 GenEval 从 63% 拉到 95%、视觉文字渲染从 59% 到 92%。 需要说明：Flow-GRPO 的实验是在文生图上做的（SD3.5-M），不是视频。它之所以是视频 RL 的关键前置，是因为主流视频模型也是 flow matching 训的，ODE→SDE 那套推导可以直接搬过来——但视频上的实证还远没有图像充分。 08. 常见误解 「奖励涨了就是训好了」。 01 节整节都在讲这件事。奖励是你自己写的目标函数，模型最大化它是本分。奖励上升唯一证明的是优化在工作，不是效果在变好。 「reward hacking 是 RL 算法的问题，换个算法就好了」。04.4 节那个对照实验就是为了拆掉这个想法：同一个 GRPO，只改奖励权重，行为从 $m = 0.966$ 翻到 $m = 0.008$。算法只是执行者。出了 hacking 该去审奖励函数。 「GRPO 比 PPO 好」。 它主要是更省——不用训 critic，少一个网络和一堆超参。省掉 critic 的代价是必须一次采一组样本才能估出基线，采样开销更大。在视频这种采样极贵的场景里，这个权衡并不是无脑赚。 「多加几个奖励维度就能防 hacking」。 加维度确实提高了套利难度，但只要还是线性加权，套利空间就存在——4.2 节那个例子里有 29 个维度，静态依然套出了 1.1 的净收益。真正起作用的是权重结构（指令遵循占 37.4% 才压住了它），以及 DPOK 那种把策略拴在原始分布附近的 KL 护栏。 「flow matching 模型可以直接套用扩散模型的 RL 方法」。按随机转移似然做 PPO/GRPO 时，需要先定义非退化随机转移，例如 3.2 节的 ODE→SDE 路线。确定性 ODE 的输出边缘密度、可微奖励路径导数等另有方法，并不是所有 RL 都必须给每一步加噪。 09. 动手验证 三个实验，用到的脚本全文都在文末附录，纯 NumPy，没有 GPU 依赖，复制存成同名文件就能跑。 实验一：给自己的奖励函数算套利空间。 跑 reward_hacking.py，看第 2 段那笔账。然后把你自己项目里奖励各维度的权重填进 DIMS，重算一遍「静止能白拿 / 要放弃」的净收益。预期结论是：只要净收益为正，你的模型早晚会发现它。 实验二：找出翻转的临界点。 跑 grpo_minimal.py，把指令遵循的三个权重乘一个系数 $c$（改 weight_without_instruction 为按比例缩放），从 $c = 1.0$ 逐步降到 $0$，看收敛的 $m$ 在哪一步从 0.97 掉到 0.01。预期是个比较陡的转折而不是平滑过渡——这意味着奖励权重的微小改动可能导致行为质变，调权重时值得扫一遍而不是只试一个值。 实验三：验证组大小对方差的影响。 还是 grpo_minimal.py，把 group 从 32 改成 4、8、64，各跑三个不同 seed。预期看到组越小、收敛点越不稳定——因为组内均值和标准差是用这几个样本估的，样本太少基线就不可靠。还要一起权衡采样成本、同分组概率、奖励稀疏度和并行吞吐。 本文三张配图由 make_figures.py 生成（这个需要额外装 pip install matplotlib，其余脚本只要 NumPy）。想复现就把权重换成你自己的：左边的权重账本换成你的 29 维，右边的曲面和 GRPO 轨迹会直接告诉你你的奖励函数在纵容哪种躺平。 10. 延伸阅读 这篇在知识树里是 L4（前沿专题），前面还有三级台阶。如果 03 节读起来吃力，按下面顺序补更省时间： 策略梯度与 PPO 基础——3.1 节那个对数导数技巧、为什么要减基线、clip 是干什么的，都属于那一篇 从 DPO 到 GRPO：去掉价值网络——GRPO 本身的完整讨论在那篇。本文只用到它「组内相对优势」这一个机制，没有展开和 PPO、DPO 的对比 把 RL 用到扩散模型上——DDPO 那套「把去噪当多步决策」的建模 流匹配与 Rectified Flow——3.2 节 ODE→SDE 转换的前提 读完这篇可以继续看： 人类偏好数据怎么收（还没写）——4.1 节那 29 个权重是怎么从成对比较数据里拟合出来的 视频质量评测指标（还没写）——06 节说的那三个检测手段该怎么落实，尤其是「奖励模型看不见的指标」具体该记哪些 附录：完整代码 09 节用到的脚本全文如下（make_figures.py、reward_hacking.py、grpo_minimal.py）。复制到本地存成同名文件，按各脚本开头的依赖说明准备环境后即可运行。 make_figures.py #!/usr/bin/env python3 # -*- coding: utf-8 -*- """生成「视频生成中的强化学习与奖励模型」的三张解释图。 数值一律从同目录两个脚本里取（reward_hacking 的 29 个真实权重、 grpo_minimal 的奖励曲面与 GRPO 轨迹），这里只负责画——那边改了这里要重跑， 免得图和正文数字打架。 三张图分别回答： 1. 29 个奖励权重的账本长什么样（为什么「有没有照做」压过「好不好看」） 2. 组内优势到底把什么信息抽出来了（原始奖励看着都一样，标准化后才有正负） 3. 只改奖励权重、不动算法，GRPO 为什么会从「照做」翻转成「躺平」 只依赖 numpy + matplotlib。跑法：python make_figures.py """ import textwrap from pathlib import Path import matplotlib.pyplot as plt import numpy as np import grpo_minimal as G from reward_hacking import NAMES, WEIGHT ROOT = Path(__file__).resolve().parents[1] OUT = ROOT / "figures" OUT.mkdir(parents=True, exist_ok=True) plt.rcParams.update({ "font.sans-serif": ["PingFang SC", "Arial Unicode MS", "DejaVu Sans"], "axes.unicode_minus": False, "figure.dpi": 160, "savefig.bbox": "tight", }) INK = "#1f2937" MUTE = "#6b7280" C_INSTR = "#d1495b" # 指令遵循：红 C_OTHER = "#2f6fb0" # 其余维度：蓝 C_AES = "#c9ced6" # 构图/色彩：灰 C_FULL = "#2f6fb0" # 完整权重 C_NOIN = "#d1495b" # 去掉指令遵循 C_POS = "#2f9e6f" # 优势为正 C_NEG = "#d1495b" # 优势为负 # 维度分组（下标对应 reward_hacking.DIMS 的顺序） GROUPS = [ ("指令遵循", [0, 1, 2], C_INSTR), ("构图＋色彩", [3, 4, 6], C_AES), ("光照", [7, 8, 9, 10, 11, 12], C_OTHER), ("形状保持", [13, 14, 15, 16, 17], C_OTHER), ("运动", [5, 18, 19, 20, 21], C_OTHER), ("画质＋细节", [22, 23, 24, 25], C_OTHER), ("文字＋内容", [26, 27, 28], C_OTHER), ] def style(ax, title, xlabel=None, ylabel=None): ax.set_title(title, fontsize=11.5, color=INK, pad=10, loc="left") if xlabel: ax.set_xlabel(xlabel, fontsize=10, color=MUTE) if ylabel: ax.set_ylabel(ylabel, fontsize=10, color=MUTE) ax.tick_params(colors=MUTE, labelsize=9) for s in ("top", "right"): ax.spines[s].set_visible(False) for s in ("left", "bottom"): ax.spines[s].set_color("#d1d5db") ax.grid(alpha=0.25, linewidth=0.6, axis="y") ax.set_axisbelow(True) def footer(fig, text, width=118): """把「这张图要看什么」放到坐标轴下方。 必须放在 y<0：bbox_inches="tight" 会把负坐标的 artist 一起收进来， 放在 0~0.05 之间会和 x 轴标签叠在一起。 """ wrapped = "\n".join(textwrap.wrap(text, width=width)) fig.text(0.012, -0.13, wrapped, fontsize=8.5, color=MUTE, va="top", ha="left", linespacing=1.6) # ────────────────────────────────────────────────────────────────────── # 图 1：29 个奖励权重的账本 # ────────────────────────────────────────────────────────────────────── def fig_weight_ledger(): order = np.argsort(WEIGHT) # 从小到大，画出来是升序 colors = [C_INSTR if i in (0, 1, 2) else C_AES if i in (3, 4, 6) else C_OTHER for i in order] labels = [NAMES[i] for i in order] fig, axes = plt.subplots(1, 2, figsize=(12.6, 6.4), gridspec_kw={"width_ratios": [1.45, 1]}) fig.subplots_adjust(wspace=0.34) ax = axes[0] ax.barh(np.arange(29), WEIGHT[order], color=colors, height=0.72) ax.set_yticks(np.arange(29)) ax.set_yticklabels(labels, fontsize=7.6) ax.set_xlim(0, WEIGHT.max() * 1.18) ax.tick_params(axis="y", length=0) style(ax, "29 个维度的权重（升序）", "权重 $w_k$", None) # 最高 / 最低标注 ax.text(WEIGHT[order][-1] + 0.02, 28, f"{WEIGHT.max():.4f} 最高", fontsize=8.5, color=INK, va="center") ax.text(WEIGHT[order][0] + 0.02, 0, f"{WEIGHT.min():.4f} 最低", fontsize=8.5, color=MUTE, va="center") # 图例：红=指令遵循，灰=构图色彩，蓝=其余 from matplotlib.patches import Patch ax.legend(handles=[ Patch(color=C_INSTR, label="指令遵循（3 维）"), Patch(color=C_AES, label="构图＋色彩（3 维）"), Patch(color=C_OTHER, label="其余（23 维）"), ], fontsize=8.2, frameon=False, loc="lower right") ax = axes[1] gcolor = {name: c for name, _, c in GROUPS} gs = [(name, float(WEIGHT[idx].sum())) for name, idx, _ in GROUPS] gs.sort(key=lambda t: -t[1]) total = WEIGHT.sum() ys = np.arange(len(gs)) # 颜色必须按排序后的名字取：按 GROUPS 顺序 zip 会跟排序后的柱子错位 ax.barh(ys, [v for _, v in gs], color=[gcolor[n] for n, _ in gs], height=0.7) ax.set_yticks(ys) ax.set_yticklabels([n for n, _ in gs], fontsize=9) ax.tick_params(axis="y", length=0) for i, (_, v) in enumerate(gs): ax.text(v + 0.03, i, f"{v:.3f} ({v / total * 100:.1f}%)", fontsize=8.2, color=INK, va="center") ax.set_xlim(0, max(v for _, v in gs) * 1.45) style(ax, "按语义归堆：谁在主导总分", "权重合计", None) fig.suptitle("图 1　VisionReward 的 29 个权重不是均衡的：" "指令遵循一家占了 37.4%，构图加色彩只有 1.5%", fontsize=12, color=INK, x=0.012, ha="left", y=1.02) footer(fig, "要看什么：左图是 29 个权重从低到高排开——最高的「指令-未完全失败」1.1418，" "最低的「细节-不粗糙」0.0085，相差 134 倍，根本不是「各维度都重要」的均衡奖励。" "右图把它们按语义归堆后更刺眼：指令遵循三档合计 2.3486，一家吃掉 37.4%；" "而构图加色彩只有 0.0924、占 1.5%。真实人类偏好数据里，" "「有没有照着 prompt 做」比「好不好看」重要二十多倍——这个结构是后面所有套利的源头。") fig.savefig(OUT / "reward_weight_ledger.png") plt.close(fig) print(f"[图1] 29 维权重合计 {total:.4f}；最高 {WEIGHT.max():.4f}／最低 {WEIGHT.min():.4f}" f" = {WEIGHT.max() / WEIGHT.min():.0f} 倍；指令遵循占 " f"{WEIGHT[[0, 1, 2]].sum() / total * 100:.1f}%") # ────────────────────────────────────────────────────────────────────── # 图 2：组内优势把「谁比谁好」抽出来 # ────────────────────────────────────────────────────────────────────── def fig_group_advantage(): rng = np.random.default_rng(2024) rs = 8.0 + rng.random(8) * 0.2 # 同一个 prompt 采 8 条，分都挤在 8.0~8.2 adv = (rs - rs.mean()) / (rs.std() + 1e-8) idx = np.arange(1, 9) fig, axes = plt.subplots(1, 2, figsize=(12.6, 4.6)) fig.subplots_adjust(wspace=0.26) ax = axes[0] ax.bar(idx, rs, color=C_OTHER, width=0.62) ax.set_ylim(0, 8.9) ax.set_xticks(idx) style(ax, "原始奖励：8 条都在 8.0~8.2 之间", "组内第 i 条", "奖励 $r_i$") for i, v in zip(idx, rs): ax.arrow(i, v + 0.42, 0, 0.30, width=0.055, color=INK, length_includes_head=True, head_width=0.16, head_length=0.12) ax.text(0.99, 8.72, "每一条都在说：往我这边推（只是力度略有差别）", fontsize=8.6, color=MUTE, ha="right", va="top") ax.text(0.02, 0.30, f"极差只有 {rs.max() - rs.min():.2f}", fontsize=8.6, color=INK, transform=ax.transAxes) ax = axes[1] cols = [C_POS if a >= 0 else C_NEG for a in adv] ax.bar(idx, adv, color=cols, width=0.62) ax.axhline(0, color=INK, lw=1.0) ax.set_xticks(idx) ax.set_ylim(-2.0, 2.0) for i, v in zip(idx, adv): ax.text(i, v + (0.09 if v >= 0 else -0.09), f"{v:+.2f}", ha="center", va="bottom" if v >= 0 else "top", fontsize=8, color=INK) style(ax, "减均值除标准差后：谁该推、谁该压，一眼分清", "组内第 i 条", "优势 $A_i$") ax.text(0.98, 0.94, "绿的往上推，红的往下压；组内均值自动当基线", fontsize=8.6, color=MUTE, ha="right", va="top", transform=ax.transAxes) fig.suptitle("图 2　GRPO 的组内基线：把「绝对分多高」换成「谁比谁好」", fontsize=12, color=INK, x=0.012, ha="left", y=1.04) footer(fig, "要看什么：左图是同一个 prompt 采出的 8 条视频的奖励，全挤在 8.0 到 8.2 之间——" "如果直接把 $r_i$ 乘在梯度上，8 条样本会一起喊「往我这边推」，只是力气略有差别，" "有用的信息被绝对值淹没了（箭头方向完全一致）。右图减掉组内均值、除以组内标准差之后，" "同样的 8 个数字变成了有正有负的优势：一半往上推、一半往下压。" "这就是 GRPO 不用 critic 也能训的原因——基线是从同组样本里白捡的。") fig.savefig(OUT / "group_advantage.png") plt.close(fig) print(f"[图2] 8 条奖励 {rs.min():.3f}~{rs.max():.3f}（极差 {rs.max() - rs.min():.3f}）；" f"优势 {adv.min():+.2f}~{adv.max():+.2f}") # ────────────────────────────────────────────────────────────────────── # 图 3：只改奖励权重，同一个 GRPO 翻转 # ────────────────────────────────────────────────────────────────────── def _trace_grpo(weight=None, seed=0, steps=400, group=32, lr=0.35, sigma=0.15): """复刻 grpo_minimal.grpo 的循环，只是把每步的 m 和组内平均奖励记下来。 抽随机的顺序必须和原函数完全一致，否则数字对不上正文。 """ rng = np.random.default_rng(seed) theta = 0.0 ms, rr = [], [] for _ in range(steps): mu = 1 / (1 + np.exp(-theta)) eps = rng.normal(0, sigma, group) m_clip = np.clip(mu + eps, 0.0, 1.0) r = np.array([G.reward(m, rng, weight) for m in m_clip]) adv = (r - r.mean()) / (r.std() + 1e-8) grad = float(np.mean(adv * eps / sigma ** 2) * mu * (1 - mu)) theta += lr * grad ms.append(mu) rr.append(float(r.mean())) return np.array(ms), np.array(rr), 1 / (1 + np.exp(-theta)) def fig_surface_flip(): w_no = G.weight_without_instruction() # 奖励曲面：与 grpo_minimal.main 同口径（rng=7，每点 4000 次采样） rng = np.random.default_rng(7) grid = np.linspace(0, 1, 21) surf_full, surf_no = [], [] for m in grid: surf_full.append(float(np.mean([G.reward(m, rng) for _ in range(4000)]))) surf_no.append(float(np.mean([G.reward(m, rng, w_no) for _ in range(4000)]))) surf_full = np.array(surf_full) surf_no = np.array(surf_no) # 正文引用的那张 11 点表，峰值标注以它为准 rng2 = np.random.default_rng(7) tbl = np.linspace(0, 1, 11) tf, tn = [], [] for m in tbl: # 抽随机顺序与 grpo_minimal.main 完全一致（逐点交错），否则峰值数字对不上正文 tf.append(float(np.mean([G.reward(m, rng2) for _ in range(4000)]))) tn.append(float(np.mean([G.reward(m, rng2, w_no) for _ in range(4000)]))) tf, tn = np.array(tf), np.array(tn) pk_f, pk_n = tbl[int(tf.argmax())], tbl[int(tn.argmax())] ms_full, rr_full, m_end_full = _trace_grpo(weight=None, seed=0) ms_no, rr_no, m_end_no = _trace_grpo(weight=w_no, seed=0) fig, axes = plt.subplots(1, 2, figsize=(12.6, 5.0)) fig.subplots_adjust(wspace=0.30) ax = axes[0] ax.plot(grid, surf_full, lw=2.4, color=C_FULL, label="完整 29 维权重") ax.plot(grid, surf_no, lw=2.4, color=C_NOIN, ls="--", label="去掉指令遵循三档") ax.axhline(0, color="#9ca3af", lw=0.8) ax.scatter([pk_f], [tf.max()], s=70, color=C_FULL, zorder=5) ax.scatter([pk_n], [tn.max()], s=70, color=C_NOIN, zorder=5) ax.annotate(f"最高点 m={pk_f:.1f}", (pk_f, tf.max()), textcoords="offset points", xytext=(10, -30), fontsize=8.6, color=C_FULL, arrowprops=dict(arrowstyle="->", color=C_FULL, lw=1.0)) ax.annotate(f"最高点 m={pk_n:.1f}", (pk_n, tn.max()), textcoords="offset points", xytext=(18, 26), fontsize=8.6, color=C_NOIN, arrowprops=dict(arrowstyle="->", color=C_NOIN, lw=1.0)) ax.legend(fontsize=8.6, frameon=False, loc="lower center") style(ax, "同一组物理约束，两种权重下的奖励曲面形状完全相反", "运动强度 $m$", "期望奖励 $r$") ax = axes[1] ax.plot(ms_full, lw=2.4, color=C_FULL, label="m（完整权重）") ax.plot(ms_no, lw=2.4, color=C_NOIN, ls="--", label="m（去掉指令遵循）") ax.set_ylim(-0.05, 1.05) ax.set_xlabel("训练步", fontsize=10, color=MUTE) ax.set_ylabel("运动强度 $m$", fontsize=10, color=MUTE) ax2 = ax.twinx() ax2.plot(rr_full, lw=1.1, color=C_FULL, alpha=0.45) ax2.plot(rr_no, lw=1.1, color=C_NOIN, alpha=0.45, ls="--") ax2.set_ylabel("组内平均奖励（浅色细线）", fontsize=9, color=MUTE) ax2.tick_params(colors=MUTE, labelsize=8.5) ax2.spines["top"].set_visible(False) h1, l1 = ax.get_legend_handles_labels() ax.legend(h1, l1, fontsize=8.6, frameon=False, loc="center right") style(ax, "同一个 GRPO 跑 400 步：一条冲到 0.97，一条躺到 0.01", None, None) ax.set_ylabel("运动强度 $m$", fontsize=10, color=MUTE) ax.set_xlabel("训练步", fontsize=10, color=MUTE) ax.text(0.02, 0.06, f"完整权重收敛 m={m_end_full:.3f}\n" f"去掉指令遵循收敛 m={m_end_no:.3f}", fontsize=8.4, color=INK, transform=ax.transAxes, va="bottom", bbox=dict(fc="white", ec="#d1d5db", lw=0.7, pad=4)) fig.suptitle("图 3　奖励曲线一样漂亮，学到的东西完全相反", fontsize=12, color=INK, x=0.012, ha="left", y=1.03) footer(fig, "要看什么：左图是同一组物理约束下的奖励曲面。完整权重（蓝实线）在 m=0.8 附近见顶，" "去掉指令遵循三档后（红虚线）整条曲线翻转成单调递减，最高点直接落到 m=0——彻底不动才是最优解。" "右图是同一个 GRPO、同一个初值跑 400 步：蓝线冲到 0.966 选择照做 prompt，" "红线掉到 0.008 选择躺平。而两条浅色细线是各自的组内平均奖励——它们都在稳步上涨，" "形状和一次成功的训练没有任何区别。奖励曲线不能告诉你模型学到了什么。") fig.savefig(OUT / "reward_surface_flip.png") plt.close(fig) print(f"[图3] 曲面峰值：完整 m={pk_f:.1f} ({tf.max():+.4f})，" f"去掉指令 m={pk_n:.1f} ({tn.max():+.4f})；" f"GRPO 收敛 {m_end_full:.4f} vs {m_end_no:.4f}") def main(): fig_weight_ledger() fig_group_advantage() fig_surface_flip() print(f"[OK] 三张图已写入 {OUT}") for f in sorted(OUT.glob("*.png")): print(f" {f.name} {f.stat().st_size / 1024:.0f} KB") if __name__ == "__main__": main() reward_hacking.py """用 VisionReward 的真实权重，算一遍「静态视频能不能骗到高分」。 权重与问题取自 THUDM/VisionReward（2026-08 的 main 分支）： VisionReward_Video/weight.json 29 个线性权重 VisionReward_Video/VisionReward_video_qa_select.txt 29 个 yes/no 问题 打分公式同 inference-video.py 的 score()：yes=+1，no=-1，乘权重后取均值。 注意：下面两个视频的 yes/no 画像是**按各维度语义做的人工假设**， 不是真实模型的输出。它用来说明权重结构本身允许什么样的套利， 不构成对 VisionReward 实际表现的评测。 运行 `python reward_hacking.py`。 """ import numpy as np # (权重, 维度简称) —— 顺序与上游 weight.json 严格一致 DIMS = [ (0.9544, "指令-全部满足"), (0.2524, "指令-大部分满足"), (1.1418, "指令-未完全失败"), (0.0350, "构图-美观"), (0.0252, "构图-无明显缺陷"), (0.1260, "镜头运动-无明显缺陷"), (0.0322, "色彩-不难看"), (0.1629, "光照-完全准确"), (0.2167, "光照-无明显错误"), (0.0197, "光照-存在"), (0.1360, "光照-极美"), (0.0965, "光照-美"), (0.1549, "光照-不难看"), (0.1294, "首帧形状-完全准确"), (0.0989, "首帧形状-无明显错误"), (0.1884, "首帧形状-不混乱"), (0.1844, "全程形状-完美保持"), (0.2636, "全程形状-不混乱"), (0.1117, "镜头运动-高度动态"), (0.0517, "镜头运动-不微弱"), (0.0256, "物体运动-非常平滑"), (0.4390, "物体运动-完全真实"), (0.2686, "画质-非常稳定"), (0.4293, "细节-非常精细"), (0.0085, "细节-不粗糙"), (0.1276, "细节-不显著粗糙"), (0.0580, "文字-全部正确"), (0.1446, "文字-存在"), (0.3942, "内容-属于物理世界"), ] WEIGHT = np.array([w for w, _ in DIMS]) NAMES = [n for _, n in DIMS] Y, N = 1, -1 # 视频 A：几乎静止的精致画面。prompt 要求「人物走过街道」，它只给了个站着的人。 STATIC = [ N, N, Y, # 指令：没走，但也不算完全没关系 Y, Y, # 构图：精心构图 Y, # 镜头运动无缺陷（没动，自然没缺陷） Y, # 色彩 Y, Y, Y, Y, Y, Y, # 光照：静态画面容易做好 Y, Y, Y, # 首帧形状：完美 Y, Y, # 全程形状：不动＝完美保持 N, N, # 镜头动态：没有 Y, Y, # 物体运动：平滑（无运动可挑）、真实（无违反物理） Y, # 画质稳定：不动＝极稳 Y, Y, Y, # 细节：静态可以渲染很精细 Y, Y, # 文字 Y, # 属于物理世界 ] # 视频 B：真的走起来了，但运动带来形变、抖动和轻微模糊。 DYNAMIC = [ Y, Y, Y, # 指令：确实照做了 N, Y, # 构图：运动中构图一般 Y, # 镜头运动无明显缺陷 Y, # 色彩 N, Y, Y, N, Y, Y, # 光照：运动中光照没那么完美 Y, Y, Y, # 首帧形状：还行 N, N, # 全程形状：走动中有形变 ← 运动的代价 Y, Y, # 镜头动态：有 Y, Y, # 物体运动：平滑且真实 N, # 画质稳定：运动带来抖动 ← 运动的代价 N, Y, Y, # 细节：动态模糊，不够精细 ← 运动的代价 Y, Y, # 文字 Y, # 属于物理世界 ] def score(answers): """复刻 inference-video.py 的 score()。""" a = np.array(answers) assert len(a) == len(WEIGHT), f"维度不匹配: {len(a)} vs {len(WEIGHT)}" return float(np.mean(a * WEIGHT)) def main(): print("=== 1. 权重结构 ===") print(f" 29 个维度，权重总和 {WEIGHT.sum():.4f}，均值 {WEIGHT.mean():.4f}") print(f" 最高 {WEIGHT.max():.4f}（{NAMES[int(WEIGHT.argmax())]}）") print(f" 最低 {WEIGHT.min():.4f}（{NAMES[int(WEIGHT.argmin())]}）") print(f" 最高／最低 = {WEIGHT.max() / WEIGHT.min():.0f} 倍") instr = WEIGHT[0] + WEIGHT[1] + WEIGHT[2] print(f"\n 指令遵循三档合计 {instr:.4f}，占总权重 {instr / WEIGHT.sum() * 100:.1f}%") aes = WEIGHT[3] + WEIGHT[4] + WEIGHT[6] print(f" 构图＋色彩合计 {aes:.4f}，占总权重 {aes / WEIGHT.sum() * 100:.1f}%") print("\n=== 2. 静态套利的账 ===") shape_stable = WEIGHT[[16, 17, 22]].sum() motion = WEIGHT[[18, 19]].sum() print(f" 静止能白拿：全程形状保持＋不混乱＋画质稳定 = {shape_stable:.4f}") print(f" 静止要放弃：镜头高度动态＋不微弱 = {motion:.4f}") print(f" 净收益 {shape_stable - motion:+.4f}" f" （yes/no 相差 2 个单位，实际影响 ×2 = {(shape_stable - motion) * 2:+.4f}）") print("\n=== 3. 两个视频的实际得分 ===") s_static, s_dynamic = score(STATIC), score(DYNAMIC) print(f" 视频 A（几乎静止的精致画面，没照做 prompt） {s_static:+.4f}") print(f" 视频 B（照做了 prompt，但有形变和抖动） {s_dynamic:+.4f}") winner = "A（静态）" if s_static > s_dynamic else "B（动态）" print(f" 奖励模型更偏好：{winner} 差值 {abs(s_static - s_dynamic):.4f}") print("\n=== 4. 差异来自哪几个维度 ===") diff = (np.array(STATIC) - np.array(DYNAMIC)) * WEIGHT order = np.argsort(-np.abs(diff)) print(" (正值＝这一维让静态视频占便宜)") for i in order[:8]: if diff[i] != 0: print(f" {diff[i]:+.4f} {NAMES[i]}") if __name__ == "__main__": main() grpo_minimal.py """GRPO 的最小可运行实现，跑在 VisionReward 的真实权重上。 策略只有一个参数：运动强度 m。这是刻意的——把维度压到 1 维， 才能看清 GRPO 到底把策略推向哪里，而不是被高维噪声淹没。 人为设定的玩具权衡（不是视频物理定律或真实 VLM 的测量）： m 越大 → 指令遵循、动态性维度更容易拿 yes m 越大 → 形状保持、画质稳定、细节精细更容易拿 no（运动带来形变/抖动/模糊） prompt 明确要求运动（"a person walking across the street"）， 所以「照做」在语义上就等于 m 要大。看 GRPO 最后选了什么。 运行 `python grpo_minimal.py`。 """ import numpy as np from reward_hacking import WEIGHT, NAMES # 各维度对运动强度 m 的敏感方向 # +1: m 越大越容易 yes -1: m 越大越容易 no 0: 与 m 无关 SENSITIVITY = np.zeros(29) SENSITIVITY[[0, 1, 2]] = +1.0 # 指令遵循（prompt 要求运动） SENSITIVITY[[18, 19]] = +1.0 # 镜头动态 SENSITIVITY[[20, 21]] = +0.3 # 运动平滑/真实：要有运动才谈得上，但过猛会假 SENSITIVITY[[16, 17]] = -1.0 # 全程形状保持 SENSITIVITY[22] = -1.0 # 画质稳定 SENSITIVITY[[23, 24, 25]] = -0.7 # 细节精细（动态模糊） SENSITIVITY[[7, 8, 10, 11]] = -0.3 # 光照准确/美观：运动中更难保持 def yes_prob(m): """给定运动强度 m∈[0,1]，各维度答 yes 的概率。 基线 0.5 表示「说不准」，敏感度把它往两边拉。 """ return np.clip(0.5 + SENSITIVITY * (m - 0.5) * 1.6, 0.02, 0.98) def sample_video(m, rng): """按概率采一个 yes/no 画像，yes=+1 no=-1。""" return np.where(rng.random(29) < yes_prob(m), 1, -1) # 指令遵循的三个维度。把它们置零就得到一个「只看画质」的奖励模型， # 用来对照权重结构对 reward hacking 的影响。 INSTRUCTION_DIMS = [0, 1, 2] def reward(m, rng, weight=None): """复刻 inference-video.py 的 score()：np.mean(answers * weight)。""" w = WEIGHT if weight is None else weight return float(np.mean(sample_video(m, rng) * w)) def weight_without_instruction(): w = WEIGHT.copy() w[INSTRUCTION_DIMS] = 0.0 return w def grpo(steps=400, group=32, lr=0.35, sigma=0.15, theta0=0.0, seed=0, weight=None, verbose=False): """单参数 GRPO，返回收敛后的运动强度 m。 每步采 group 个动作，用组内均值/标准差做 baseline 得到 advantage， 再按 advantage 加权更新——不需要 critic，这是 GRPO 的核心简化。 为了看清裸的优化行为，这里省掉了 PPO 式的 clip 和对参考模型的 KL 惩罚。 """ rng = np.random.default_rng(seed) theta = theta0 for step in range(steps): # 策略：m = clip(sigmoid(theta) + noise)，高斯探索 mu = 1 / (1 + np.exp(-theta)) eps = rng.normal(0, sigma, group) ms = np.clip(mu + eps, 0.0, 1.0) rs = np.array([reward(m, rng, weight) for m in ms]) # 组内相对优势：减均值除标准差 adv = (rs - rs.mean()) / (rs.std() + 1e-8) # 潜在动作 a = mu + eps 的高斯 score，m=clip(a) 只是环境映射。 # 这不是含边界原子质量的裁剪后 m 分布的逐点 log-density。 # ∂mu/∂theta = mu(1-mu) grad = float(np.mean(adv * eps / sigma ** 2) * mu * (1 - mu)) theta += lr * grad if verbose and (step % 80 == 0 or step == steps - 1): print(f" {step:>4} {mu:.4f} {float(rs.mean()):+.4f}") return 1 / (1 + np.exp(-theta)) def main(): print("=== 1. 优化前后的运动强度 ===") print(" step 运动强度 m 组内平均奖励") m_final = grpo(verbose=True) print(f"\n GRPO 收敛到 m = {m_final:.4f}") verdict = "静态（放弃了 prompt 要求的运动）" if m_final < 0.35 else \ "动态（照做了 prompt）" if m_final > 0.65 else "折中" print(f" → 策略选择：{verdict}") print("\n=== 2. 从不同初值出发是否都收敛到同一处 ===") for name, t0 in (("偏静态 (m≈0.27)", -1.0), ("中性 (m=0.50)", 0.0), ("偏动态 (m≈0.73)", 1.0)): mf = grpo(theta0=t0, seed=1) print(f" 初值 {name:<18} → 收敛 m = {mf:.4f}") print("\n=== 3. 奖励曲面：各运动强度的期望得分 ===") rng = np.random.default_rng(7) w_no_instr = weight_without_instruction() print(" m 完整权重 去掉指令遵循 (每点 4000 次采样)") best, best_no = (-9e9, None), (-9e9, None) for m in np.linspace(0, 1, 11): r = float(np.mean([reward(m, rng) for _ in range(4000)])) r2 = float(np.mean([reward(m, rng, w_no_instr) for _ in range(4000)])) best = max(best, (r, m)) best_no = max(best_no, (r2, m)) print(f" {m:.1f} {r:+.4f} {r2:+.4f}") print(f" → 完整权重的最高点在 m = {best[1]:.1f}（{best[0]:+.4f}）") print(f" → 去掉指令遵循后最高点在 m = {best_no[1]:.1f}（{best_no[0]:+.4f}）") print("\n=== 4. 对照：奖励模型不看指令遵循时，GRPO 选什么 ===") print(" step 运动强度 m 组内平均奖励") m_hack = grpo(weight=w_no_instr, seed=0, verbose=True) print(f"\n 收敛到 m = {m_hack:.4f}" f"（完整权重下是 {m_final:.4f}）") print(" → 同一个算法、同一个物理约束，只改奖励的权重结构，" "\n 策略就从「照做」翻转成「静止不动」。这就是 reward hacking 的来源：" "\n 不是算法坏了，是它精确地最大化了你真正写下来的那个目标。") print("\n=== 5. 两个收敛点在各维度上的差别 ===") p_ok, p_hack = yes_prob(m_final), yes_prob(m_hack) gap = (p_hack - p_ok) * WEIGHT order = np.argsort(np.abs(gap))[::-1] print(" (正值＝hack 后的策略在这一维更容易拿 yes)") for i in order[:6]: print(f" {gap[i]:+.4f} {NAMES[i]}" f" (yes 概率 {p_ok[i]:.2f} → {p_hack[i]:.2f})") if __name__ == "__main__": main()

AIGC 基本功|旋转位置编码 RoPE 的原理与实现-RoPE 旋转位置编码 RoPE 的原理与实现 所属方向：注意力与核心零件 ｜ 难度：核心组件 ｜ 前置知识：MHA 关键词：位置编码、RoPE、旋转位置编码、相对位置、复数旋转、长度外推 01. 为什么需要它 先说一个具体的失败场景。 你训了一个视频 DiT，训练时用 480p、49 帧。推理时想出 720p、97 帧，结果画面从第 50 帧开始漂——人物五官慢慢挪位，背景纹理开始流动。你换了采样器、加了 CFG、调了 shift，都没用。 先把这次改动的量级算清楚。按常见配置（VAE 空间 8 倍、时间 4 倍压缩，再做 2×2 patchify）： 480p/49 帧 → 13 个 latent 帧 × 1560 tokens/帧 = 20280 tokens 720p/97 帧 → 25 个 latent 帧 × 3600 tokens/帧 = 90000 tokens 总量涨了 4.44 倍，其中空间方向 2.31 倍、时间方向 1.92 倍。也就是说，模型要给四倍多的位置编号赋予意义，新增的较远坐标与注意力关系超出了训练范围（已有坐标仍然见过）。 这是一个待诊断的假设场景：位置编码、训练尺度、VAE、注意力规模和运动分布都可能造成失败，不能仅凭症状归因给 RoPE。 注意力机制本身是置换等变的：把输入序列的顺序打乱，输出也只是跟着打乱，每个 token 看到的上下文完全不变。换句话说，纯注意力分不清「第 3 帧」和「第 40 帧」。位置信息必须额外注入。 早期做法是可学习的绝对位置嵌入（learned absolute PE）：给每个位置分配一个可训练向量，加到 token 上。这在训练长度内工作良好，但有个硬伤——没见过的位置没有对应向量。你训练时只见过 49 个位置，推理要 97 个，后面 48 个位置的嵌入是凭空来的，模型自然就崩了。 RoPE 解决的正是这件事：它不给位置分配「向量」，而是给位置分配「旋转角度」。角度公式可计算到新位置；生成质量能否外推仍需训练或扩展策略验证。 02. 最小可用理解 三句话： 把 query 和 key 的特征维度两两配对看成复数，每个位置 $m$ 把它们旋转 $m\theta$ 弧度。 两个向量做内积时，指数项相减，绝对位置 $m$、$n$ 自动消掉，只剩相对距离 $(m-n)$。 所以对固定的 q、k，位置项只依赖 token 间的相对距离；完整分数仍依赖内容特征，而旋转角度可以外推到训练时没见过的位置。 如果只记一件事：RoPE 是乘上去的旋转，不是加上去的向量。加法会污染特征的模长，旋转不会——这是它比绝对位置嵌入更「干净」的地方（3.4 节会把这句话量化）。 还有一个位置要记准：它作用在每一层注意力的 $q$ 和 $k$ 上，不是输入层加一次，通常不旋转 $v$。这是本文及常用实现的约定，某些变体也会旋转 value，不能一概判错。 03. 数学推导 3.1 我们到底想要什么性质 先把目标写成式子。设 $x_m$ 是位置 $m$ 上的 token 特征，我们要找一个变换 $f$，使得变换后的 query 和 key 的内积满足： $$\langle f(q, m), f(k, n) \rangle = g(q, k, m-n)$$ 左边依赖两个绝对位置 $m$ 和 $n$，右边只依赖它们的差值 $m-n$。这就是「相对位置不变性」的精确表述。注意 $g$ 的具体形式我们不关心，只要求它不含 $m$、$n$ 的单独出现。 3.2 复数域上的现成答案 二维情况下有个现成的东西满足这个条件：旋转。 把二维向量 $q = (q_0, q_1)$ 看成复数 $q_0 + iq_1$。复数乘法 $e^{im\theta}$ 的几何意义就是逆时针旋转 $m\theta$ 弧度，且不改变模长。 现在令 $f(q, m) = q e^{im\theta}$，算一下内积。复数域上的内积要取共轭： $$\langle f(q,m), f(k,n) \rangle = \mathrm{Re}\left[(q e^{im\theta}) \overline{(k e^{in\theta})}\right] = \mathrm{Re}\left[q \bar{k} e^{i(m-n)\theta}\right]$$ 关键一步在指数：$e^{im\theta} \cdot \overline{e^{in\theta}} = e^{im\theta} e^{-in\theta} = e^{i(m-n)\theta}$。绝对位置就这么消掉了，只留下 $(m-n)$。这就是 RoPE 的全部动机，后面所有工程细节都只是把这个二维结论铺到高维。 3.3 铺到高维 实际的 head_dim 是 64 或 128，不是 2。做法是把 $d$ 维切成 $d/2$ 组，每组两个相邻维度构成一个复数平面，第 $i$ 组用自己的旋转频率： $$\theta_i = \text{base}^{-2i/d}, \quad i = 0, 1, \dots, d/2 - 1$$ 这里的 base 就是代码里的 rope_theta，通常取 10000。 这个频率设计是理解 RoPE 行为的关键。展开看两端： $i = 0$ 时 $\theta_0 = 1$，位置每 +1 就转 1 弧度，转得最快 $i = d/2 - 1$ 时 $\theta_i \approx \text{base}^{-1} = 10^{-4}$，位置每 +1 只转 0.0001 弧度，转得最慢 也就是说，高频维度负责区分近距离，低频维度负责区分远距离。一个位置差 1 的 token 对，主要靠前面几维的相位差来分辨；位置差较大时，高频相位已多次绕圈，而低频相位仍提供变化较慢的尺度；各频段并非互斥地负责某一距离。 这张图要看什么：左图是 $\cos(m\theta_i)$ 的热力图（横轴位置、纵轴维度组）——256 长度内转满一整圈的只有前 13 组；这些高频组提供较大的近距离相位变化，维度索引越大，相位变化越慢，但并非不携带近距离位置信息；右图把位置拉到 32768，高频组早就转了几万弧度，低频组才刚积累出可分辨的相位。这是 32 组频率提供的多尺度结构，不是互斥的距离分工；后面所有外推方法（PI、NTK、YaRN）做的都是重新分配这个分工（图由 code/make_figures.py 生成）。 这个结构直接解释了长度外推的两种主流做法：位置插值（PI）把所有位置除以一个缩放因子，等价于整体降低旋转速度；NTK-aware 缩放只调 base，让低频维度转得更慢而尽量保住高频维度的分辨率。它们改的都是同一组 $\theta_i$。 写成矩阵形式，RoPE 就是一个分块对角的旋转矩阵 $R_m$，第 $i$ 个 $2\times2$ 块为： $$\begin{pmatrix} \cos m\theta_i & -\sin m\theta_i \\ \sin m\theta_i & \cos m\theta_i \end{pmatrix}$$ 于是注意力分数变成 $q_m^\top R_{n-m} k_n$——形式上是把相对位置信息塞进了双线性形式的中间，而没有动 $q$、$k$ 本身的模长。 3.4 「不改模长」为什么要紧 02 节说过 RoPE 是乘上去的旋转而不是加上去的向量，这里补上为什么这件事重要。 $R_m$ 是正交矩阵（$R_m^\top R_m = I$），所以 $\|R_m x\| = \|x\|$——位置编码前后，token 的模长分毫不变。加法式绝对位置嵌入做不到这点：$\|x + p_m\| \neq \|x\|$，而且改变量取决于 $p_m$ 的方向，是随位置乱跳的。 这有实际后果。注意力分数里 $q$、$k$ 的模长直接决定 softmax 前的量级，而模长会影响打分量级，但不能普遍解释成 token 的重要性或置信度。加法式 PE 会把位置信息和这个量级搅在一起：同一个词放在句首和句中，模长不同，softmax 后的分布也跟着偏移。旋转则是干净的——位置只影响方向（相位），不碰长度。 实测一下（code/rope_decay.py 第 1 段）： 原始向量模长 7.315344 RoPE 后（位置 0/1/100/511） 7.315344 7.315344 7.315344 7.315344 跨位置模长极差 8.88e-16 加法式 PE 后 9.843323 10.080425 10.808198 10.190631 跨位置模长极差 0.9649 ← 同一个 token 的模长被位置改写了 RoPE 那行四个位置的模长在浮点精度内完全一致（偏差 8.88e-16 是机器 epsilon 量级）；加法式则在 9.84 到 10.81 之间晃，同一个 token 的「强度」被它所在的位置改写了将近 10%。 04. 代码实现 下面是完整可跑的最小实现。刻意用 numpy 而不是 PyTorch：讲原理时框架的抽象反而是噪声，而且 numpy 谁都能跑。 import numpy as np def build_rope_cache(seq_len: int, dim: int, base: float = 10000.0): """预计算 RoPE 的 cos/sin 表。 dim 必须是偶数：每两个相邻维度配成一个复数参与旋转。 base 就是常说的 rope_theta，调它会按维度改变旋转速度；最高频 i=0 始终为 1。 """ assert dim % 2 == 0, f"dim 必须为偶数，收到 {dim}" # theta_i = base^(-2i/dim)：i 越大转得越慢，对应越低的频率 inv_freq = 1.0 / (base ** (np.arange(0, dim, 2) / dim)) pos = np.arange(seq_len) freqs = np.outer(pos, inv_freq) # [seq_len, dim/2] return np.cos(freqs), np.sin(freqs) def apply_rope(x: np.ndarray, cos: np.ndarray, sin: np.ndarray) -> np.ndarray: """把旋转作用到 x 上。x: [seq_len, dim] 最后一维两两配对看成复数 (x_even + i*x_odd)，乘上 e^{i*m*theta}： 实部 = x_even * cos - x_odd * sin 虚部 = x_even * sin + x_odd * cos """ x_even, x_odd = x[..., 0::2], x[..., 1::2] out = np.empty_like(x) out[..., 0::2] = x_even * cos - x_odd * sin out[..., 1::2] = x_even * sin + x_odd * cos return out 验证核心性质：固定一对向量，只改赋予它们的绝对位置，间距始终为 1。如果推导是对的，内积应该完全不变。 rng = np.random.default_rng(0) seq_len, dim = 8, 64 cos, sin = build_rope_cache(seq_len, dim) q = rng.standard_normal(dim) k = rng.standard_normal(dim) for shift in range(4): qr = apply_rope(q[None, :], cos[shift:shift + 1], sin[shift:shift + 1]) kr = apply_rope(k[None, :], cos[shift + 1:shift + 2], sin[shift + 1:shift + 2]) print(f"位置 ({shift}, {shift + 1}) → 内积 {float(qr[0] @ kr[0]):.6f}") 真实输出： cos/sin 形状: (8, 32) (8, 32) 位置 1 的前 4 维 cos: [0.5403, 0.7318, 0.846, 0.9124] 位置 1 的末 2 维 cos: [1.0, 1.0] 同一对向量，平移绝对位置（间距恒为 1）: 位置 (0, 1) → 内积 -11.297869 位置 (1, 2) → 内积 -11.297869 位置 (2, 3) → 内积 -11.297869 位置 (3, 4) → 内积 -11.297869 对照组，改变间距: 间距 1 → 内积 -11.297869 间距 2 → 内积 -12.637566 间距 3 → 内积 -13.049657 未加 RoPE 的原始内积: -9.701179 三个数字值得盯一下。 第一，四行内积完全相同（-11.297869），小数点后六位都不差。这就是 3.1 节那个式子在数值上的体现：绝对位置从 (0,1) 挪到 (3,4)，分数一点没变。而间距一变（1→2→3），内积立刻变化。 第二，位置 1 的前 4 维 cos 是 0.5403、0.7318、0.846、0.9124，递增。对应 $\theta_i$ 递减：第 0 维转了整整 1 弧度（cos 1 ≈ 0.5403），第 3 维只转了约 0.42 弧度。这就是「高频维度在前」的直接观测。 第三，末 2 维 cos 打印出来是 1.0。不是真的等于 1，而是 $\theta_{31} \approx 10^{-4}$，位置 1 只转了 0.0001 弧度，$\cos(0.0001) = 0.999999995$，六位小数下显示成 1.0。这些维度在近距离上几乎不携带位置信息——它们是留给远距离的。也正因为如此，NTK-aware 缩放对这些低频维度施加更大压缩；实际质量代价仍取决于训练后的频段使用方式。 上面这张图（code/make_figures.py 图 3）把本节的两件事画在了一起。左图是平移不变性的可视化：同一对 $q$、$k$，把绝对起点放在 0 / 500 / 2000，逐点画出内积随相对距离的曲线——三条曲线完全重合（最大偏差 5.2e-13），这就是上面输出里「四行内积完全相同」的图形版。右图是另一个常被混着说的现象：距离越远，加了 RoPE 的分数与不加位置编码的原始内积的相关性在所采样的区间整体降低，但并不单调（4000 对随机向量实测，$\Delta=1$ 时 0.97、$\Delta=4096$ 时接近 0；对独立各向同性输入，虚线为 $\frac{2}{d}\sum_i\cos(\Delta\theta_i)$，d 是特征维度、Δ 是相对距离）。注意衰减的是这个相关性，不是内积本身的大小——旋转保模长，本文独立各向同性高斯向量的内积分布与距离无关（实验三会验证这一点）。 完整脚本见 code/rope_minimal.py。 4.1 顺手把两种外推方法也实现掉 3.3 节说过 PI 和 NTK-aware「改的都是同一组 $\theta_i$」。这句话用代码写出来就是两个函数，各三行——把它们并排放一起，两种方法的取舍立刻就清楚了。 def theta_original(dim=64, base=10000.0): return 1.0 / (base ** (np.arange(0, dim, 2) / dim)) def theta_pi(dim=64, base=10000.0, scale=8): """Position Interpolation：位置索引整体除以 scale。 等价于把每个 theta_i 都乘上 1/scale——所有频率被同等减速。""" return theta_original(dim, base) / scale def theta_ntk(dim=64, base=10000.0, scale=8): """NTK-aware：把 base 放大，让减速量随维度索引递增（频率越低减速越多）。""" new_base = base * scale ** (dim / (dim - 2)) return theta_original(dim, new_base) 跑 code/rope_extrapolate.py 看它们对各维度做了什么（外推 8 倍）： 维度 i 原始 PI(÷8) NTK-aware PI减速比 NTK减速比 0 1.000e+00 1.250e-01 1.000e+00 8.00x 1.00x 4 3.162e-01 3.953e-02 2.418e-01 8.00x 1.31x 8 1.000e-01 1.250e-02 5.847e-02 8.00x 1.71x 16 1.000e-02 1.250e-03 3.419e-03 8.00x 2.92x 31 1.334e-04 1.667e-05 1.667e-05 8.00x 8.00x PI 那一列减速比是齐刷刷的 8.00x，从最高频到最低频一视同仁。NTK 那一列从 1.00x 单调爬到 8.00x——最高频维度完全没动，最低频维度承担了全部的 8 倍压缩。 这张图要看什么：左图把上表画成曲线——PI 是一条平线（每维都慢 8 倍），NTK 从最高频的 1.00x 单调爬到最低频的 8.00x，两条线夹出的面积就是「NTK 少减的部分」；右图是代价，PI 把每一维的近距离分辨率都砍到 12.5%，NTK 保住高频 100% 的代价是中低频一路掉到 34.2%。两种方法改的是同一组 $\theta_i$，区别只在把这次外推的代价放在哪个频段。 代价体现在近距离分辨率上： 维度 i 原始 PI NTK PI保留 NTK保留 0 1.000000 0.125000 1.000000 12.5% 100.0% 4 0.316228 0.039528 0.241809 12.5% 76.5% 8 0.100000 0.012500 0.058472 12.5% 58.5% 16 0.010000 0.001250 0.003419 12.5% 34.2% PI 让每维相邻位置的相位增量降到 12.5%；这可能影响局部分辨，但不等于任务指标必然下降，微调与模型适应会改变实际结果。NTK 在第 0 维保住了 100%，代价是第 16 维只剩 34.2%。 还有一个容易看漏的点。NTK 外推后，最高频维度在位置 32767 处的相位是 32767 弧度，仍然远超训练时见过的 4095： 训练时见过的最大相位（位置 4095） 4095.0 rad 原始（直接外推） 位置 32767 32767.0 rad ← 超出 8.0 倍 PI 位置 32767 4095.9 rad ← 接近但略超训练端点 4095 NTK-aware 位置 32767 32767.0 rad ← 超出 8.0 倍 这是该缩放的设计取舍：最高频保留原始步长，低频逐渐压缩。未展开相位超过训练最大值不等于必然失效，相位本身具有周期性；只有最低频获得完整 8 倍压缩，不能据此保证所有频段已回到训练范围或质量安全。PI 则把所有维度的相位范围压到接近训练区间，代价是相邻位置相位步长整体下降；精确端点对齐还需使用端点跨度比。YaRN 的贡献就是把「哪些频段该压、压多少」变成一组可调策略，而不是这两种极端。 05. 工业级实现对照 参考实现（以 2026-08 的 main 分支为准，上游重构频繁）： huggingface/transformers → modeling_llama.py：LlamaRotaryEmbedding、rotate_half、apply_rotary_pos_emb huggingface/diffusers → embeddings.py：get_3d_rotary_pos_embed 生产代码和上面的最小实现有四处不一样，每一处都有理由。 5.1 配对方式不同，而且不能混用 这是最容易踩的坑。本文把第 $2i$ 和 $2i+1$ 维配成一个复数（奇偶交错），而 HF 的写法是把第 $i$ 和第 $i + d/2$ 维配成一对（前后半切分）： def rotate_half(x): """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2 :] return torch.cat((-x2, x1), dim=-1) q_embed = (q * cos) + (rotate_half(q) * sin) 配合 emb = torch.cat((freqs, freqs), dim=-1)（cos/sin 表沿最后一维复制两遍），展开就是 $x_i \cos - x_{i+d/2} \sin$ 和 $x_{i+d/2} \cos + x_i \sin$——和 3.3 节的旋转矩阵一模一样，只是谁跟谁配对换了。 那两者是等价的吗？我写了个脚本实测（code/rope_pairing.py）： === 1. 两种配对都满足相对位置不变性 === 奇偶交错: 间距恒为 3 时内积 -13.049657，四个绝对位置间的最大波动 3.55e-15 前后半切分: 间距恒为 3 时内积 -7.706177，四个绝对位置间的最大波动 1.78e-15 === 2. 同一份权重下，两种配对给出的分数不同 === 奇偶交错 -13.049657 前后半切分 -7.706177 差值 5.343480 ← 不为 0，实现和权重必须配套 === 3. 重排维度后两者完全一致 === 奇偶交错(原始排布) -13.049657 前后半切分(重排后权重) -13.049657 差值 0.00e+00 ← 只差一个固定的维度置换 三段结论连起来才是完整答案：作为位置编码方案两者等价（都满足相对位置不变性，波动 1e-15 是纯浮点误差），但对同一份权重不可互换（差了 5.34），因为它们只差一个固定的维度置换（重排后差值精确为 0）。 实践意义：移植权重时如果配对方式对不上，模型不会报错，只会变笨。这种 bug 极难查——loss 能降，输出也通顺，就是长距离能力莫名其妙地差。判断办法很简单：看它的 cos/sin 表是 cat((freqs, freqs)) 还是 repeat_interleave(freqs, 2)，前者是前后半切分，后者是奇偶交错。 5.2 inv_freq 存成 buffer，而且不进 checkpoint self.inv_freq = nn.Buffer(inv_freq, persistent=False) self.original_inv_freq = nn.Buffer(inv_freq.clone(), persistent=False) 缓存是为了省掉每步的幂运算。而 persistent=False 这个细节更值得说：它让 inv_freq 不写进 state_dict。好处是改 rope_theta 做长度外推时，不必担心 checkpoint 里存着旧频率把新配置覆盖掉——频率永远由 config 现算。留一份 original_inv_freq 是给 dynamic NTK 用的：那类方法会随序列长度临时改频率，改完得能退回原值。 5.3 接受 position_ids 而不是内部 arange def forward(self, x, position_ids): inv_freq_expanded = self.inv_freq[None, :, None].expand(position_ids.shape[0], -1, 1) position_ids_expanded = position_ids[:, None, :].float() 最小实现里位置就是 arange(seq_len)，生产上不行，至少三个场景要显式传： KV cache 增量解码：第 101 个 token 单独前向，它的位置是 100 而不是 0 序列打包（packing）：多条短样本拼成一条长序列喂进去，每条的位置要各自从 0 起算 多模态与视频：图像 token 和文本 token 的位置编号规则不同，位置得由外部算好 5.4 强制 fp32 算 cos/sin with maybe_autocast(device_type=device_type, enabled=False): freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2) emb = torch.cat((freqs, freqs), dim=-1) cos = emb.cos() * self.attention_scaling 注意它显式关掉了 autocast。原因是 $m\theta_i$ 在长序列下会长到几万，bf16 只有 8 位有效尾数，在这个量级上间隔已经大于 1——相邻位置会算出同一个角度，位置信息直接丢失。所以三角函数必须在 fp32 里算完，再转回低精度参与矩阵乘。 5.5 读陌生代码库时的四个问题 接手一个没见过的模型，想搞清它的 RoPE 怎么配的，按这个顺序问四个问题最快，也刚好对应上面四小节： cos/sin 表是怎么铺开的？ 搜 cat 还是 repeat_interleave——torch.cat((freqs, freqs)) 是前后半切分，repeat_interleave(freqs, 2) 是奇偶交错。这决定了权重能不能和别的实现互换。 inv_freq 从哪来、会不会变？ 搜 inv_freq 的赋值点。如果只在 __init__ 里算一次，就是静态 RoPE；如果 forward 里还会重算，大概是某种 dynamic NTK，那么推理时的行为会随序列长度变。 position_ids 谁给的？ 如果是内部 arange，还要核对是否加了 cache offset 或传入各段位置；从零重启才会破坏 KV cache 增量解码，多模态场景基本要改。 三角函数在什么精度下算的？ 搜附近有没有 float()、autocast、enabled=False。没有强制 fp32 的话，长序列下有精度隐患（5.4 节）。 这四个问题都能靠 grep 在几分钟内回答，比通读整个 attention 实现快得多。 5.6 视频侧：三个轴分走维度 视频 DiT 的 token 有时间、高、宽三个坐标，做法是把 head_dim 切成三份，每份跑一遍 1D RoPE 再拼起来。diffusers 里的分配是： dim_t = embed_dim // 4 # 时间轴，占 25% dim_h = embed_dim // 8 * 3 # 高，占 37.5% dim_w = embed_dim // 8 * 3 # 宽，占 37.5% 这个 1:1.5:1.5 的分配是 CogVideoX 相关实现的一种维度预算。不能从空间占 75%、时间占 25% 推断哪个轴更先失败或失败症状；需结合训练范围、每轴频率和实际消融。详见3D RoPE。 06. 代价与边界 RoPE 不是免费的。 它对固定内容向量只保证位置项的相对位置性质，不保证外推可用。 3.1 节那个性质对任意 $m$、$n$ 都成立，包括训练时没见过的位置。但「公式成立」和「效果好」是两件事：模型只在 $m\theta_i$ 落在见过的区间里学过怎么解释这些相位，超出去就是在外推一个没有监督过的函数。这也是为什么直接放长序列会崩，得配 PI 或 YaRN 这类手段。 它引入了远程衰减，而且衰减偏偏打在你最需要的地方。 这条值得展开，因为流传的说法不太准确。 「RoPE 自带远程衰减」几乎成了共识，但实测下来它只对一部分输入成立。我分两种情况量了一遍（code/rope_decay.py 第 2 段）。 先看彼此无关的随机 $q$、$k$： 不加 RoPE 6.2076 (基线) 相对距离 1 6.3182 (101.8% of 基线) 相对距离 64 6.2415 (100.5% of 基线) 相对距离 1024 6.2583 (100.8% of 基线) 相对距离 4000 6.4830 (104.4% of 基线) 完全没有衰减。 道理也简单：$R_m$ 是正交变换，对本文独立、各向同性高斯向量，不改变其内积的统计分布；任意非各向同性输入并无这个保证，转多少度都一样。 再看 $q = k$（一对高度相似的 token，比如相邻帧里的同一个主体）： 全 1 向量（gap=0 时内积 64.00）: 相对距离 1 +61.8337 ( 96.6% of gap=0) 相对距离 16 +38.7456 ( 60.5% of gap=0) 相对距离 64 +27.8122 ( 43.5% of gap=0) 相对距离 1024 +16.8272 ( 26.3% of gap=0) 相对距离 4000 +0.0989 ( 0.2% of gap=0) 衰减掉了 99.8%。 换个随机向量做同样的实验，结论一致（50.96 → 0.97）。 这两组实验只展示了各向同性独立输入与所选相似输入的差别。RoPE 的相位加权和依赖内容在各频段的分布；既不能推出“只有相似 token 才受影响”，也不能据此将真实视频的身份漂移归因于 RoPE，后者需要固定模型的消融实验。 顺带一个细节：衰减不是单调的。上面那个随机向量在 gap=256 处回升到 51.3%，比 gap=64 的 40.9% 还高。它是震荡衰减，因为各频率分量的相位在不同距离上会偶然重新对齐。所以拿单个距离点的数值下结论是不可靠的。 它无法表达绝对位置。 「这是第一帧」这类信息 RoPE 表达不了，因为它被设计成只看差值。需要绝对锚点的任务（比如首帧条件生成、指定时间戳编辑）得靠别的机制补，通常是额外的条件注入而不是位置编码。 算力和显存几乎不是代价。 这点顺手算清楚，免得担错心。拿 01 节那个 90000 token 的 720p 配置、假设 40 层 24 头、head_dim 64： 旋转本身约 33.2 GFLOPs，而同规模的注意力矩阵乘是 995.3 TFLOPs——RoPE 占 0.0033%，四舍五入就是零 cos/sin 表按半维存 fp32 是 22.0 MB，铺满 dim 是 43.9 MB（bf16 减半）。这个表在各层频率配置相同时可共享，算一次就够，相比激活值的开销可以忽略 所以选不选 RoPE 从来不是性能问题，纯粹是「要不要相对位置性和外推能力」的取舍。真正吃资源的是 5.4 节那个 fp32 约束——它强迫这一小段计算不能跟着整个网络一起降精度。 什么时候不该用：序列长度固定且很短、任务本身依赖绝对位置、或者模型是纯双向编码器且已有成熟的绝对位置方案时，换 RoPE 收益不大。RoPE 的价值集中在「训练短、推理长」这个场景上。 07. 经典论文脉络 五篇串起来就是位置编码近几年的主线： ALiBi（arXiv:2108.12409，2021-08）不做旋转，直接给注意力分数加一个随距离线性递减的偏置。外推能力很强，但把「远处不重要」写死成了先验，表达力受限。它是 RoPE 的主要对照组。 RoFormer（arXiv:2104.09864，2021-04）RoPE 的出处。贡献是用复数旋转同时拿到相对位置性和乘法式注入；留下的问题是没解决外推，那时也还没人需要 128K 上下文。 A Length-Extrapolatable Transformer（arXiv:2212.10554，2022-12）把「外推」这件事本身形式化，提出用注意力分辨率来度量，给后续方法提供了评价标准而不只是刷分。 Position Interpolation（arXiv:2306.15595，2023-06）最简单有效的一招：把位置索引整体除以缩放因子，让新长度挤回训练区间。代价是所有频率一起被压缩，近距离分辨率跟着下降。 YaRN（arXiv:2309.00071，2023-08）按频率分段处理——高频维度少动或不动以保住近距离分辨率，低频维度多压缩以覆盖新长度，再配一个注意力温度修正。这是目前长度外推的实用基线。 主线很清楚：RoPE 定下了「相对位置靠旋转」的框架，后面几年都在调同一组 $\theta_i$。理解了 3.3 节的频率分配，这些方法的差别只是一句话。 7.1 视频侧这条支线 上面五篇都出自语言模型。视觉生成这边是另一条时间线，起点晚但收敛很快： DiT（arXiv:2212.09748，2022-12）把扩散模型的骨干从 U-Net 换成 Transformer，位置信息用的还是固定的 sincos 绝对嵌入。它证明了 Transformer 在扩散上能 scale，但也把绝对位置嵌入那套外推问题一起继承了过来。 CogVideoX（arXiv:2408.06072，2024-08）视频 DiT 转向 3D RoPE 的代表作。diffusers 里那个 get_3d_rotary_pos_embed（5.6 节引的那段维度分配代码）就是为它写的。 HunyuanVideo（arXiv:2412.03603，2024-12）与 Wan（arXiv:2503.20314，2025-03）两个开源大模型延续了 3D RoPE 的路线。到这一步，「时间/高/宽三轴各分走一段 head_dim」基本成了视频 DiT 的默认选择。 值得注意的是这条支线继承了语言侧的全部外推工具：PI 和 NTK 的思路可以直接套到某一个轴上——比如只对空间两轴做插值来支持更高分辨率，时间轴不动。反过来，语言侧没有的新问题也在这里出现：三个轴的维度预算怎么分、变分辨率训练时坐标怎么归一化。这些留给「视频 DiT 里的 3D RoPE 与分辨率外推」。 08. 常见误解 「RoPE 是加在 embedding 上的」。 不是。它作用在每一层注意力的 $q$、$k$ 上，而且是乘法（旋转）。$v$ 在本文与常用实现里不旋转；存在其他变体，不能把所有旋转 value 的方案都判为错误。把它当成输入层的一次性操作，是最常见的理解错误。 「调大 rope_theta 就能延长上下文」。 只对了一半。调大 base 会让除最高频 $\theta_0=1$ 之外的频率变慢，能塞进更长序列，但不重新训练的话，高频维度的近距离分辨率会一起掉，表现为长文本能读了、短距离推理反而变差。YaRN 之所以要分频段处理，就是为了避免这个副作用。 「相对位置不变性意味着模型对绝对位置完全无感」。 不成立。注意力有 causal mask，第 0 个 token 只能看到自己，第 100 个能看到 100 个——这种「可见范围」本身就泄露了绝对位置。研究里管这叫隐式位置信息。 「间距相同分数就一定相同」。 还必须固定同一对内容向量 q、k；真实模型里的它们也会随上下文变化。真实模型里中间还夹着 softmax 归一化和 causal mask，同一对 token 在不同上下文里拿到的注意力权重是不同的。04 节验证的是旋转这一步的数学性质，不是整个注意力层的行为。 「RoPE 自带远程衰减」需要限定输入与统计口径。 独立各向同性高斯输入的内积分布不随旋转改变；所选相似向量可能在某些距离明显降低、在另一些距离回升。示例不是普遍衰减定律，也不能单独证明长视频质量的因果关系。 09. 动手验证 四个五分钟能跑完的小实验，用到的脚本全文都在文末附录，只依赖 NumPy，复制存成同名文件就能跑。另外跑一遍 make_figures.py 可以把本文三张配图重新生成（这个需要 matplotlib，其余脚本只要 NumPy）。 实验一：看外推在哪一步崩。 跑 rope_minimal.py，把 seq_len 从 8 改成 4096，然后对比位置 (0, 1) 和位置 (4000, 4001) 的内积。预期结果是两者仍然完全相等——因为数学性质对任意位置都成立。这恰好说明：外推失败不是公式失效，而是模型没在那个相位区间学过。这是本文最容易被误解的一点，值得亲手确认一次。 实验二：验证配对方式不能混用。 跑 rope_pairing.py，观察第 2、3 段输出。预期看到同一份权重下两种配对差 5.34，而重排维度后差值精确为 0。把 dim 改成 128 再跑一次，结论不变——这是结构性的，与维度大小无关。 实验三：亲手推翻一个常识。 跑 rope_decay.py，重点看 2a 和 2b 的差别。预期看到随机 $q$、$k$ 不衰减、$q = k$ 时衰减 99% 以上。然后把 2b 里的向量换成两个弱相关的向量试试（比如 k = 0.5 * q + 0.5 * rng.standard_normal(dim)），看衰减幅度怎么随相似度变化——这是理解 06 节那笔账的最快路径。 实验四：给自己的模型算外推预算。 跑 rope_extrapolate.py，把 SCALE 从 8 改成你实际需要的倍数（比如 480p→720p 大约是 2.25 倍 token），看 PI 和 NTK 各自在近距离分辨率上要付多少。预期规律是 PI 的保留率恒等于 $1/\text{scale}$，而 NTK 的保留率在高频维度接近 1、随维度下降。这些频率变化只能说明位置编码的几何差异，不能直接预测生成质量。如果任务依赖纹理、口型等近距离细节，应在目标分辨率上比较 PI、NTK 与必要的适配训练，结合质量和长距离一致性决定方案。 10. 延伸阅读 读这篇之前建议先看： 自注意力机制的计算与显存账本——RoPE 作用在 $q$、$k$ 上，得先清楚它们在注意力里的位置 读完这篇可以继续看： 视频 DiT 里的 3D RoPE 与分辨率外推——5.6 节那个 1:1.5:1.5 的示例维度分配，以及分辨率和帧数变化如何分别影响各轴相位 FlashAttention 为什么不需要存下注意力矩阵——RoPE 之后紧接着的那一步计算 附录：完整代码 09 节用到的脚本全文如下（make_figures.py、rope_decay.py、rope_minimal.py、rope_extrapolate.py、rope_pairing.py）。复制到本地存成同名文件，按各脚本开头的依赖说明准备环境后即可运行。 make_figures.py #!/usr/bin/env python3 # -*- coding: utf-8 -*- """生成「旋转位置编码 RoPE」的三张解释图。 数值全部来自同目录的三个脚本（rope_minimal / rope_extrapolate / rope_decay）， 这里只负责画——改了那边这里要重跑，免得图和正文数字打架。 三张图分别回答： 1. 32 个频率是怎么分层铺在 64 维上的（高频相位在近距离变化更快） 2. PI 与 NTK-aware 把这 8 倍外推的代价分别摊到了哪些频段 3. 「相对位置不变」和「随距离衰减」在数值上长什么样 只依赖 numpy + matplotlib。跑法：python make_figures.py """ import textwrap from pathlib import Path import matplotlib.pyplot as plt import numpy as np import rope_decay as RD import rope_extrapolate as RE ROOT = Path(__file__).resolve().parents[1] OUT = ROOT / "figures" OUT.mkdir(parents=True, exist_ok=True) plt.rcParams.update({ "font.sans-serif": ["PingFang SC", "Arial Unicode MS", "DejaVu Sans"], "axes.unicode_minus": False, "figure.dpi": 160, "savefig.bbox": "tight", }) INK = "#1f2937" MUTE = "#6b7280" C_ORIG = "#2f6fb0" # 原始 RoPE：蓝 C_PI = "#d1495b" # PI：红 C_NTK = "#2f9e6f" # NTK-aware：绿 C_BAD = "#d1495b" C_OK = "#2f9e6f" def style(ax, title, xlabel=None, ylabel=None): ax.set_title(title, fontsize=11.5, color=INK, pad=10, loc="left") if xlabel: ax.set_xlabel(xlabel, fontsize=10, color=MUTE) if ylabel: ax.set_ylabel(ylabel, fontsize=10, color=MUTE) ax.tick_params(colors=MUTE, labelsize=9) for s in ("top", "right"): ax.spines[s].set_visible(False) for s in ("left", "bottom"): ax.spines[s].set_color("#d1d5db") ax.grid(alpha=0.25, linewidth=0.6) ax.set_axisbelow(True) def footer(fig, text, width=118): """把「这张图要看什么」放到坐标轴下方。 必须放在 y<0 的位置：bbox_inches="tight" 会把负坐标的 artist 一起收进来， 放在 0~0.05 之间的话会和 x 轴标签叠在一起。 """ wrapped = "\n".join(textwrap.wrap(text, width=width)) fig.text(0.012, -0.13, wrapped, fontsize=8.5, color=MUTE, va="top", ha="left", linespacing=1.6) # ────────────────────────────────────────────────────────────────────── # 图 1：32 个频率怎么分层铺在 64 维上 # ────────────────────────────────────────────────────────────────────── def fig_freq_layers(): dim = RE.DIM th = RE.theta_original() # [dim/2] pos_short = np.arange(0, 256) pos_long = np.arange(0, 32768, 8) cos_short = np.cos(np.outer(pos_short, th)) # [256, dim/2] phase_long = np.outer(pos_long, th) # 累积相位（不取模） fig, axes = plt.subplots(1, 2, figsize=(12.8, 4.4)) fig.subplots_adjust(wspace=0.42) ax = axes[0] im = ax.imshow(cos_short.T, aspect="auto", origin="lower", cmap="RdBu_r", vmin=-1, vmax=1, extent=[0, len(pos_short), -0.5, len(th) - 0.5]) ax.set_yticks([0, 8, 16, 24, 31]) cb = fig.colorbar(im, ax=ax, fraction=0.035, pad=0.02) cb.ax.tick_params(labelsize=8, colors=MUTE) cb.set_label(r"$\cos(m\theta_i)$", fontsize=9, color=MUTE) style(ax, "近距离（0~255）：前 13 组完成至少一圈", "位置 m", "维度组 i") # 在 256 个位置内转满一整圈（2π）的维度组数 n_moving = int((pos_short[-1] * th / (2 * np.pi) >= 1).sum()) ax.annotate(f"256 长度内转满一整圈的只有前 {n_moving} 组", xy=(128, n_moving - 1), xytext=(12, n_moving + 7), fontsize=8.5, color="white", fontweight="bold", arrowprops=dict(arrowstyle="->", color=MUTE, lw=0.9)) ax = axes[1] for i, c in zip([0, 4, 8, 16, 24, 31], ["#1f6feb", "#2f9e6f", "#e0a03c", "#d1495b", "#8b5cf6", "#6b7280"]): ax.plot(pos_long, phase_long[:, i], lw=1.8, color=c, label=f"i={i}") trained = (RE.TRAIN_LEN - 1) * th[0] ax.axhline(trained, color=INK, ls="--", lw=1.2) ax.text(600, trained * 1.35, f"最高频训练最大相位 {trained:.0f} rad", fontsize=8.5, color=INK) ax.set_xscale("log") ax.set_yscale("log") ax.legend(fontsize=8.5, frameon=False, loc="upper left", ncol=2) style(ax, "远距离（到 32768）：各频段的累积相位", "位置 m（对数轴）", r"累积相位 $m\theta_i$（rad，对数轴）") fig.suptitle("图 1　RoPE 把 64 维切成 32 组，从「每步 1 弧度」一路排到「每步 1e-4 弧度」", fontsize=12, color=INK, x=0.012, ha="left", y=1.02) footer(fig, "左图：0 到 255 的相位差中，只有前 13 组超过一整圈；低频也有非零相位变化。" "右图：不同频率提供不同尺度，不是互斥的距离分工。虚线只代表最高频的训练相位上限；" "每个频段的训练相位上限都不同。") fig.savefig(OUT / "rope_freq_layers.png") plt.close(fig) print(f"[图1] 相位变化>=2π 的维度组数（0→255）: {n_moving}/{len(th)}") # ────────────────────────────────────────────────────────────────────── # 图 2：PI 与 NTK-aware 把 8 倍外推的代价摊在哪 # ────────────────────────────────────────────────────────────────────── def fig_extrapolation(): t_orig, t_pi, t_ntk = (RE.theta_original(), RE.theta_pi(), RE.theta_ntk()) idx = np.arange(len(t_orig)) slow_pi = t_orig / t_pi slow_ntk = t_orig / t_ntk keep_pi = t_pi / t_orig * 100 keep_ntk = t_ntk / t_orig * 100 fig, axes = plt.subplots(1, 2, figsize=(11.6, 4.3)) ax = axes[0] ax.plot(idx, slow_pi, color=C_PI, lw=2.4, label=r"PI：各维一律减速 8.00x") ax.plot(idx, slow_ntk, color=C_NTK, lw=2.0, label=r"NTK-aware：1.00x $\to$ 8.00x") ax.fill_between(idx, slow_ntk, slow_pi, color=C_PI, alpha=0.10) ax.set_yscale("log") ax.set_yticks([1, 2, 4, 8]) ax.yaxis.set_minor_formatter(plt.NullFormatter()) ax.get_yaxis().set_major_formatter(plt.FuncFormatter(lambda v, _: f"{v:g}x")) ax.legend(fontsize=9, frameon=False, loc="center right") style(ax, "谁被减速了多少（对数轴）", "维度组 i", r"$\theta_i^{\mathrm{orig}} / \theta_i^{\mathrm{new}}$") ax = axes[1] ax.plot(idx, keep_pi, color=C_PI, lw=2.4, label="PI：全线只剩 12.5%") ax.plot(idx, keep_ntk, color=C_NTK, lw=2.0, label=rf"NTK：{keep_ntk[0]:.0f}% $\to$ {keep_ntk[16]:.1f}%（i=16）") ax.fill_between(idx, keep_pi, keep_ntk, color=C_NTK, alpha=0.10) ax.axhline(100, color=MUTE, ls=":", lw=1.0) ax.set_yscale("log") ax.set_yticks([10, 20, 50, 100]) ax.get_yaxis().set_major_formatter(plt.FuncFormatter(lambda v, _: f"{v:g}%")) ax.legend(fontsize=9, frameon=False, loc="lower left") style(ax, "代价：相邻 token 的相位分辨率还剩多少", "维度组 i", r"$\theta_i^{\mathrm{new}} / \theta_i^{\mathrm{orig}}$") fig.suptitle("图 2　PI 把减速平均分给每一维，NTK 从高频不压缩逐渐过渡到低频压缩 8 倍", fontsize=12, color=INK, x=0.012, ha="left", y=1.02) footer(fig, "要看什么：左图 PI 是一条平线（每维都慢 8 倍），NTK 从最高频的 1.00x 单调爬到" "最低频的 8.00x。右图是这个分配的代价——PI 把每一维的近距离分辨率都砍到 12.5%，" "这只是相位步长，不能直接预测任务质量；NTK 保住高频（100%），最低频与 PI 同为 12.5%。" "两者改的是同一组 theta，区别只在把代价放在哪个频段。") fig.savefig(OUT / "rope_extrapolation.png") plt.close(fig) print(f"[图2] NTK 减速比: i=0 {slow_ntk[0]:.2f}x -> i=31 {slow_ntk[-1]:.2f}x；" f"分辨率保留 i=16: PI {keep_pi[16]:.1f}% / NTK {keep_ntk[16]:.1f}%") # ────────────────────────────────────────────────────────────────────── # 图 3：相对位置不变 + 随距离衰减 # ────────────────────────────────────────────────────────────────────── def fig_relativity(): dim, seq_len = 64, 8192 cos, sin = RD.build_rope_cache(seq_len, dim) rng = np.random.default_rng(0) # ── 左：同一对 q,k，只平移绝对位置，内积应当逐点不变 ── q = rng.standard_normal(dim) k = rng.standard_normal(dim) gaps = np.arange(0, 512) starts = [0, 500, 2000] curves = {} for s in starts: qr = RD.apply_rope(q[None, :], cos[s:s + 1], sin[s:s + 1])[0] kr_all = RD.apply_rope( np.repeat(k[None, :], len(gaps), 0), cos[s:s + len(gaps)], sin[s:s + len(gaps)]) curves[s] = kr_all @ qr # ── 右：距离越远，RoPE 把「无位置编码的原始内积」改写得多彻底 ── # 独立各向同性高斯输入：corr(q^T R_gap k, q^T k) = (2/dim) Σ_i cos(gap θ_i)。 # 注意：这是「与原始内积的相关性」衰减，不是内积的绝对值变小—— # 独立各向同性高斯输入的内积分布与距离无关（RMS 为 sqrt(dim)）。 Q = rng.standard_normal((4000, dim)) K = rng.standard_normal((4000, dim)) s0 = np.einsum("ij,ij->i", Q, K) th = 1.0 / (10000.0 ** (np.arange(0, dim, 2) / dim)) probe = np.unique(np.round(np.geomspace(1, 4096, 60)).astype(int)) emp, theo = [], [] for g in probe: qr = RD.apply_rope(Q, cos[0:1], sin[0:1]) kr = RD.apply_rope(K, cos[g:g + 1], sin[g:g + 1]) sg = np.einsum("ij,ij->i", qr, kr) emp.append(float(np.corrcoef(sg, s0)[0, 1])) theo.append(float(2.0 / dim * np.sum(np.cos(g * th)))) fig, axes = plt.subplots(1, 2, figsize=(11.6, 4.3)) ax = axes[0] cols = {0: C_ORIG, 500: C_NTK, 2000: "#e0a03c"} for s in starts: ax.plot(gaps, curves[s], lw=1.6, color=cols[s], label=f"绝对起点 {s}") ax.legend(fontsize=9, frameon=False, loc="upper right") style(ax, "同一对向量、同一组间距，换绝对位置后曲线完全重合", "相对距离 (n-m)", r"$q_m^\top R_{n-m} k_n$") ax = axes[1] ax.plot(probe, emp, marker="o", ms=3.5, lw=1.8, color=C_ORIG, label="实测（4000 对随机向量）") ax.plot(probe, theo, lw=1.4, ls="--", color=C_BAD, label=r"理论 $\frac{2}{D}\sum_i\cos(\Delta\theta_i)$") ax.axhline(0, color=MUTE, ls=":", lw=1.0) ax.set_xscale("log") ax.legend(fontsize=9, frameon=False, loc="upper right") style(ax, "独立高斯输入：与原始内积的相关性随距离振荡", "相对距离（对数轴）", r"$\mathrm{corr}(q^\top R_d k, q^\top k)$") fig.suptitle("图 3　固定内容的相对位置不变性，与高斯输入下的分数相关性", fontsize=12, color=INK, x=0.012, ha="left", y=1.02) max_dev = max(float(np.abs(curves[s] - curves[0]).max()) for s in starts) footer(fig, f"左图：固定 q/k 的三条曲线在浮点误差内重合，最大偏差 {max_dev:.1e}，不是逐比特相等。" "右图：D=64、Δ 为间距，理论相关性是有限余弦和；它会振荡，不保证单调或在无穷远趋于零。" "本图的独立各向同性高斯输入下，内积分布不随旋转改变。") fig.savefig(OUT / "rope_relativity.png") plt.close(fig) print(f"[图3] 不同绝对起点的最大逐点偏差: {max_dev:.2e}；" f"corr: d=1 {emp[0]:.3f} → d={probe[-1]} {emp[-1]:.3f}") def main(): fig_freq_layers() fig_extrapolation() fig_relativity() print(f"[OK] 三张图已写入 {OUT}") for f in sorted(OUT.glob("*.png")): print(f" {f.name} {f.stat().st_size / 1024:.0f} KB") if __name__ == "__main__": main() rope_decay.py """观测 RoPE 的两个常被口头断言、但很少有人实测的性质。 1. 旋转不改变模长（对比加法式绝对位置嵌入） 2. 所选输入的内积随距离变化；并非普遍衰减定律 运行 `python rope_decay.py` 复现。 """ import numpy as np def build_rope_cache(seq_len, dim, base=10000.0): inv_freq = 1.0 / (base ** (np.arange(0, dim, 2) / dim)) freqs = np.outer(np.arange(seq_len), inv_freq) return np.cos(freqs), np.sin(freqs) def apply_rope(x, cos, sin): out = np.empty_like(x) out[..., 0::2] = x[..., 0::2] * cos - x[..., 1::2] * sin out[..., 1::2] = x[..., 0::2] * sin + x[..., 1::2] * cos return out def demo_norm_preserved(rng, dim=64, seq_len=512): """旋转保模长，加法不保。""" cos, sin = build_rope_cache(seq_len, dim) x = rng.standard_normal(dim) print("=== 1. 模长是否被位置编码改变 ===") print(f" 原始向量模长 {np.linalg.norm(x):.6f}") norms = [np.linalg.norm(apply_rope(x[None, :], cos[p:p + 1], sin[p:p + 1])) for p in (0, 1, 100, 511)] print(f" RoPE 后（位置 0/1/100/511） " f"{' '.join(f'{n:.6f}' for n in norms)}") print(f" 跨位置模长极差 {max(norms) - min(norms):.2e}") # 加法式绝对位置嵌入：同尺度的可学习向量直接相加 pe = rng.standard_normal((seq_len, dim)) add_norms = [np.linalg.norm(x + pe[p]) for p in (0, 1, 100, 511)] print(f" 加法式 PE 后 " f"{' '.join(f'{n:.6f}' for n in add_norms)}") print(f" 跨位置模长极差 " f"{max(add_norms) - min(add_norms):.4f} ← 同一个 token 的模长被位置改写了") def _score(cos, sin, q, k, gap): qr = apply_rope(q[None, :], cos[0:1], sin[0:1]) kr = apply_rope(k[None, :], cos[gap:gap + 1], sin[gap:gap + 1]) return float(qr[0] @ kr[0]) GAPS = (1, 4, 16, 64, 256, 1024, 4000) def demo_long_range_decay(rng, dim=64, seq_len=4096, trials=2000): """「RoPE 自带远程衰减」到底成不成立——分两种输入分布看。 结论：本文独立各向同性高斯 q/k 的内积分布不变；所选 q=k 示例会振荡。 """ cos, sin = build_rope_cache(seq_len, dim) print("\n=== 2a. 随机无关的 q/k：取 |内积| 均值 ===") baseline = float(np.mean([abs(float(rng.standard_normal(dim) @ rng.standard_normal(dim))) for _ in range(trials)])) print(f" 不加 RoPE {baseline:.4f} (基线)") for gap in GAPS: m = float(np.mean([abs(_score(cos, sin, rng.standard_normal(dim), rng.standard_normal(dim), gap)) for _ in range(trials)])) print(f" 相对距离 {gap:>5} {m:.4f} ({m / baseline * 100:5.1f}% of 基线)") print(" → 没有衰减趋势。独立各向同性高斯输入在正交旋转下分布不变。") print("\n=== 2b. q = k（高度相似的一对 token）===") for name, vec in (("全 1 向量", np.ones(dim)), ("某个随机向量", np.random.default_rng(7).standard_normal(dim))): # gap=0 时位置 0 的旋转角为 0，内积就是未加 RoPE 的原值 head = _score(cos, sin, vec, vec, 0) print(f" {name}（gap=0 时内积 {head:.2f}）:") for gap in GAPS: s = _score(cos, sin, vec, vec, gap) print(f" 相对距离 {gap:>5} {s:+9.4f} ({s / head * 100:5.1f}% of gap=0)") print(" → 这些采样点多低于 gap=0，也有回升；不代表任意输入或距离都衰减。") if __name__ == "__main__": demo_norm_preserved(np.random.default_rng(0)) demo_long_range_decay(np.random.default_rng(1)) rope_minimal.py #!/usr/bin/env python3 # -*- coding: utf-8 -*- """RoPE 的最小可运行实现，验证「注意力分数只依赖相对位置」这一核心性质。 只依赖 numpy，直接 python rope_minimal.py 即可运行。 刻意不用 PyTorch：讲原理时框架的抽象反而是噪声，而且 numpy 谁都能跑。 生产实现见文章第 05 节对 transformers 的引用。 """ import numpy as np def build_rope_cache(seq_len: int, dim: int, base: float = 10000.0): """预计算 RoPE 的 cos/sin 表。 dim 必须是偶数：每两个相邻维度配成一个复数参与旋转。 base 就是常说的 rope_theta，调它会按维度改变旋转速度；最高频 i=0 始终为 1。 """ assert dim % 2 == 0, f"dim 必须为偶数，收到 {dim}" # theta_i = base^(-2i/dim)：i 越大转得越慢，对应越低的频率 inv_freq = 1.0 / (base ** (np.arange(0, dim, 2) / dim)) pos = np.arange(seq_len) freqs = np.outer(pos, inv_freq) # [seq_len, dim/2] return np.cos(freqs), np.sin(freqs) def apply_rope(x: np.ndarray, cos: np.ndarray, sin: np.ndarray) -> np.ndarray: """把旋转作用到 x 上。x: [seq_len, dim] 最后一维两两配对看成复数 (x_even + i*x_odd)，乘上 e^{i*m*theta}： 实部 = x_even * cos - x_odd * sin 虚部 = x_even * sin + x_odd * cos """ x_even, x_odd = x[..., 0::2], x[..., 1::2] out = np.empty_like(x) out[..., 0::2] = x_even * cos - x_odd * sin out[..., 1::2] = x_even * sin + x_odd * cos return out if __name__ == "__main__": rng = np.random.default_rng(0) seq_len, dim = 8, 64 cos, sin = build_rope_cache(seq_len, dim) print(f"cos/sin 形状: {cos.shape} {sin.shape}") print(f"位置 1 的前 4 维 cos: {np.round(cos[1, :4], 4).tolist()}") print(f"位置 1 的末 2 维 cos: {np.round(cos[1, -2:], 6).tolist()}") # 固定一对向量，只改赋予它们的绝对位置，间距始终为 1 q = rng.standard_normal(dim) k = rng.standard_normal(dim) print("\n同一对向量，平移绝对位置（间距恒为 1）:") for shift in range(4): qr = apply_rope(q[None, :], cos[shift:shift + 1], sin[shift:shift + 1]) kr = apply_rope(k[None, :], cos[shift + 1:shift + 2], sin[shift + 1:shift + 2]) print(f" 位置 ({shift}, {shift + 1}) → 内积 {float(qr[0] @ kr[0]):.6f}") print("\n对照组，改变间距:") for gap in (1, 2, 3): qr = apply_rope(q[None, :], cos[0:1], sin[0:1]) kr = apply_rope(k[None, :], cos[gap:gap + 1], sin[gap:gap + 1]) print(f" 间距 {gap} → 内积 {float(qr[0] @ kr[0]):.6f}") print(f"\n未加 RoPE 的原始内积: {float(q @ k):.6f}") rope_extrapolate.py """两种长度外推方法（PI / NTK-aware）到底改了什么。 结论先说：它们改的是同一组 theta_i，区别只在「怎么分配这次减速」。 运行 `python rope_extrapolate.py` 复现。 """ import numpy as np DIM = 64 BASE = 10000.0 TRAIN_LEN = 4096 # 训练时见过的长度 SCALE = 8 # 想外推到 8 倍，即 32768 def theta_original(dim=DIM, base=BASE): return 1.0 / (base ** (np.arange(0, dim, 2) / dim)) def theta_pi(dim=DIM, base=BASE, scale=SCALE): """Position Interpolation：位置索引整体除以 scale。 等价于把每个 theta_i 都乘上 1/scale——所有频率被同等减速。 """ return theta_original(dim, base) / scale def theta_ntk(dim=DIM, base=BASE, scale=SCALE): """NTK-aware：把 base 放大，让减速量随维度索引递增（频率越低减速越多）。 base' = base * scale^(d/(d-2))，于是高频维度几乎不受影响， 最低频维度与 PI 同样减速 scale 倍，其余维度的减速更少。 """ new_base = base * scale ** (dim / (dim - 2)) return theta_original(dim, new_base) def main(): t_orig, t_pi, t_ntk = theta_original(), theta_pi(), theta_ntk() print("=== 1. 三种方案下各维度的旋转速度 theta_i ===") print(" 维度 i 原始 PI(÷8) NTK-aware PI减速比 NTK减速比") for i in (0, 1, 4, 8, 16, 24, 31): print(f" {i:>4} {t_orig[i]:.3e} {t_pi[i]:.3e} {t_ntk[i]:.3e}" f" {t_orig[i] / t_pi[i]:>7.2f}x {t_orig[i] / t_ntk[i]:>7.2f}x") print(f" → PI 对每一维都减速 {t_orig[0] / t_pi[0]:.2f}x（均匀）；" f"NTK 从最高频的 {t_orig[0] / t_ntk[0]:.2f}x 递增到最低频的 " f"{t_orig[-1] / t_ntk[-1]:.2f}x（越低频压得越狠）") print("\n=== 2. 外推到 8 倍长度后，各方案在第 0 维（最高频）的相位 ===") target = TRAIN_LEN * SCALE trained_max = (TRAIN_LEN - 1) * t_orig[0] print(f" 训练时见过的最大相位（位置 {TRAIN_LEN - 1}） {trained_max:9.1f} rad") for name, th in (("原始（直接外推）", t_orig), ("PI", t_pi), ("NTK-aware", t_ntk)): v = (target - 1) * th[0] flag = ("落在训练区间内" if v <= trained_max else f"接近但略超训练端点 {trained_max:g}" if v <= trained_max * 1.01 else f"超出 {v / trained_max:.1f} 倍") print(f" {name:<18} 位置 {target - 1} {v:9.1f} rad ← {flag}") print(" → NTK 在最高频维度上「不作为」是刻意的，它保留高频相位步长，逐渐压缩更低频的相位步长") print("\n=== 3. 代价：近距离分辨率还剩多少 ===") print(" (相邻位置差 1 时转过的角度，按维度看。越小＝越难区分相邻 token)") print(" 维度 i 原始 PI NTK PI保留 NTK保留") for i in (0, 4, 8, 16): print(f" {i:>4} {t_orig[i]:.6f} {t_pi[i]:.6f} {t_ntk[i]:.6f}" f" {t_pi[i] / t_orig[i] * 100:5.1f}% {t_ntk[i] / t_orig[i] * 100:5.1f}%") print(f" → PI 在每一维都只保留 {t_pi[0] / t_orig[0] * 100:.1f}% 的分辨率；" f"NTK 在第 0 维保留 {t_ntk[0] / t_orig[0] * 100:.1f}%、" f"第 16 维只剩 {t_ntk[16] / t_orig[16] * 100:.1f}%。") print(" 两者都在同一组 theta 上做手脚，只是把代价放在了不同频段——" "\n YaRN 进一步把这个分配变成可调的分段策略。") if __name__ == "__main__": main() rope_pairing.py """对照 RoPE 的两种维度配对方式，并验证它们只差一个固定置换。 奇偶交错（本文最小实现）：第 2i 维和第 2i+1 维配成一个复数 前后半切分（HuggingFace）：第 i 维和第 i+d/2 维配成一个复数 结论先说：两者都满足相对位置不变性，但对同一份权重不可互换。 运行 `python rope_pairing.py` 复现下面的数值。 """ import numpy as np def rope_interleaved(x, pos, dim, base=10000.0): """奇偶交错配对：(x_0,x_1), (x_2,x_3), ... 每对是一个复平面。""" inv_freq = 1.0 / (base ** (np.arange(0, dim, 2) / dim)) ang = pos * inv_freq # [dim/2] cos, sin = np.cos(ang), np.sin(ang) out = np.empty_like(x) out[0::2] = x[0::2] * cos - x[1::2] * sin out[1::2] = x[0::2] * sin + x[1::2] * cos return out def rope_half_split(x, pos, dim, base=10000.0): """前后半切分配对：(x_i, x_{i+d/2})。等价于 HF 的 rotate_half 写法。 HF 把 cos/sin 表沿最后一维复制两遍（cat((freqs, freqs))）， 于是 q*cos + rotate_half(q)*sin 展开后就是下面两行。 """ half = dim // 2 inv_freq = 1.0 / (base ** (np.arange(0, dim, 2) / dim)) ang = pos * inv_freq # [dim/2] cos, sin = np.cos(ang), np.sin(ang) x1, x2 = x[:half], x[half:] return np.concatenate([x1 * cos - x2 * sin, x2 * cos + x1 * sin]) def interleave_to_half(x): """把奇偶交错的排布重排成前后半切分的排布。 交错版的第 2i 维（复数实部）搬到前半的第 i 位， 第 2i+1 维（虚部）搬到后半的第 i 位。 """ return np.concatenate([x[0::2], x[1::2]]) def attn_score(rope_fn, q, k, m, n, dim): return float(rope_fn(q, m, dim) @ rope_fn(k, n, dim)) def main(): rng = np.random.default_rng(0) dim = 64 q = rng.standard_normal(dim) k = rng.standard_normal(dim) print("=== 1. 两种配对都满足相对位置不变性 ===") for name, fn in (("奇偶交错", rope_interleaved), ("前后半切分", rope_half_split)): scores = [attn_score(fn, q, k, m, m + 3, dim) for m in range(4)] spread = max(scores) - min(scores) print(f" {name}: 间距恒为 3 时内积 {scores[0]:.6f}，" f"四个绝对位置间的最大波动 {spread:.2e}") print("\n=== 2. 同一份权重下，两种配对给出的分数不同 ===") s_i = attn_score(rope_interleaved, q, k, 0, 3, dim) s_h = attn_score(rope_half_split, q, k, 0, 3, dim) print(f" 奇偶交错 {s_i:.6f}") print(f" 前后半切分 {s_h:.6f}") print(f" 差值 {abs(s_i - s_h):.6f} ← 不为 0，实现和权重必须配套") print("\n=== 3. 重排维度后两者完全一致 ===") qp, kp = interleave_to_half(q), interleave_to_half(k) s_h_perm = attn_score(rope_half_split, qp, kp, 0, 3, dim) print(f" 奇偶交错(原始排布) {s_i:.6f}") print(f" 前后半切分(重排后权重) {s_h_perm:.6f}") print(f" 差值 {abs(s_i - s_h_perm):.2e}" f" ← 只差一个固定的维度置换") if __name__ == "__main__": main()