Running RL on Academic budget at high-end speed

2 September 2026

Here I share how I took verl from “cannot start” to a significant step-time reduction, and only ~two times the 2XB200 speed for Qwen3-4B GRPO on two A6000s with no NVLink

Introduction

I have a real tragedy. I am doing some post-training RL research now, but I don’t have much compute. I have some ephemeral access here and there to high-end GPUs (B200, H200, H100), but not much. On the other hand, my uni’s cluster has a lot of L40s/A600 gpus, which are not that occupied because everyone tries to run on the higher end cards. As my only resort, I fought to make these gpus usable for RL research, and I survived to tell you and share the outputs so anyone reading it and running on an academic budget would be able to make use of this older hardware for their next RL project.

A real goal throughout the process below is to do all the optimization we can, but nothing that touches the training dynamics.

Setup

Our working setup is 2XA6000, and we will focus on training Qwen3-4B on the DAPO-17K dataset with batch size 32 and group size 8. PPO mini batch size will be 256 (we will do a fully-on-policy run; this is an important detail for later). We don’t have NVLink on our setup, so devices are connected only using PCIe and can communicate through host memory only.

Step -1

As mentioned, we have no NVLink in our setup. Therefore, the 2 gpus can only communicate via the host memory and their PCIe connectors. This makes the fsdp communication very slow and it slows down training to unbearable step time. So we resort to DDP using the supported FSDP engine in verl (i.e. fsdp_size = 1). Note that we could have used other strategies such as HSDP to deal with the non-existence of the NVLink and the communication overhead, but that is left for another post.

Therefore, the configuration is ordinary with an important distinction, a colocated vLLM rollout engine. The unusual part is fsdp_size=1: we don’t shard the model (to avoid expensive communications). This setting makes fsdp replicate the model on each card, which costs more HBM memory. Therefore we lean on param_offload plus optimizer_offload to keep the memory bill payable.

param_offload offloads the fsdp model weights (the “trained” weights) during inference to the host memory. The model is pretty small and loading it back to GPU memory is fast and definitely worth the extra ±8GB of HBM memory for decode kv cache. optimizer_offload does the same for the optimizer tensors which are double the size of the model weights (since fsdp updates fp32 weights, they are called “master weights”).

These cards are very short on memory and we want to utilize every bit of it. As a typical RL loop goes: Decode sequences of tokens -> update the model. The update and the inference are totally separable, so we use a feature of our inference engine, vLLM in this case, called “sleep mode”.

Sleep mode allows us to tell vllm to release the KV cache HBM memory allocation it holds and offload the model weights to CPU RAM (which is much larger most of the time). Since at every inference iteration the weights change (we do on-policy RL!), we don’t even need that and we can tell vllm to drop that entirely. So we want vllm to be able to use the entire HBM memory for allocating kv cache during inference (so all sequences could be decoded in parallel). Note that we do not need anything of the training phase on the GPU when at the inference phase, so we can use the entire card memory for that and release it entirely when it’s done and we move to the training phase.

Step 0 - A good old OOM for the beginning

vLLM sizes its KV cache at startup by reading how much device memory is free. Apparently, it read the wrong number and quit:

ValueError: Free memory on device cuda:0 (31.33/47.4 GiB) on startup is less
than desired GPU memory utilization (0.85, 40.29 GiB).

The message points at gpu_memory_utilization, which we set very high (0.85) in order to utilize it fully. After enabling the DEBUG logging level:

After FSDP,                                     allocated 15.02 GB
After offload model/optimizer/grad during init, allocated 15.02 GB

Look! param_offload ran and freed nothing.

State: step time — · device 15.02 / 47.40 GB

Step 1 — Offloading model in fsdp_size=1 does not free the memory

First, a few terms:

ℹ️ use_orig_params is a flag with which verl tells fsdp to use the original model tensors (e.g. model.weights.q) as the registered parameters (i.e. the parameters that are returned by model.parameters() - the ones which the optimizer updates). When False, fsdp uses the flat param (next paragraph) tensor as the data storage, and the named parameters (e.g. model.weights.q) are just views (pointers) into that storage.

ℹ️ flat_param is how fsdp stores the model params. It is named after the fact that it stores them in a 1D tensor. It groups multiple params into a single flat tensor - by default an entire layer’s params. Then, the model parameters (e.g. model.weights.q) become views into that storage - i.e. a pointer and a shape (the tensor’s metadata) to the flat param offset for which the actual parameters are stored.

