# LoGRA：用低秩梯度 Sketch 扩展大模型强化学习——中文精读笔记

**原题**：*LoGRA: Scaling LLM Reinforcement Learning with Low-Rank Gradient Sketches*  
**作者**：Shaokun Zhang、Yifan Zhang、Jian Hu、Yueying Li、Hao Zhang、Binfeng Xu、Jan Kautz、Yi Dong（NVIDIA）  
**版本**：[arXiv:2610.06647v3](https://arxiv.org/abs/2610.06647v3)，2026-10-17 修订；本次获取于 2026-10-23。  
**原文**：[本地 PDF](../source/paper.pdf) · [作者 TeX 源码](../source/tex/files/main.tex) · [代码链接（作者提供）](https://github.com/skzhang1/labs-molt/tree/logra/examples/scripts/logra)

## 论文总览

**问题**：大模型强化学习后训练需要反复生成、更新和同步策略。即使模型权重放得下，稠密梯度与 Adam 状态仍会把训练 GPU 显存推到容量之外。LoGRA 把选定权重矩阵的梯度直接累计为低秩 sketch，并用 sketch 构造权重更新；另用预测 KL 调整步长，避免一次更新使策略概率突变。它的重点是**保留完整权重矩阵、压缩更新所需的信息**。（Sec. 1–2，PDF pp. 1–5）

在论文的八张 H100 80GB、其中四张训练的配置下，1.5B 与 7B 模型的平均训练显存分别降低 **21.8%** 和 **45.7%**；7B 的 MATH-500 峰值 Pass@1 与稠密 Adam 基本相同。27B 的密集 Adam 在首次分配优化器状态时 OOM，而 LoGRA 完成 200 次数学任务更新，并在另一组 Reasoning-Gym 任务上持续训练超过 1,100 个记录步。（Table 1，PDF p. 7；Fig. 4，PDF p. 7）

## 结构地图

| 原文章节 | 论证作用 | PDF 位置 |
| --- | --- | --- |
| 1. Introduction | 说明 RL 显存瓶颈与压缩梯度的动机 | pp. 1–2；Fig. 1 |
| 2.1. Low-rank gradient sketches | 定义 sketch、直接累计与低秩权重更新 | pp. 2–3；Eqs. (1)–(4) |
| 2.2. Applying LoGRA to RL post-training | 训练/生成端更新流程与同步载荷 | pp. 3–4；Algorithm 1 |
| 2.3. Step-size control with predicted KL | 用局部 KL 曲率选择更新倍率 | pp. 4–5；Eqs. (5)–(7) |
| 3.1–3.2. Setup / Main Experiments | 模型、数据、基线、显存及准确率 | pp. 5–7；Figs. 2–3，Table 1 |
| 3.3. Sustained 27B training | 单节点长程训练 | p. 7；Fig. 4 |
| 3.4. Further Analysis | 秩/投影消融、LoRA 比较、KL 预算 | pp. 7–10；Figs. 5–6，Table 2 |
| 4–5. Related Works / Conclusion | 与参数高效及系统方法的关系和总结 | p. 10 |
| Appendix A–D | KL 推导、RowAdam、数据与超参数 | pp. 13–17；Tables 3–4 |

## 一：动机——为什么压缩梯度，而不是只压缩权重

以 7B 模型为例，BF16 权重约需 14 GB；Adam 的两个 FP32 动量缓冲区另需约 56 GB，还没算梯度和激活。FSDP、检查点重算和卸载能重新分配这些开销；LoRA 则冻结原权重、只训练适配器。LoGRA 选择另一条路径：原矩阵仍直接接收更新，但其**梯度累计、优化器状态和更新表示**尽可能保持紧凑。这也适合训练器与生成器反复同步权重的 RL 工作流。（Sec. 1，PDF pp. 1–2）

作者用二值结果奖励说明直觉：每个回答只收到很少的外部反馈，未必需要保存高维梯度的每个坐标才能有效改进策略。这个直觉提出压缩的可能性；实际保真度和训练效果仍由后面的消融与任务实验检验。（Sec. 1，PDF p. 2）

## 二：方法——直接累计低秩梯度并控制策略变化

### 2.1 Sketch 的维度、计算与含义

对一层权重 $W\in\mathbb R^{d\times k}$，稠密梯度为 $G=\nabla_W\mathcal L\in\mathbb R^{d\times k}$。LoGRA 为它抽取固定于本次累计过程的随机投影 $A\in\mathbb R^{r\times k}$，其中 $r\ll\min(d,k)$，得到

$$
S=GA^\top\in\mathbb R^{d\times r},\qquad
\widehat G=SA=(GA^\top)A,\qquad
W\leftarrow W-\alpha\eta UA.
$$

$S$ 是梯度 sketch，$U$ 是优化器处理后的同形状 sketch；基础 SGD 取 $U=S$。$\eta$ 是标称步长，$\alpha$ 是下节控制的全局倍率。$SA$ 与 $G$ 同形状，但单步更新的秩最多为 $r$；这是**近似梯度**，不是精确重建。默认 $A_{ij}$ 是独立的 $\pm1/\sqrt r$。由这个采样定义可推得 $\mathbb E[A^\top A]=I$，因此固定 $G$ 时 $\mathbb E[GA^\top A]=G$；单次采样仍有重建误差。（Eqs. (1)–(4)，PDF pp. 2–3；期望等式为依据原式的推导）

关键实现不是先算出 $G$ 再压缩。以线性层的一个样本为教学例子，若输入 $x\in\mathbb R^k$、反传导数 $g\in\mathbb R^d$，该样本贡献 $gx^\top$，则其 sketch 可直接算作 $g(Ax)^\top$。对每个微批相加时，只保留 FP32 的 $d\times r$ 累计器。权重正常参与前向计算；压缩的是选定矩阵的梯度缓存。对 $d=k=4096,r=64$，一个 FP32 累计器由 64 MiB 降至 1 MiB；实际总显存还包含权重、投影、激活等。（Sec. 2.1–2.2，PDF p. 3）

![原文 Fig. 1：不同模型规模下的训练显存](figures/table1_memory_comparison.png)

原图 Fig. 1（PDF p. 1）展示**每张训练 GPU**的平均/峰值显存；27B 的稠密条只表示 OOM 门槛，不能读作测得的稠密显存值。图中横轴在 34–48 GiB 处断开。

### 2.2 多步更新、RowAdam 与生成端同步

每个选定的 attention/MLP 投影矩阵各有 $A,S,U$，其它参数固定；整步共用同一个 $\alpha$。每步重采样 $A$，因此各步低秩增量可落在不同子空间，合并到完整 $W$ 后，累计变化无需保持秩 $r$。训练器可以只发送因子 $\alpha\eta U$ 与随机种子，让生成端重建 $A$ 并施加同一更新，理论载荷由 $O(dk)$ 降至 $O(dr)$ 加少量元数据。（Algorithm 1，PDF pp. 3–4）

主实验为与 Adam 保留一定自适应性，使用 **RowAdam**：对 sketch 的每一行追踪平方均值 $v_{t,i}$，按 $1/(\sqrt{\widehat v_{t,i}}+\epsilon_t)$ 缩放，再把处理后矩阵的 Frobenius 范数调回原 sketch 的范数。它只需每行一个二阶状态，$\beta_2=0.95$，没有 Adam 的一阶动量，也不是逐参数 Adam。因此 Table 1 的对照同时改变了梯度表示和优化器。（Appendix B，PDF pp. 14–15）

### 2.3 预测 KL：让步长适应策略敏感度

只限定权重更新范数不能保证输出分布变化小。令所有选定矩阵的候选更新组成 $D$，其中 $D_W=-\eta UA$；在训练批中抽取若干已生成 token 的上下文 $\mathcal P$，比较同一上下文、同一温度 $T$ 下的旧/新 next-token 分布。局部展开给出

$$
\frac1{|\mathcal P|}\sum_{s\in\mathcal P}
D_{\mathrm{KL}}\!\left(\pi_\theta(\cdot\mid s)\middle\|\pi_{\theta+\alpha D}(\cdot\mid s)\right)
=\alpha^2q(D)+O(\alpha^3),
$$

$$
q(D)=\frac{1}{2|\mathcal P|T^2}\sum_{s\in\mathcal P}
\operatorname{Var}_{v\sim p(s)}[\dot z_v(s)],\qquad
\alpha=\min\left\{\alpha_{\max},\sqrt{\delta/q(D)}\right\}.
$$

这里 $p(s)=\operatorname{softmax}(z(s)/T)$，$\dot z_v(s)$ 是 token $v$ 的 logit 沿 $D$ 的方向导数，$q(D)$ 是未缩放候选步的局部 KL 曲率系数，$\delta$ 是平均预测 KL 预算。若 $q(D)=4\delta$，步长至多缩到原来一半；若预测很小，倍率也受 $\alpha_{\max}$ 限制。方差项只关心 token 间**相对** logit 变化：给所有 logit 加同一常数不改变 softmax。附录推导先用 $K_s(0)=K'_s(0)=0$ 得到二次项，再将其写为方向导数的概率加权方差。局部近似约束的是预测 KL，实际更新后的 KL 可能不同。（Eqs. (5)–(7)，PDF pp. 4–5；Appendix A，pp. 13–14）

27B 长程实验为效率使用直接 KL 评估的**有限差分探针**估计 $q(D)$，而非 Eq. (6) 的 JVP 形式；因此“公式中的估计量”和“所有实验的具体实现”不能简单画等号。（Sec. 3.3，PDF p. 7）

## 三：实验设计与指标

主比较在一台八 H100 80GB 节点上进行：四张 GPU 做 FSDP 训练，四个独立生成引擎占用其余 GPU；BF16、激活检查点、无 CPU 卸载。模型为 Qwen2.5-Math-1.5B/7B 和 Qwen3.8-27B。DAPO-Math-7.5K 取 7,500 个去重提示；每步 128 个提示、每提示一个回答，训练上下文上限 4,096 token。MATH-500 测试集有 500 题，每题生成四个回答；Pass@1 是这 2,000 个回答的平均二值正确率，数值越高越好。（Sec. 3.1–3.2，PDF p. 5；Appendix C–D，pp. 15–16）

主实验每个可运行设置有三个随机种子。LoGRA 用秩 256、逐步刷新 Rademacher 投影、RowAdam、从 $2\times10^{-4}$ 余弦退火到 $2\times10^{-5}$ 的预测 KL 预算；稠密基线用 Adam 和学习率退火。原文称 1.5B/7B 跑 20,000 次更新、27B 跑 200 次；**图 2 与峰值准确率只统计前 200 次**，Table 1 的显存和吞吐覆盖完整运行。显存是训练 GPU 的每步峰值分配量，表中“平均”再对时间取均值；“峰值”是三种子均值曲线的最大值，不是单卡瞬时最大值。吞吐为每小时完成更新数。（Sec. 3.2，PDF pp. 5–7；Table 1）

## 四：结果——能省多少，学习是否保留

### 4.1 主实验：显存大幅下降，吞吐近似持平

| 模型 | 方法 | 平均显存 GiB/GPU ↓ | 峰值显存 GiB/GPU ↓ | 更新/小时 ↑ | 峰值 Pass@1 % ↑ | 峰值 Pass@4 % ↑ |
| --- | --- | ---: | ---: | ---: | ---: | ---: |
| 1.5B | Dense | 9.18 | 9.19 | 54.28 | 63.77 | 80.33 |
| 1.5B | LoGRA | **7.18** | **8.66** | 55.44 | **67.87** | **82.47** |
| 7B | Dense | 31.82 | 31.85 | 25.28 | **72.48** | **85.33** |
| 7B | LoGRA | **17.29** | **22.26** | 25.73 | 72.33 | 85.13 |
| 27B | Dense | OOM | OOM | — | — | — |
| 27B | LoGRA | 51.54 | 53.88 | 6.99 | 71.52 | 81.87 |

数值为 Table 1 的三种子均值（PDF p. 7）；原表还给出样本标准差。1.5B 的平均显存相对 Dense 降 **21.8%**，7B 降 **45.7%**。7B Pass@1 只差 0.15 个百分点，吞吐 25.73 对 25.28 更新/小时；因此在这一规模的主要收益是显存容量，而不是明显提速。27B Dense 在首次 Adam 状态分配时 OOM，无法计算同条件下的显存节省百分比或吞吐比。Fig. 2 的曲线补充显示前 200 次更新的学习过程，Fig. 3 展示每步显存而不是整个节点的总显存。（PDF pp. 6–7）

![原文 Fig. 2：前三个模型规模的 MATH-500 学习曲线](figures/flora_pass1.png)

原图 Fig. 2（PDF p. 6）：左为 Pass@1、右为 Pass@4，横轴是实际优化器更新，曲线/阴影为三种子的均值/样本标准差。

### 4.2 另一组 27B 长程实验：超过 1,100 个记录步

单独的 Reasoning-Gym Hard 实验包含八种可验证任务，训练集 38,948 题、留出集 193 题；指标是八任务 verifier 分数的**等权宏平均**。该运行用秩 256、SGD sketch、16,384-token 上下文和逐渐收紧的 KL 预算，只有一条 seed 42 轨迹，约每 20 个记录步评估一次。（Sec. 3.3，PDF p. 7；Appendix C–D，pp. 15–16）

留出宏平均从 **39.69%** 升至 **62.94%**，第 1,060 步峰值 **65.52%**；最后五次评估处于 59.11%–64.16%。同一轨迹的平均回答长度从 8,312 降至 4,705 token，达到上下文上限的比例从 34.72% 降至 12.95%。这些读数说明所报告配方在这条轨迹中持续学习，且没有观察到持续性分数崩塌。（Fig. 4，PDF p. 7）

![原文 Fig. 4：27B 长程训练的留出分数、回答长度与截断率](figures/long27_stability.png)

原图 Fig. 4（PDF p. 7）：三面板共用“训练步”横轴；分数是八任务宏平均，长度与截断率帮助判断分数变化是否伴随输出行为异常。

### 4.3 压缩设计：提高秩的效应最明显

在 1.5B 的单次消融中，秩 $r\in\{4,16,64,256\}$ 越高，MATH-500 Pass@1 越高；从 4 到 256，三种投影配置提升 **3.02–3.95 个百分点**。秩 256 时固定 Rademacher、逐步刷新 Rademacher、逐步刷新 Gaussian 分别为 **68.50%、68.31%、67.42%**，短程内没有稳定的“刷新必优”排序。对应的梯度重建余弦均值约从 0.046 升至 0.348；同层同秩的刷新与分布差异远小于秩差异。余弦只测单个微批的重建相似度，任务表现仍以图 5(a) 为依据。（Sec. 3.4.1，Fig. 5，PDF p. 8；Appendix D.3，p. 17）

![原文 Fig. 5：秩、基底刷新及投影分布的消融](figures/rank_refresh.png)

原图 Fig. 5（PDF p. 8）：(a) 是任务准确率；(b)(c) 是逐层重建余弦热图，色条统一为 0–0.4，不能把热图读成准确率。

### 4.4 LoRA 比较：显存更低，但速度与 Pass@1 未占优

在 1.5B、秩 256、三种子设置中，LoGRA / LoRA 的平均训练显存为 **7.18 / 13.21 GiB**，峰值为 **8.66 / 13.38 GiB**；Pass@1 是 **68.47% / 70.45%**，训练调用耗时是 **11.31 / 10.63 秒/更新**。LoGRA 的 Pass@4 高 0.40 个百分点。这里的秒/更新**不含生成与同步**；LoRA 与 LoGRA 使用不同优化器、步长计划与计数方式，因此这张表支持“此配方下 LoGRA 更省训练显存”，不能单独归因于适配器参数化差异。（Sec. 3.4.2，Table 2，PDF p. 9；Appendix D.4，p. 17）

### 4.5 KL 预算：保守预算比过大预算更可靠

1.5B 长程消融采用两种子、SGD sketch，比较无 KL 控制、固定预算 $1.5\times10^{-3}$、固定预算 $1.5\times10^{-4}$、以及 $2\times10^{-4}\to2\times10^{-5}$ 退火。第 700–800 个记录步的原始评估均值 Pass@1 依次为 **70.30%、63.43%、71.03%、71.47%**；对应平均回答长度为 **774、976、858、832 token**。过大的固定预算出现性能退化；小预算与退火预算维持较高尾段准确率。不过无控制组也保持稳定，而且退火组的初始预算与固定小预算不同，不能从这组两种子实验断言 KL 控制是稳定的必要条件，或将差异仅归因于“退火”本身。（Sec. 3.4.3，Fig. 6，PDF pp. 9–10）

![原文 Fig. 6：不同 KL 控制下的长程准确率与回答长度](figures/long_horizon.png)

原图 Fig. 6（PDF p. 9）：粗线为两种子均值、细线为单次运行，横轴是*记录步*；展示曲线经过五次评估移动平均，上述尾段数值则取未平滑的原始评估。

## 五：结论与读论文时要保留的区分

LoGRA 的有效设计单元是“直接累计 $S=GA^\top$，用 $UA$ 更新完整权重，再按预测 KL 缩放”。实验最强的证据是单卡训练显存降低和同一八卡节点上 27B 可训练；1.5B/7B 的端到端更新吞吐仅小幅变化。预测 KL 的消融支持**合适的小预算**，而非“加上控制必然稳定”。

还要分清**可实现的同步机制**与**实际测量范围**：方法节给出通过低秩因子与随机种子同步生成端的方案，但附录 D.1 明确写主实验两组均使用 *full-weight policy synchronization*。因此 Table 1 的吞吐和显存结果不能当作低秩同步节省带宽的实测数字；这部分仍需要单独通信测量。主比较的 RowAdam 与 Dense Adam 也不是同一优化器，准确率差异不能只归因于 sketch。（Sec. 2.2、Table 1、Appendix B/D.1，PDF pp. 3、7、14–16）