PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2610.06647v1cs.CL

TL;DR

RL post-training is memory-intensive, with optimizer states and other training data adding substantially to model-weight storage. LoGRA uses low-rank gradient sketches for updates and policy synchronization, paired with predicted-KL step control. It reduces average memory by up to 45.7% without sacrificing performance and sustains 27B training for over 1,100 steps on one eight-GPU node where Dense Adam runs out of memory.

  • Problem

    RL post-training’s memory requirements, including gradients and optimizer states, remain a barrier to broader adoption.

    Source: §1 · pp. 1–2

  • Method

    LoGRA accumulates low-rank gradient sketches for model updates and policy synchronization, while predicted-KL step control scales proposed updates to limit policy change.

    Source: §1 · p. 2

  • Results

    LoGRA reduces average training memory by up to 45.7% without sacrificing task performance and sustains 27B training for over 1,100 steps on one eight-GPU node where Dense Adam runs out of memory.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The findings identify compact gradient representations as a practical lever for lowering the hardware barrier to RL post-training.

    Source: §5 · p. 10

Abstract

from arXiv · show

Reinforcement learning (RL) has greatly advanced the capabilities of large language models (LLMs), but its memory demands remain a barrier to broader adoption. We introduce LoGRA, an approach to RL post-training that reduces memory by retaining useful learning signals in low-rank gradient sketches. These compact representations support both model updates and efficient policy synchronization. To prevent overly large updates from disrupting learning, we complement gradient compression with predicted-KL step control, which estimates policy changes before applying each update and adjusts its magnitude accordingly. Across reasoning tasks, LoGRA reduces average training memory by up to 45.7\% without sacrificing performance. It also enables stable training of a 27B-parameter model for over 1,100 steps on a single eight-GPU node, where dense Adam runs out of memory, making previously memory-infeasible RL training practical. Code is available in the \href{https://github.com/skzhang1/labs-molt/tree/logra/examples/scripts/logra}{Molt library}.

1. Introduction

RL post-training can demand far more memory than model weights alone, limiting who can train LLMs with feedback. LoGRA compresses gradient signals and controls update size; evaluations report lower memory use and sustained 27B training where Dense Adam runs out of memory.

Motivation

  • RL memory demands extend beyond weights: Adam’s two FP32 moment buffers add 56 GB to the approximately 14 GB of BF16 weights in a 7B Qwen model.Gradients and activations add further memory demands.Source: §1 · pp. 1–2

Related work

  • Existing responses include sharding, checkpointing, CPU offloading, parameter-efficient methods, and gradient compression that retains full-weight training.LoRA freezes pretrained weights and trains small low-rank adapters.Source: §1 · p. 2

LoGRA

  • LoGRA accumulates low-rank gradient sketches during backpropagation, merges their updates into model weights, and reuses the sketches for rollout-policy synchronization.The approach targets both training memory and communication costs.Source: §1 · p. 2

Predicted-KL step control

  • Predicted-KL step control estimates how a proposed update changes next-token probabilities and scales the update to stay within a prescribed KL budget.It complements compression because an overly large update can disrupt previously learned behavior.Source: §1 · p. 2

Results

  • Average training memory falls by 21.8% at 1.5B and 45.7% at 7B without sacrificing performance; LoGRA sustains 27B training for over 1,100 steps where Dense Adam OOMs.The 27B run uses a single eight-GPU node.Source: §1 · p. 2

2. Method

LoGRA replaces full gradient buffers with compact sketches, uses them to update selected weights and synchronize the rollout policy, and scales proposed steps using predicted KL. Its method comprises low-rank gradient sketching, RL post-training integration, and predicted-KL step control.

2.1. Low-rank gradient sketches

  • For a d×k weight matrix, LoGRA projects gradients into a d×r sketch and accumulates projected microbatch contributions directly, avoiding a full gradient buffer.The sketch stores dr rather than dk values; the projection is fixed during accumulation and is not learned.Source: §2.1 · pp. 2–3

2.2. Applying LoGRA to RL post-training

  • LoGRA applies per-matrix sketches to attention and MLP projections, merges updates into weights, and refreshes projections so successive updates can use different subspaces.Each matrix has its own projection and sketches; the method does not retain a persistent adapter branch.Source: §2.2 · p. 3
  • Synchronization regenerates each projection from its seed and sends the update factors; at d=k=4096 and r=64, fp32 accumulator storage falls from 64 MiB to 1 MiB.The synchronization payload is O(dr) values plus a seed and metadata, rather than a full d×k update.Source: §2.2 · p. 3

2.3. Step-size control with predicted KL

  • Predicted KL compares current and proposed next-token probabilities on a small subset of generated-response contexts at the same temperature.The proposal is an update to all parameters, with the compressed-weight component formed from the sketch factors and nominal step size.Source: §2.3 · p. 4
  • The controller scales the proposal to fit a KL budget and caps the multiplier; realized KL can differ because the estimate is a local approximation.Predicted KL scales quadratically with the update multiplier, and an annealed budget permits less policy change later in training.Source: §2.3 · p. 5

3. Experiments

Across reasoning tasks and model scales, the experiments evaluate LoGRA’s accuracy, memory use, and sustained training, then probe compression choices, LoRA comparison, and KL budgets. LoGRA preserves competitive performance while reducing memory and enabling 27B training under the tested hardware budget.

3.1. General Experimental Setup

  • The experiments use PPO with verifiable rewards, one response per prompt, and a global running reward baseline, primarily on Qwen models from 1.5B to 27B.The deliberately simple setup is intended to isolate the method’s effects from potential gains of more sophisticated RL algorithms.Source: §3.1 · p. 5, §3 · p. 5

3.2. Main Experiments

  • At 27B, LoGRA reaches 71.52% Pass@1 and 81.87% Pass@4, while Dense exhausts GPU memory at its first Adam-state allocation.At 1.5B, LoGRA reaches 67.87% peak Pass@1 versus 63.77% for Dense; at 7B, peak Pass@1 is 72.33% versus 72.48%.Source: §3.2 · pp. 5–6, Figure 2 · p. 6, Table 1 · p. 7
  • Average memory falls from 9.18 to 7.18 GiB at 1.5B and from 31.82 to 17.29 GiB at 7B; at 27B, LoGRA averages 51.54 GiB.Figure 3 compares per-update peak allocated memory; at 27B Dense fails before an update, so no reduction ratio or throughput comparison is available.Source: §3.2 · p. 6, Figure 3 · p. 6, Table 1 · p. 7

3.3. Sustained 27B training on a single eight-GPU node

  • On one eight-GPU node, the 27B held-out macro score rises from 39.69% to 62.94%, peaking at 65.52% at step 1,060, without sustained score collapse.Mean response length falls from 8,312 to 4,705 tokens, and truncation falls from 34.72% to 12.95% over the trajectory.Source: §3.3 · p. 7, Figure 4 · p. 7

3.4.1. Design Choices in Low-Rank Gradient Compression

  • Increasing projection rank from 4 to 256 improves pass@1 by 3.02–3.95 percentage points, while refreshing bases shows no consistent benefit across ranks.At rank 256, fixed Rademacher, refreshed Rademacher, and refreshed Gaussian achieve 68.50%, 68.31%, and 67.42%; reconstruction cosine rises from approximately 0.046 at rank 4 to 0.348 at rank 256.Source: §3.4.1 · pp. 7–8, Figure 5 · p. 8

3.4.2. Comparison with LoRA

  • LoGRA uses 7.18 GiB average and 8.66 GiB peak memory versus LoRA’s 13.21 and 13.38 GiB, but LoRA has higher mean Pass@1 and a shorter training call.LoGRA’s mean Pass@4 is 0.40 percentage points higher; the comparison uses the recorded recipes and does not establish an accuracy or speed advantage.Source: §3.4.2 · pp. 8–9, Table 2 · p. 9

3.4.3. KL Budget Control

  • Across the first 800 logged steps, tail pass@1 is 70.30% without KL control, 63.43% with the large constant budget, 71.03% with the small constant budget, and 71.47% with annealing.The small constant budget improves tail accuracy over no control while maintaining stable learning; the no-control baseline also remains stable, so these two-seed results do not establish that KL control is necessary.Source: §3.4.3 · pp. 9–10, Figure 6 · p. 9

4. Related Works

The section situates LoGRA among approaches to RL’s coordination and memory demands. It distinguishes methods that manage resources or reduce trainable parameters from methods that compress gradients and updates.

4. Related Works

  • RL repeatedly coordinates generation, reward evaluation, and policy optimization, creating computation and memory demands beyond those of supervised fine-tuning.These stages require careful coordination of computation and memory.Source: §4 · p. 10
  • Asynchronous execution, resource scheduling, sharding, activation checkpointing, and CPU offloading address RL execution or memory costs.Asynchronous RL overlaps rollout with policy training, while other methods allocate resources or move, shard, or recompute training state.Source: §4 · p. 10
  • LoRA freezes pretrained weights and trains low-rank adapters, while GaLore and Flora compactly represent gradients and optimizer states with original weights trainable.These approaches reduce memory through parameter-efficient training or compressed training-state representations.Source: §4 · p. 10
  • LoGRA compresses gradients during backpropagation and reuses compressed updates for rollout-policy synchronization, extending compression to computation and communication.It targets the RL training–generation loop without materializing full gradients for selected weights.Source: §4 · p. 10

5. Conclusion

LoGRA uses low-rank gradient sketches throughout RL post-training and policy synchronization, with predicted-KL control regulating policy changes. The paper reports reduced memory without sacrificing reasoning performance, sustained 27B training, and supplementary statements on AI use, ethics, and reproducibility.

5. Conclusion

  • LoGRA uses low-rank gradient sketches for gradient accumulation, weight updates, and policy synchronization; predicted-KL control regulates policy changes.The conclusion reports that conservative KL budgets support stable long-horizon learning and that the method lowers the hardware barrier to RL post-training.Source: §5 · p. 10

AI Use Statement

  • The authors used generative AI for copyediting and clarity improvements; authors carried out the research questions, method design, experiments, and interpretation.Source: AI Use Statement · p. 11

Ethics Statement

  • The experiments used public models and benchmarks, with no human subjects or sensitive personal information; downstream use requires safeguards and evaluation.The statement says lowering hardware requirements may broaden access to model development.Source: Ethics Statement · p. 11

Reproducibility Statement

  • The main text specifies the algorithm and protocol, appendices document derivations and experimental details, and main comparisons are repeated across seeds with variability reported.Source: Reproducibility Statement · p. 11

A. Theoretical Derivations

The appendix derives a local predicted-KL estimate from changes in token scores and describes scaling a proposed update to a KL budget. It also specifies RowAdam’s row-wise adaptive scaling of gradient sketches.

A.1. From policy KL to changes in token scores

  • The local KL approximation is quadratic in update scale, with curvature determined by probability-weighted variance in token-score change rates.The derivation assumes logits have continuous, bounded third derivatives near the current policy; its probe uses sampled contexts and full-vocabulary score changes without dense gradients or curvature matrices.Source: §A.1 · pp. 13–14

A.2. Scaling an update to a KL budget

  • Predicted-KL step control scales the proposed update to meet a KL budget, capping enlargement when the estimate is small and leaving update direction unchanged.The multiplier is chosen from the budget and predicted KL; the estimate is based on evaluating the current and proposed policies on the same text contexts.Source: §A.1 · p. 13, §A.2 · p. 14

B. RowAdam

  • RowAdam adapts each sketch row using a running squared-magnitude estimate, unlike Adam’s per-parameter scaling and first-moment momentum.After row scaling it restores the sketch’s overall Frobenius norm; its row statistics persist across projection refreshes and require d second-moment values per selected matrix.Source: B. RowAdam · pp. 14–15

C.1. Mathematical reasoning

The mathematical-reasoning experiments construct a deduplicated 7,500-prompt training subset from DAPO-Math-17k and evaluate on MATH-500. The protocol specifies response generation, answer verification, and Pass@1 measurement.

C.1. Mathematical reasoning

  • DAPO-Math-7.5K takes the first 7,500 unique prompts in source order from the DAPO-Math-17k training split, deduplicating by first-message content.The source is the corresponding Hugging Face release (Yu et al., 2026).Source: §C.1 · p. 15
  • Evaluation uses all 500 MATH-500 test problems; training-subset deduplication is not cross-dataset semantic decontamination.Source: §C.1 · p. 15
  • Pass@1 is mean binary correctness across 2,000 generated evaluation responses, with unparseable final answers receiving no credit.Each problem gets four responses at temperature 0.7 and top-p 0.7; training generates one response per prompt at temperature 1.0 and top-p 1.0.Source: §C.1 · p. 15

C.2. Reasoning-Gym Hard mixture

The Reasoning-Gym Hard mixture selects eight tasks with intermediate verifier scores while excluding tasks whose failures mainly reflect truncation or formatting. It specifies training and evaluation problem counts, verifier scoring, and evaluation settings.

C.2. Reasoning-Gym Hard mixture

  • Eight tasks with preliminary mean verifier scores of approximately 0.25–0.60 were selected from 89, excluding tasks whose failures mainly reflect response truncation or formatting errors.The selection is intended to provide a challenging mixture with room for improvement.Source: §C.2 · p. 15
  • Each task contributes up to 5,000 candidate training problems and 25 default-difficulty evaluation problems, with duplicate training questions and train–evaluation overlaps removed.Table 3 reports the final counts for the mixture.Source: §C.2 · p. 15, Table 3 · p. 15
  • Task verifiers return scores in [0,1], and results use the equally weighted mean of eight task-average scores.Evaluation uses one response per problem, temperature 0.7, top-p 0.95, and a 16,384-token context limit.Source: §C.2 · p. 15

D.1. Main mathematical-reasoning experiments

The main experiments use one eight-GPU node split between FSDP training and generation, with shared precision and activation-checkpointing settings. Their KL budgeting includes a mismatch-based adjustment that the paper characterizes as a heuristic, not an exact KL decomposition.

D.1. Main mathematical-reasoning experiments

  • Each run uses one node with eight H100 80GB GPUs and 128 CPU cores, split into four FSDP training GPUs and four generation engines.Both use BF16 computation, activation checkpointing, no CPU offloading, and full-weight policy synchronization; schedule horizons count actual optimizer updates.Source: §D.1 · p. 16, Table 4 · p. 16
  • The main-experiment KL budget uses H=300, δmax=2 × 10^-4, and δmin=2 × 10^-5; 27B runs stop at 200 updates without shortening the schedule.The supplied passages do not state the intervening schedule formula.Source: §D.1 · p. 16
  • The effective KL budget subtracts an estimate of rollout-engine versus trainer log-probability mismatch, averaged over completed nontruncated responses and smoothed with weight 0.3.The estimate is clamped below at zero, and no subtraction is applied without a valid measurement.Source: §D.1 · p. 16
  • The mismatch adjustment is a budgeting heuristic, not an exact decomposition of KL.Source: §D.1 · p. 16

D.2. Long-horizon 27B training

The long-horizon 27B run uses a specified training configuration, while the reconstruction heatmaps average measurements across matrices, batches, states, and projection realizations. The LoRA comparison uses different counters, optimizers, and schedules, limiting what it isolates.

D.2. Long-horizon 27B training

  • The Figure 4 Reasoning-Gym run uses rank-256 refreshed projections, SGD sketch updates, and a KL budget annealed from 2 × 10^-4 to 2 × 10^-5 over 1,000 merges, then held constant.It uses seed 42, 128 prompts with one response each, microbatch size one, a 16,384-token context limit, and four training GPUs plus four generation engines on one node.Source: §D.2 · p. 16

D.3. Reconstruction heatmaps.

  • Reconstruction heatmap cells average seven attention/MLP matrices per layer, 32 batches, three model states, and eight random projection realizations.The measurements use initialization and steps 160 and 280 of one fixed-basis rank-64 trajectory; fixed bases are reused across batches, while refreshed bases are redrawn.Source: §D.3 · pp. 16–17

D.4. LoRA

  • The rank-256 Qwen2.5-Math-1.5B LoRA comparison uses three runs per method, but different counters, optimizers, and schedules prevent isolating adapter parameterization alone.LoRA statistics cover 300 logged steps, excluding records without an update from memory/time averages, whereas LoGRA covers 300 actual updates.Source: §D.4 · p. 17

Related papers

Loading 2610.06647v1…