LoGRA：用低秩梯度草图扩展LLM强化学习 LoGRA: Scaling LLM Reinforcement Learning with Low-Rank Gradient Sketches
低秩梯度草图+预测KL步长控制，RL训练内存省45.7%，27B单机8卡稳定训练
前置知识
RL 后训练（PPO + 可验证奖励）
在预训练/SFT 模型上用强化学习继续优化。典型循环：rollout 引擎对 prompt 采样回复 → 验证器打分（如数学答案对错，即二值可验证奖励）→ 估计优势（本文用全局运行奖励基线）→ 反向传播更新策略 → 把新权重同步回 rollout 引擎。本文用刻意简单的 PPO 设置，以便隔离方法本身的效果。
LoGRA 的训练内存瓶颈和策略同步通信开销都发生在这一循环中，理解各阶段的数据流才能看清压缩点在哪里。
Adam 优化器状态与训练内存构成
Adam 为每个可训练参数维护一阶动量 $m$ 和二阶动量 $v$ 两个缓冲区，通常以 FP32 存储。以 7B Qwen 为例：BF16 权重约 14 GB，而 Adam 两个动量缓冲就要 56 GB，还没算梯度和激活。训练内存远大于推理内存，这正是 RL 后训练的硬件门槛来源。
LoGRA 压缩的正是梯度累积缓冲和优化器状态这两块最大的开销，必须知道内存在哪里才能理解节省从何而来。
随机投影与低秩梯度草图
用固定随机矩阵 $A \in \mathbb{R}^{r \times k}$（元素为独立 $\pm 1/\sqrt{r}$，即 Rademacher 分布）把高维对象压缩到 $r$ 维，其思想源于 Johnson–Lindenstrauss 引理：随机投影能近似保持内积与距离。对梯度 $G \in \mathbb{R}^{d\times k}$，草图 $S = GA^\top \in \mathbb{R}^{d\times r}$ 只需存 $dr$ 个数而非 $dk$ 个，例如 $d=k=4096, r=64$ 时 fp32 存储从 64 MiB 降到 1 MiB。
低秩草图是 LoGRA 的核心数据结构，重建 $\hat{G}=SA$ 及其近似质量决定了方法成立与否。
KL 散度与信赖域方法（TRPO）
KL 散度 $D_{KL}(\pi_\theta \| \pi_{\theta+\alpha D})$ 衡量新旧策略在下一 token 概率分布上的差异。TRPO 的信赖域思想指出：只有当策略更新足够小（KL 在预算内）时，局部代理目标的改进才能转化为真实性能提升；过大的更新会破坏已学到的行为。
论文的 predicted-KL 步长控制就是信赖域思想的轻量化在线实现，用于约束低秩草图更新的幅度，是稳定长程训练的关键。
系统级省内存技术（FSDP / 激活检查点 / offload）
FSDP 把参数、梯度、优化器状态分片到多张 GPU；激活检查点在前向时丢弃部分激活、反向时重算以换内存；CPU offload 把部分训练状态挪到主机内存。它们的共同点是把状态搬运或摊薄，但状态总量没有减少。
论文明确将 LoGRA 定位为与这些系统级手段互补的算法级路线——直接压缩梯度与更新的表示本身，两者正交可叠加。
研究动机
RL 已成为 LLM 能力提升的关键驱动力，但其巨大的内存需求是更广泛采用的硬性障碍，且瓶颈远不止参数本身。论文给出具体账本：7B 的 Qwen 模型 BF16 权重约 14 GB，而 Adam 的两个 FP32 动量缓冲再加 56 GB，这还没算梯度和激活。结果是“能推理的硬件未必能训练”：在单机 8×H100 80GB 的配置下，27B 模型的 dense Adam 训练在第一次分配 Adam 状态时就 OOM。现有应对手段各有局限：系统级优化（FSDP 分片、梯度检查点、CPU offload）只是搬运或摊薄状态，不减少总量；参数高效方法（LoRA）冻结预训练权重、只训练低秩适配器，牺牲了全参数更新的表达能力；而已有梯度压缩工作（GaLore、Flora）主要面向预训练/SFT，需要先形成完整梯度再投影。
本文的目标是本文目标是让 RL 后训练在保持全参数可训练、任务性能不损失的前提下，大幅降低两处开销：一是训练侧的梯度累积缓冲与优化器状态内存，二是训练器与 rollout 引擎之间的策略同步通信量。量化目标是：在 1.5B/7B 上实现可观内存节省且 MATH-500 精度持平或更好，并把原本 dense Adam 完全无法训练的 27B 模型，在单机 8 卡节点上稳定训练超过 1,100 步，使此前因内存不可行的 RL 训练变得实际可行。
与已有工作不同的是，论文的独特切入是一个关于 RL 学习信号的观察：采用二值结果奖励时，每个采样回复只携带一个正确性比特，从奖励中学习所需的信息可能远小于一个完整梯度所能表达的信息量，因此学习信号可以被紧凑表示承载。基于此，LoGRA 不是“先算完整梯度再压缩”，而是在反向传播内部直接累积投影后的草图，完整梯度从头到尾不被构造；更进一步，把同一压缩表示复用于 rollout 策略同步（只需传输草图因子和投影种子），把压缩从训练内存扩展到通信；再配上 predicted-KL 步长控制来约束更新幅度。这一“RL 训练-生成循环专用”的组合是以往预训练导向的梯度压缩工作没有覆盖的。
核心方法
直觉上，RL 的学习信号远小于全梯度维度：二值结果奖励下每条回复只携带一个正确性信号，因此可用低秩草图承载。技术上，对注意力/MLP 投影矩阵 $W \in \mathbb{R}^{d\times k}$，选秩 $r \ll \min(d,k)$ 和固定随机投影 $A \in \mathbb{R}^{r\times k}$（默认 Rademacher $\pm 1/\sqrt{r}$），定义草图 $S = GA^\top \in \mathbb{R}^{d\times r}$，存储从 $dk$ 降到 $dr$（$d=k=4096, r=64$ 时 fp32 从 64 MiB 降到 1 MiB）。实现上不先算 $G$ 再压缩，而是把权重梯度计算中的输入先投影、再乘反向导数，以 fp32 直接逐 microbatch 累积进 $S$。更新时用 $\hat{G}=SA$ 代替 $G$：$W \leftarrow W - \alpha\eta U A$，其中 $U$ 是可选经优化器调整的草图，$\alpha$ 由 predicted-KL 控制器给出。每步从新种子重采样 $A$，使 successive 更新张成不同子空间；同一 $\alpha\eta U$ 加种子发给 rollout 引擎重建 $A$ 并施加相同更新，通信量从 $O(dk)$ 降到 $O(dr)$。
三个环环相扣的创新。其一，直接草图累积：压缩发生在反向传播内部，完整梯度从未被实体化，与 GaLore/Flora“先形成 $\nabla_W \mathcal{L}$ 再投影”有本质区别。其二，压缩表示复用于策略同步：rollout 端由种子重建 $A$，只传 $\alpha\eta U$，把 RL 特有的 rollout 同步通信一并压缩到 $O(dr)$。其三，predicted-KL 步长控制：利用 KL 在当前策略处一阶斜率为零、曲率可由 token 分数变化率方差表示的性质，得二阶近似 $\mathrm{KL} \approx \alpha^2 q(D)$，其中 $q(D) = \frac{1}{2|\mathcal{P}|T^2}\sum_{s}\mathrm{Var}_{v\sim p(s)}[\dot{z}_v(s)]$，只需在少量上下文上算 logit 沿更新方向的变化率，无需密集梯度或曲率矩阵；再按 $\alpha = \min(\alpha_{\max}, \sqrt{\delta/q(D)})$ 缩放满足预算 $\delta$（$\alpha$ 减半、预测 KL 降为 1/4），把 TRPO 信赖域变成轻量在线探针。另配 RowAdam（每行仅一个二阶矩标量，$\beta_2=0.95$）与 mismatch subtraction 校正 rollout/trainer 概率差。
方法步骤详情
完整流程（Algorithm 1）：(1) 准备——从当前策略对 128 个 prompt 各采样一条回复，构造可验证奖励的 PPO 损失 $\mathcal{L}$。(2) 采样投影——对每个被压缩矩阵 $W$，从新种子采样 $A \in \mathbb{R}^{r\times k}$ 并把 $S$ 清零。(3) 累积草图——$\theta$、$A$ 固定下逐 microbatch 计算 $S \leftarrow S + (\nabla_W \mathcal{L}_m)A^\top$，实现上先投影输入再做权重梯度乘法，无需完整梯度缓冲。(4) 选择更新——SGD 变体取 $U=S$；RowAdam 维护每行二阶矩 $v_{t,i} = \beta_2 v_{t-1,i} + (1-\beta_2)\frac{1}{r}\sum_j (S_t)^2_{ij}$，行缩放后匹配 Frobenius 范数。(5) 步长控制——抽少量上下文估计 $q(D)$，取 $\alpha=\min(\alpha_{\max},\sqrt{\delta/q(D)})$；主实验预算余弦退火 $2\times10^{-4}\to2\times10^{-5}$（$H=300$），并做 mismatch subtraction（下限 50%）。(6) 更新与同步——trainer 端 $W \leftarrow W - \alpha\eta U A$ 直接合并进权重；发送 $\alpha\eta U$ 与种子，rollout 重建 $A$ 施加相同更新后进入下一步。
技术新颖性
与 LoRA 的本质区别：LoRA 改变参数化——冻结 $W$ 只训适配器，更新被永久限制在低秩流形上；LoGRA 保持 $W$ 全尺寸可训练，每步更新秩至多 $r$，但每步刷新投影使多步累积不被锁死在固定子空间。与 GaLore/Flora 的区别：后两者面向预训练/SFT，需先实体化完整梯度再投影，且不涉及 RL 循环；LoGRA 在反向传播中直接累积草图，并首次把压缩延伸到 rollout 同步通信（$O(dr)$ 加种子替代全量权重传输），这是 RL 训练-生成循环特有的贡献。predicted-KL 也是针对草图更新的新控制机制：附录 A 严格推导了 KL 的二阶展开与缩放律（$\alpha$ 减半、预测 KL 降为四分之一），并用 token 分数变化率的方差 $\mathrm{Var}_{v\sim p(s)}[\dot{z}_v(s)]$ 做无曲率矩阵的估计。RowAdam 则用每矩阵 $d$ 个标量近似 Adam 的 $dk$ 个二阶矩，保留自适应缩放又避免一阶动量存储。
实验结果
主实验（DAPO-Math-7.5K 训练、MATH-500 评估、3 seeds）：1.5B 上 LoGRA 峰值 Pass@1 67.87% vs Dense 63.77%（+4.10 pp）、Pass@4 82.47% vs 80.33%，平均内存 9.18→7.18 GiB（−21.8%），吞吐 55.44 vs 54.28 updates/h。7B 上精度持平（Pass@1 72.33% vs 72.48%，Pass@4 85.13% vs 85.33%），平均内存 31.82→17.29 GiB（−45.7%），峰值 31.85→22.26 GiB。27B：Dense 首次分配 Adam 状态即 OOM，LoGRA 以平均 51.54/峰值 53.88 GiB 完成训练，Pass@1 71.52%、Pass@4 81.87%，吞吐 6.99 updates/h。长程实验：Qwen3.8-27B 在 Reasoning-Gym Hard 八任务（38,948 训练题）上单机 8 卡训练 1,100+ 步，macro 验证分从 39.69% 升至 62.94%（峰值 65.52%@step 1,060），回复长度 8,312→4,705 tokens，截断率 34.72%→12.95%。消融：秩影响最大（$r$ 从 4 到 256 提升 3.02–3.95 pp；$r=256$ 时固定 Rademacher 68.50%、刷新 68.31%、Gaussian 67.42%），重构余弦从 0.046（$r=4$）升至 0.348（$r=256$），刷新/分布影响 <0.0014。KL 消融（2 seeds）：无控制 70.30%、大预算 $1.5\times10^{-3}$ 退化至 63.43%、小预算 $1.5\times10^{-4}$ 71.03%、退火最佳 71.47%。vs LoRA（1.5B）：内存 7.18 vs 13.21 GiB，Pass@1 68.47 vs 70.45%，Pass@4 81.80 vs 81.40%，11.31 vs 10.63 s/步。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| MATH-500（Qwen2.5-Math-1.5B） | 峰值 Pass@1 | 67.87% | Dense Adam：63.77% | +4.10 pp，且平均内存降 21.8%（9.18→7.18 GiB） |
| MATH-500（Qwen2.5-Math-7B） | 峰值 Pass@1 / Pass@4 | 72.33% / 85.13% | Dense Adam：72.48% / 85.33% | 精度持平（−0.15/−0.20 pp），平均内存降 45.7%（31.82→17.29 GiB） |
| MATH-500（Qwen3.8-27B） | 峰值 Pass@1 / Pass@4 | 71.52% / 81.87% | Dense Adam：OOM，无法训练 | 同等硬件下从不可行变为可训练 |
| Reasoning-Gym Hard 27B 长程训练 | held-out macro 验证分 | 39.69%→62.94%（峰值 65.52%@1,060 步） | 训练起点 39.69% | +23.25 pp，1,100+ 步稳定无崩塌，截断率 34.72%→12.95% |
| Qwen2.5-Math-1.5B vs LoRA（rank 256） | 平均训练内存 | 7.18 GiB | LoRA + AdamW：13.21 GiB | −45.6%；Pass@4 略优（81.80% vs 81.40%），Pass@1 低 1.98 pp |
| 单步训练调用耗时（1.5B） | 秒/update | 11.31 s | LoRA：10.63 s | LoGRA 慢约 6.4%，精度内存换速的代价很小 |
局限与改进
作者承认的局限：与 LoRA 的对比中两方法使用不同的计数器、优化器与调度，无法单独归因适配器参数化本身的差异，也未确立 LoGRA 在精度或速度上的优势；KL 控制消融仅 2 seeds，只能说明校准良好的 KL 控制有益，不能证明它是稳定性的必要条件，退火带来的增益解释也不唯一；投影刷新在所测短程设置中无一致收益，Gaussian 投影在秩 64/256 略差（低于 $2\times10^{-4}$）。我自己的观察：$r=256$ 时梯度重构余弦仅 0.348，意味着约九成的梯度信息被丢弃，为何 RL 仍能学得与 dense 相当缺乏理论解释；评测集中在数学推理（MATH-500、Reasoning-Gym），向代码、多轮 agent 等任务的外推未验证；27B 吞吐仅 6.99 updates/h，绝对训练成本仍很高；7B 峰值内存标准差高达 ±4.98 GiB，说明内存优势会随负载波动；实验中没有与 GaLore/Flora 在 RL 场景的直接实证对比；KL 预算的起点/终点等超参需手工设定，跨模型规模的可迁移性不明。
独立分析的弱点
第一，低秩近似的理论空白：重构余弦在 $r=256$ 仅 0.348，约九成梯度信息被丢弃，“二值奖励下低秩学习信号已足够”只是直觉，未形式化。改进方向：建立 sketch 更新下策略梯度近似误差与性能损失的定量界，或按层/按步引入误差反馈补偿投影残差。第二，验证域过窄：全部实验为数学推理与 Reasoning-Gym，多轮 agent、代码、长上下文未测，而长序列会放大激活内存并改变 KL 探针的采样需求；应在 Polar 类 agentic RL 框架上复验并重设计探针上下文选择。第三，KL 预算是手工超参且敏感，大预算 $1.5\times10^{-3}$ 直接导致退化（63.43%）；可引入基于 realized-KL 反馈的自动预算调节。第四，7B 峰值内存标准差达 ±4.98 GiB、27B 吞吐仅 6.99 updates/h，说明与激活重算、microbatch 分桶、异步同步的系统协同仍有优化空间。第五，KL 消融仅 2 seeds，统计效力弱，关键结论应提高到 5 seeds 以上并做显著性检验。第六，27B 长程实验因效率退回有限差分探针，JVP 估计器（Equation 6）与有限差分的精度/成本权衡未量化，估计器本身需要工程化。
未来方向
作者层面提出的方向：把 LoGRA 推向更大规模与多节点训练，进一步压缩 rollout 同步通信；改进 predicted-KL 探针的效率（正文 27B 长程实验因效率改用有限差分，Equation 6 的 JVP 估计器是设计目标）；结合异步 RL 框架减少同步等待。基于成果可延伸的方向：其一，与 GaLore 式投影预条件子、LoRA 混合训练组合，探索“低秩更新 + 低秩优化器状态”的最优配比；其二，把压缩同步推广到多节点拓扑感知通信（分层 all-reduce、seed 共享广播）；其三，自动秩分配与自动 KL 预算调度（按层梯度重构误差在线调 $r$、按 realized-KL 调 $\delta$）；其四，理论侧刻画二值奖励 RL 中低秩梯度子空间 sufficient 的条件；其五，在 agentic RL（工具调用、多轮）与偏好优化（DPO 类）上验证泛化，那里 rollout 同步通信占比更高，LoGRA 的通信收益会更突出。
复现评估
复现条件较好。开源方面：论文声明代码发布于 Molt 库（按论文原文），作者之一 Jian Hu 是 OpenRLHF 作者，工程可信度较高。数据全部公开：DAPO-Math-7.5K 由 HF 上 DAPO-Math-17k 按序去重取前 7,500 个唯一 prompt 构造，评估用 MATH-500 全部 500 题，长程实验用 Reasoning-Gym 八任务混合（表 3 给出精确题量 38,948/193）。协议透明：附录 C/D 完整给出评估设置（每题 4 条回复、温度 0.7、top-p 0.7）、全部超参（表 4：rank-256 Rademacher、每步刷新、RowAdam $\beta_2=0.95$、预算退火公式 16、mismatch subtraction 公式 17）、种子（1/2/42）与 3-seed 均值±标准差报告。算力门槛：主实验需单机 8×H100 80GB + 128 CPU 核心（4 卡 FSDP 训练 + 4 卡生成），27B 实验对多数实验室不可及，但 1.5B 消融在更小配置下可复现。综合难度中高：核心工程难点是在反向传播中定制投影累积（类似 Flora 的 hook），其余文档充分。
论文图表
跨模型规模（1.5B/7B/27B）的训练内存柱状图：深色为平均内存、浅色延伸为峰值内存（GiB/GPU，GPU 容量 79.2 GiB），并排对比 Dense Adam 与 LoGRA。1.5B 平均内存 9.18→7.18 GiB（省 21.8%），7B 31.82→17.29 GiB（省 45.7%），27B Dense Adam 标记为 OOM 而 LoGRA 为 51.54/53.88 GiB，纵轴在 34–48 GiB 处断轴。
这是全文的摘要图，一图看清 LoGRA 的核心卖点：内存节省随模型规模扩大，且在 27B 处实现从不可行到可行的质变，是理解论文动机与价值的入口。
每次更新的峰值分配内存随训练步数的变化曲线（3 seeds 均值±标准差），并标注时间均值：1.5B Dense 9.18 vs LoGRA 7.18 GiB（省 21.9%），7B 31.81 vs 约 16 GiB（省 45.6%），27B LoGRA 稳定在约 51.54 GiB（距 79.2 GiB 容量尚有余量）而 Dense OOM。
展示内存节省在整个训练过程中持续稳定而非偶发低谷，并直接对应 Table 1 的平均/峰值内存数字，是效率声明的核心证据。
方法主伪代码：每步先对每个被压缩矩阵从新种子采样 $A$ 并清零 $S$；逐 microbatch 直接累积 $S \leftarrow S + (\nabla_W \mathcal{L}_m)A^\top$（注释强调 Direct sketch accumulation，不形成完整梯度）；经 $U \leftarrow S$ 或 sketch 优化器、predicted-KL 选 $\alpha$ 后，trainer 端合并 $W \leftarrow W - \alpha\eta U A$，并把 $\alpha\eta U$ 与种子发给 rollout 端重建 $A$ 施加相同更新。
方法的精确规范：六步循环涵盖了直接累积、投影刷新、KL 控制与压缩同步四个创新点的落地方式，是实现者的首要参考。