大模型强化学习显存骤降 45.7%！LoGRA 论文与开源库发布：首创低秩梯度草图攻克 27B 单节点训练瓶颈

核心背景与行业痛点

随着 DeepSeek-R1、OpenAI o1 等强化学习模型展现出惊人的长思维链推理能力，RL 后训练（RL Post-Training）已成为大模型提升数学、代码与逻辑表现的核心法宝。然而，强化学习训练的显存开销远远高于常规的监督微调（SFT）：

1. 多重状态共存的显存灾难：在标准 Actor-Critic 或 PPO 范式中，显存不仅要装载 Policy 模型、Reference 模型以及 Reward 模型的权重，还需要为优化器（如 AdamW）保留多份全参数精度的梯度与一阶/二阶矩动量缓冲区；
2. 硬件门槛过高导致科研断层：主流中小团队通常仅配置单台 8 卡 GPU 服务器（如 8×A100 或 8×H100 80GB）。当尝试对 27B 以上规模的开源模型进行全量强化学习训练时，全尺寸稠密梯度瞬间撑爆显存，频发 CUDA Out-Of-Memory 报错，迫使研发人员只能使用 LoRA 等次优降级方案；
3. 多机分布式策略同步通信带宽瓶颈：跨节点全参数梯度的频繁 All-Reduce 通信耗时巨大，极易拖垮整体训练吞吐。

架构亮点与底层机制

针对强化学习特有的显存与通信瓶颈，研究团队提出了基于低秩梯度草图的全新训练方法 LoGRA（arXiv:2610.06647，开源库 github.com/skzhang1/labs-molt）：

1. 低秩梯度草图（Low-Rank Gradient Sketches）：研究团队深入分析发现，强化学习反向传播过程中的有效学习信号在频谱上具有极高的低秩特性。LoGRA 摒弃了为每个参数分配稠密梯度张量的传统做法，而是在反向传播过程中将梯度动态投影并累积到紧凑的低秩矩阵草图中。优化器直接在低秩空间中维持状态并执行参数更新，显存占用呈断崖式下降；
2. 预测 KL 步长控制（Predicted-KL Step Control）：压缩梯度的最大技术风险在于方向失真引发策略坍塌（Policy Collapse）。LoGRA 独创了预测 KL 散度步长控制机制：在正式将参数更新应用到模型权重之前，前向预估当前更新量会导致的 Policy Token 概率分布漂移；一旦预估 KL 偏离超过安全阈值，系统自动缩放并微调更新步长，从数学上保证策略演进的平滑收敛；
3. 原生集成于 Molt 强化学习库：该算法已无缝集成至轻量、PyTorch 原生的 Molt 训练库，兼容 FSDP（完全分片数据并行）与 DeepSpeed ZeRO 架构，开发者无需重写现有训练流程即可一键启用。

权威 Benchmark 与实测跑分对比

在涵盖 GSM8K 数学题库、MATH 高难度竞赛以及 HumanEval 代码推理基准的严格实测中：

1. 显存消耗直降 45.7%：在同等超参数与任务设置下，LoGRA 相比经典 Dense AdamW 优化器实现了高达 45.7% 的平均训练显存节省；
2. 单机 8 卡成功突破 27B 模型训练壁垒：在单台 8 卡 GPU 节点上，标准 Dense AdamW 在训练 27B 参数大模型时瞬间触发 OOM 终止；而启用 LoGRA 后，训练进程平稳运行超过 1,100 步无任何显存报错或数值下溢；
3. 零精度损失与卓越收敛稳定性：在完成数千步强化学习对齐后，LoGRA 训练的推理模型在准确率得分上与高成本的全量 AdamW 达到统计同等水平（±0.3% 以内），且在长程训练中表现出极佳的抗遗忘与 KL 约束稳定性。

开发者实战落地与开箱指南

LoGRA 代码已作为 Molt 框架的核心组件在 GitHub（github.com/skzhang1/labs-molt）全面开放。对于正在探索长思维链推理强化学习、RLVR 以及数学代码对齐的科研团队与 AI 企业，LoGRA 带来了立竿见影的降本增效收益：不再需要向多机集群申请成倍的昂贵节点，利用单台 8 卡机器即可直接展开 20B~30B 级别大模型的全量强化学习训练，极大平民化了大模型前沿强化学习的技术门槛。