强化学习训练内存究竟花在哪里
论文把内存问题界定为比存储权重更大的问题。对一个 7B 参数的 Qwen 模型，BF16 权重约占 14 GB，而 Adam 的两个 FP32 动量缓冲又要再加 56 GB，这还没算梯度和激活值。正是这一差距，使得能跑起模型的硬件仍可能无法通过强化学习改进它。
LoGRA 针对的是梯度和优化器这一侧，而不是权重。它让原始权重矩阵保持可训练，不同于 LoRA 式适配器冻结预训练权重、只学习小型低秩分支。论文给出的直觉是：二值结果奖励对每条生成回复只给出一个正确性信号，因此从奖励中学习所需的信息可能远少于完整梯度所表示的信息。
草图如何取代完整梯度缓冲
对一个尺寸为 d×k 的权重矩阵 W，LoGRA 选取远小于 min(d,k) 的秩 r，以及一个尺寸为 r×k 的随机投影矩阵 A。它不累积完整的 d×k 梯度，而是累积尺寸为 d×r 的草图 S = G Aᵀ。投影 A 与当前批次独立采样，在累积过程中保持固定，且不被学习；默认情况下其元素为独立随机符号，按 1/√r 缩放。
实现上，它把每个微批次的投影贡献直接以 fp32 累积进 S，做法是在与反向导数相乘之前，先对权重梯度计算中用到的输入做投影。被压缩的矩阵不会形成完整梯度缓冲。更新形式为 W ← W − αηUA，其中 U 在基础 SGD 变体中等于 S，α 是由步长控制器选择的乘子。更新直接合并进权重，没有常驻的适配器分支。
内存账很具体：当 d = k = 4096、r = 64 时，梯度累积器的 fp32 存储从 64 MiB 降到 1 MiB。论文明确指出这并非全部训练占用——模型权重、缓存的投影、激活值和其他缓冲仍占内存，带状态的草图优化器还会增加自己的存储。
同一份草图也压缩策略同步开销
压缩不只作用于训练内存，也作用于通信。同步时，rollout 引擎用种子重新生成 A，并应用收到的因子 αηU。传输负载是 O(dr) 个数值加一个种子和元数据，而不是完整的 d×k 更新。论文把这描述为在梯度累积、权重更新和策略同步之间复用同一份紧凑表示。
预测 KL 控制防止过大更新破坏学习
仅有压缩还不够。论文指出，参数空间中很小的更新仍可能显著改变策略的输出概率，而过大的更新会突然偏移策略、破坏此前学到的行为。因此 LoGRA 加入预测 KL 步长控制，在应用更新之前估计该更新会改变多少下一词元概率。
该估计从当前批次中抽取一小部分回复，在同一温度下比较当前模型的词表分布与提议更新后模型的词表分布。论文的近似预测步长呈二次依赖：平均预测 KL 按 α²q(D) 缩放，因此把 α 减半会把预测 KL 降到四分之一。给定 KL 预算 δ，控制器设定 α = min{α_max, √(δ/q(D))}。当 q(D) = 4δ 时，提议被缩放到至多原尺寸的一半，而 α_max 上限可防止预测 KL 很小时更新被过度放大。
论文对控制对象很谨慎：它调节的是预测变化，而实际 KL 可能不同，因为该估计是局部近似。预算可以是常数，也可以退火，退火预算会逐步降低 δ，使训练后期允许的策略变化更小。
内存节省随模型规模增大
在多项推理任务上，LoGRA 在 1.5B 规模把平均训练内存降低 21.8%，在 7B 规模降低 45.7%，且不牺牲任务性能。7B 的报告数值是 LoGRA 每张训练 GPU 平均 17.29 GiB、峰值 22.26 GiB，而稠密 Adam 为平均 31.82 GiB、峰值 31.85 GiB。1.5B 的对比是 LoGRA 平均 7.18 GiB、峰值 8.66 GiB，稠密 Adam 平均 9.18 GiB、峰值 9.19 GiB。
27B 的结果是可行性主张，而不是实测内存数值。LoGRA 让 27B 模型在单台八卡节点上稳定训练超过 1,100 步，而稠密 Adam 基线在同一硬件分配下耗尽内存。论文把 27B 的稠密 Adam 柱标为在 GPU 容量处 OOM，而非实测值，因此那里的比较是可行性，不是内存比例。
投影秩比基刷新更重要
设计选择实验使用 Qwen2.5-Math-1.5B 在 DAPO-Math-7.5k 上运行，秩取 r ∈ {4, 16, 64, 256}。把秩从 4 提高到 256，在各配置下使 MATH-500 pass@1 提升 3.02 到 3.95 个百分点。在秩 256 时，固定 Rademacher、刷新 Rademacher 和刷新 Gaussian 投影分别达到 68.50%、68.31% 和 67.42%，且它们的排序随秩变化。
论文报告刷新投影基没有一致收益，可能是因为固定的低秩子空间在这段较短训练周期内已捕获足够的任务相关信号。梯度重建诊断指向同一方向：稠密梯度与其重建之间的平均余弦相似度从秩 4 的约 0.046 升到秩 256 的 0.348，而对应层–秩单元之间的最大差异，固定与刷新 Rademacher 之间仅为 0.0014，刷新 Rademacher 与 Gaussian 之间仅为 0.0008。Gaussian 在秩 64 和 256 的下降幅度低于 2×10⁻⁴。主导单微批次重建质量的是秩，而不是基刷新或元素分布。
这些节省能说明什么、不能说明什么
与 LoRA 的比较受论文自身条件限定：LoRA 学习低秩权重适配器，而 LoGRA 压缩梯度并更新权重；报告的比较使用秩 256 的 Qwen2.5-Math-1.5B 在 DAPO-Math-7.5k 上训练的前 300 步，LoRA 用 AdamW 和恒定 2×10⁻⁴ 学习率，LoGRA 用从 2×10⁻⁴ 退火到 2×10⁻⁵ 的预测 KL 预算。
设计选择证据同样有范围：秩和刷新实验在 1.5B 模型上、较短周期内进行，论文自己也指出固定子空间的结果可能反映的是这一较短周期，而非普遍性质。45.7% 是 7B 规模的平均内存节省，不是峰值内存保证；27B 的主张是关于在稠密 Adam 无法启动的情况下完成 1,100 多步，而不是该规模下实测的内存下降。
这篇论文确立的实际优先级是一组配对，而非单一技巧：草图负责去掉完整梯度缓冲，而预测 KL 步长控制负责让压缩后的更新不破坏策略稳定。读者在把 LoGRA 与基于适配器或系统级的内存节省方案比较时，应按论文实际测量的轴来比——给定模型规模和秩下每张 GPU 的平均训练内存——并把 27B 结果视为在一种八卡分配下可行的证据，而不是普遍缩放规律。
本文引用的文献
LoGRA: Scaling LLM Reinforcement Learning with Low-Rank Gradient Sketches
LoGRA 用低秩草图替代选定权重矩阵的完整梯度，并加入预测 KL 步长控制，在 7B 规模把平均训练内存最多降低 45.7% 且不牺牲推理性能，并让 27B 模型在单台八卡节点上训练超过 1,100 步，而稠密 Adam 会耗尽内存。