fsdp_size=1 builds a two-dimensional (ddp, fsdp) device mesh (that is just a matrix such that each row is a set of devices responsible for one shard of the data, and each column is a set responsible for one shard of the model parameters under fsdp) whose shard group (fsdp devices) has exactly one rank. FSDP1 quietly degrades that to NO_SHARD - an fsdp way to say that this model is not sharded. Good! That’s what we wanted. Combined with verl’s default use_orig_params=False (more on that later), this happens:

FlatParamHandle.flat_param_to() moves the flat-parameter storage to the host, but it only refreshes the per-parameter views when use_orig_params=True. With False, the views registered on each module still point at the old device allocation. The storage is referenced, so nothing is freed!

A 735M-parameter reproduction test, isolating the variable:

Row 3 deserves attention. After a forward pass the bug disappears. By then FSDP has re-pointed the views at the bf16 compute shard, which it frees post-forward, so no fp32 storage is pinned (i.e. nothing is pointing to the old HBM resident weights) and the offload appears to work. The leak is only visible before the first forward — which is precisely when verl offloads, right after model construction, and is — most importantly — precisely the allocation vLLM then reads.

It compounds, too: because the stale views keep the original allocation alive, the subsequent reload allocates a second copy. The round trip left 5.4773 GB resident for a 2.7387 GB model!

The fix re-points the views after the move, gated so that neither use_orig_params=True nor any genuinely sharded strategy changes behaviour:

def _refresh_unsharded_views_after_move(handle):
    if handle._use_orig_params:
        return   # flat_param_to() already refreshed the views
    if handle.uses_sharded_strategy:
        return   # sharded strategies rebuild views on the next unshard
    handle._use_unsharded_views(as_params=False)

End to end, on the real 4B model:

After offload model/optimizer/grad during init, allocated 0.00 GB

vLLM now starts at gpu_memory_utilization=0.85.

State: step time — · device after offload 0.79 / 47.40 GB

This fix is now live and awaiting merge on the verl project as PR #7596 (Edit: It was closed due to verl’s unwillingness to support fsdp1 with size=1 :( )

Step 2 — Moving Optimizer step to CPU

First step. Generation ran. Good! Log-probs ran. Nice! Then the first optimizer step died:

engine/fsdp/transformer_impl.py:803  optimizer_step
  torch/optim/adam.py:178 _init_group
    state["exp_avg"] = torch.zeros_like(p, memory_format=torch.preserve_format)
torch.OutOfMemoryError: Tried to allocate 386.00 MiB. GPU 0 has 47.40 GiB
total capacity; 112.94 MiB free; 46.83 GiB in use.

Our NO_SHARD strategy means every rank holds the whole model, and Adam’s first step needs four full-size fp32 tensors co-resident:

Note that optimizer_offload does not help here. It parks Adam state on the host between phases (between inference and training), and loads it back to the device to step. Offloading the optimizer shuttles the Adam states to GPU every step to run optimizer.step() there, then move them back to CPU afterward to free GPU memory before the next rollout phase. So natively the states do round-trip GPU↔CPU each step. That costs host-device communication and is restricted to cases where GPU memory allows storing all the required tensors on device. That is not the case for our NO_SHARD strategy.

We will use our much more abundant host RAM by moving the step to the host. Parameters, gradients and Adam state all on CPU at optimizer.step(), so the 64 GB never lands on the card. Staging everything at once needed ~60 GB of host memory per rank and would not schedule anywhere on our cluster, so the step walks one flat parameter at a time (~34 GB/rank). Note: we checked — over five iterations, parameters, exp_avg, exp_avg_sq and the step counter all differ from a normal AdamW.step() by 0.0.

It now trains :).

State: step time 445–450 s · peak 41.28 / 47.40 GB First successful run · 4/4 steps · grad_norm 0.049–0.064 · ~625 tok/s

Step 3 — Dropping the unnecessary fully-on-policy old logprobs calculation

With a working run we could finally read a phase breakdown. On an A6000 baseline, the main step time components are: generation 222.8 s, update_actor 178.3 s, and old_log_prob 37.0 s — 8% of the step spent on a separate forward pass over the batch.

In a strictly on-policy configuration — one mini-batch — that pass computes exactly what the update’s own forward computes moments later. Same weights, same tokens. We measured the claim rather than assuming it:

max|log_prob.detach() − old_log_probs| = 0.000e+00

Bit-exact, not approximately equal. So the separate pass can be skipped and both of its consumers — the PPO ratio and the truncated importance-sampling weights — can be served from the detached output of the update’s forward. The IS weights match the ones computed out of old_log_probs to 1.192e-07 (ordinary fp32 rounding).

