跳到正文
原文
omarsar0· @omarsar0 · X·· 7 小时前AI 评分57

NVIDIA 论文提出 LoGRA:全参数 RL 训练显存低于 LoRA

AI 导读

DAIR.AI 推荐 NVIDIA 关于显存高效后训练 RL 的论文 LoGRA,其核心结论是用 RL 更新模型全部权重时显存仍低于 LoRA。1.5B 模型上 LoGRA 平均训练显存 7.2 GiB,rank-256 LoRA 需 13.2 GiB,Pass@4 分别为 81.8% 和 81.4%,Pass@1 低约 2 个点。

正文 · 原文

Recommended NVIDIA paper on memory-efficient post-training RL.

A very interesting technique for scaling RL for LLMs.

The main takeaway is that you can update all of a model's weights with RL and still use less memory than LoRA. At 1.5B, LoGRA averages 7.2 GiB of training memory, while rank-256 LoRA needs 13.2 GiB. Pass@4 is about the same (81.8% and 81.4%), and Pass@1 is about 2 points lower.

If GPU memory is what stops you from running RL on a bigger model, this method is worth a look.

LoGRA never builds a full gradient for the weight matrices it compresses.

As microbatches accumulate, each weight gradient is compressed with a rank-256 random projection that is redrawn every step, and the update is rebuilt from two small matrices. Before each step, it predicts how much the policy will change and scales the step to stay within a KL budget.

Average training memory per GPU drops 45.7% at 7B (31.8 to 17.3 GiB), with peak accuracy close to dense Adam (72.3% and 72.5% Pass@1).

It also trains a 27B model for more than 1,100 steps on a single eight-GPU node, where dense Adam runs out of memory.

Paper: arxiv.org/abs/2610.06647

Chat with Paper: academy.dair.ai/papers/logra…

来源:omarsar0 · x.com