This is guarded hard since it might change training dynamics, which we want to avoid. It fires only when ppo_mini_batch_size == train_batch_size, ppo_epochs == 1, and roughly a dozen other preconditions hold; anything else prints its reasons and falls back. A skip that is wrong here would silently corrupt training rather than fail.

Note: It is not free under verl’s engine: two metric groups (actor/entropy and the rollout_corr/* family) were derived from the skipped pass. Entropy comes back via calculate_entropy=True.

State: old_log_prob 37.0 → 0.084 s We reduced 37.0 s of the step time. Good progress!

Step 4 — Pinned memory for the master weights/gradient tensors

Moving a 16 GB fp32 master between host and device every step is a lot of PCIe traffic, and the CUDA copy engine can only DMA directly out of page-locked memory. From pageable memory the runtime has to stage through a hidden pinned buffer first — an extra hop, repeated every transfer.

So we allocate the pinned buffers once and reuse them. One growable shared buffer per tag, sized to the largest flat parameter, rather than one buffer per parameter — the per-parameter version pins a second full copy of the model, roughly 30 GB of host memory, for no additional speed.

This turned out to be a big win.

State: update_actor 87.31 → 65.35 s · win −21.96 s

Step 5 — Reducing host memory demand

By default, when using fsdp, even though NO_SHARD is configured, in this replicated setup, both ranks all-reduce their gradients and then perform the identical Adam step on identical state. One of the two is pure waste — 30 GB of host memory holding a second copy of exp_avg and exp_avg_sq that will never differ from the first.

So we change things so that rank 0 owns the optimizer and broadcasts the updated fp32 master. Rank 1 never allocates Adam state at all. The broadcast must carry the fp32 master rather than only the bf16 compute copy, because rank 1’s master is a live source for the vLLM weight sync and the log-prob path.

max|rank1 weights − rank0| = 0.000e+00
adam_state  rank0 = 29.97 GiB / 37 params    rank1 = 0.00 GiB / 0 params

The time saving is modest. The main contribution of this step is the 30 GB of host memory: in our setup it allowed us to schedule jobs on the cluster more easily.

State: step time 375.5 s · peak 40.33 / 47.40 GB Steps 3–5 together: 445–450 s → 375.5 s, a 16% reduction. Independently reproduced on H200 at −15.6%.

Step 6 — Avoiding master weights on GPU

One more thing we tried to do: release the resident fp32 master weights from GPU memory in the hope that it will allow us to increase the micro-batch size.

During forward and backward the fp32 master is idle — compute runs on the bf16 copy. Park it on the host for that window and you get 16 GB of device memory back. We dropped the master weights to the host and measured a −8.75 GB reduction in peak VRAM — and a +18.7 s increase in step time. Why would this increase step time? The cost comes from a single line of FSDP. Every forward, for every handle, it rebuilds the bf16 shard for every micro-batch!

flat_param._mp_shard.copy_( flat_param._local_shard.to(self.device) )

We measured the number of such casts (and transfers): We found it was ~800–1300 times per step, about 350 GiB of fp32 read. That is reducible since we can just cast the entire model to bf16 once and transfer it to gpu once for the gradient computation forward-backward.

Note that this, again, is tied to our no-shard strategy since in the sharding case, FSDP really doesn’t keep the BF16 compute weights to save memory.

So cast once per step and keep the bf16 shard resident across all micro-batches, invalidating it after the optimizer step. The instrumentation shows the mechanism doing exactly what it claims:

Casting once recovers 15.4 s of the 18.7 s penalty — 82% — for 1.8 GB of the memory saving. The package nets out at −6.95 GB of peak for +3.3 s, about 2%.

The lesson hides in the comparison of B and A. Cast-once is worthless on its own: with the master already device-resident those 863 casts are device-to-device and cost almost nothing. It only becomes valuable once you have moved the master to the host — an optimization whose entire worth is created by a different decision — saving memory.

Correctness is checked at steps 1 and 2 — the second zero is the one that matters, since it proves the resident shard was re-cast after the optimizer updated the master rather than training on stale weights:

max|resident_bf16 − cast(current_master)| = 0.000e+00   (step 1 and step 2)

State: step time 369.4 s · peak 33.38 / 47.40 GB

The whole arc, one card type

All four rows are two A6000s, gpu_memory_utilization=0.85, max_token_len=24576, mean of steps 2–4.

Every arm passed the same sanity bar — grad_norm between 0.036 and 0.066, finite pg_loss, no NaN — and the two changes that alter what is computed were held to bit-exactness rather than tolerance.

Reproduction

The code will be open-sourced soon (I need to clean the instrumentation and all that stuff), so we can all run more RL experiments.