Running async RL on my DGX Spark cluster

After tuning the rollout server, I wanted to run the full async RL loop on my own hardware. The learner needed to keep training while the sampler produced new trajectories, with fresh weights moving between them and enough information in each trajectory to decide whether it was still valid to train on.

I used Qwen3-4B so one DGX Spark could train and the other could serve rollouts through vLLM. NeMo RL owns the learner, and Polar runs the agent harness and collects trajectories. This is a different setup from the previous recipe, where both machines were needed to serve DeepSeek-V4-Flash.

The loop worked, but a training step took about 270 seconds, and about 84% of that was weight sync. I expected to optimize the transfer, but the time was going somewhere else: the trainer spent about 227 seconds moving optimizer state out of the way of an inference engine that was already on the other machine. Skipping that offload cut mean sync time from 226.5 seconds to 1.73 and mean step time to about 42 seconds.

• Profiling the wallclock
  • Testing the transfer
  • Keeping the optimizer on the trainer
• The trainer-sampler contract
• The rollout server collects, the trainer decides
• Defining your config
• The recipe

Training step

Weight sync Rest of step

Baseline≈270 s total

Weight sync 226.5 s

Skip offload≈42 s total

Weight sync 1.73 s

0100200300 s

Mean durations across three seeds. Both bars use the same 0–300 s scale.

Figure 1. Skipping the optimizer offload reduced mean step time about 6.4x. The three fixed runs took 35.0, 55.8, and 35.0 seconds; weight sync stayed near 1.73 seconds.

Profiling the wallclock

Generation and policy-training compute looked healthy, so I started with the weight transfer. Even after moving the broadcast to RoCE, the fp8 rollout configuration spent a mean of 226.5 seconds in weight sync across three seeds.

I had two places to look first: transport, which determines how the bytes move, and payload, which determines how much I send. This weight-sync guide explains that split well. My loop broadcast the full bf16 weights over NCCL, so I tried changing both.

Testing the transfer

RoCE had helped in the last recipe. Here, moving the broadcast off TCP cut weight-sync time by only 5 to 7%. I also tried enabling GPUDirect RDMA, but nvidia_peermem failed to load on both hosts with the installed driver and kernel, so I could not measure that path.

Then I tested the payload. FP8 rollout precision should have cut the transfer, but it did not, because the broadcast casts to the training dtype, bf16, and vLLM requantizes to fp8 after receipt. I made the broadcast itself fp8, matching the engine’s blockwise cast in all 41 tested shape and edge cases. That dropped the payload 45%, from 8.04 to 4.41 GB. Weight sync did not move: 231 seconds against 226, inside the noise.

If transfer time were the problem, a 45% smaller payload should have made a visible difference. At the link rate, 8 GB would take about 0.3 seconds before protocol and software overhead, against more than 200 seconds of reported weight sync. So I stopped tuning the transfer and profiled what the refit, the step that loads new weights into the inference engine, was doing around it.

The profile put almost all of the time outside the broadcast. Producer packing and broadcast took about 1.7 seconds, while offload_before_refit took about 227. That phase moved roughly 32 GB of Adam optimizer state from GPU to CPU, one tensor at a time, before every sync.

Refit profile

offload_before_refit227.3 s

broadcast1.7 s

load1.5 s

050100150200250 s

Three measured phases from one refit. All bars use the same 0–250 s scale.

Figure 2. The optimizer offload accounted for 227.3 of the 231 seconds spent in this refit. Phase names come from the profiler.

That offload is a guardrail for colocated serving, where the inference engine shares the trainer’s GPU. NeMo RL uses it to clear trainer-card memory before refit, but this run serves rollouts on the other Spark. On GB10’s unified memory, the guardrail becomes a 227-second copy inside the same physical memory. NeMo RL’s own code comment warns against running it in non-colocated mode.

Keeping the optimizer on the trainer

I added a guard to skip the optimizer offload when the inference engine is remote. The optimizer stays on the training GPU, so there is also nothing to move back before the next update. The loss and update rule stay the same.

The three seeds came back at 1.73, 1.73, and 1.74 seconds for weight sync. Total step time varied more, at 35.0, 55.8, and 35.0 seconds, averaging about 42. Sync was now a small part of the step, and policy training the largest cost.

All three runs passed the training-data audit, a script that checks each run exited cleanly, that the saved training batch carries token ids, loss masks, and both sampler and trainer logprobs, and that the log records Polar rollouts, replay-buffer adds and samples, a policy update, and a weight sync. Their mean importance ratios, the trainer’s probability divided by the sampler’s for each trained token and averaged over the batch, were 1.0018, 1.0006, and 0.9966. That checked the data path after the change; I had not run a learning-curve comparison or a bitwise comparison of optimizer trajectories.

Figure 3. Polar’s rollout architecture (Apache 2.0). NeMo RL trains on one Spark; vLLM and Polar run on the other. The collector turns Polar’s trajectories into replay-buffer groups for the learner.

The trainer-sampler contract

The other part of standing up the loop was keeping the trainer and sampler in agreement. I found Hugging Face’s survey of 16 open-source RL libraries useful for comparing how frameworks handle those boundaries.

1. The rollout server generates trajectories from a specific policy version. It needs high throughput, but it also has to preserve the exact training data the learner will consume, down to the token ids and logprobs.

2. The learner updates the policy and decides which samples are still valid. Rollout age, replay, off-policy correction, and rewards determine how those samples enter the update.

3. The weight-sync path moves fresh weights from the learner into the inference engine. Its timing affects how old the sampled policy is when a trajectory reaches the learner, as well as how long either side has to wait.

4. The trajectory contract preserves tokens, logprobs, rewards, masks, and policy-version metadata across the boundary. I check those fields at the boundary, because a trajectory can look correct as text while carrying the wrong training data.

The rollout server collects, the trainer decides

Polar is a trainer-agnostic rollout layer. It separates agent harnesses (OpenHands, mini-swe-agent, Terminus2) and environment backends (Docker, Modal, EC2) from trainer logic, so each can change independently. Evaluation, distillation, and RL all sit behind one rollout endpoint where a client submits a task and receives a trajectory. NeMo RL fills that client role here, but nothing on the rollout side is tied to NeMo.

The collector’s own docstring draws the boundary.

This module intentionally keeps the integration surface small: NeMo still owns
the trainer, replay buffer, staleness window, importance-sampling correction,
and vLLM weight sync. The collector only replaces NeMo's default rollout actor
with one that submits Polar tasks and converts the resulting traces into the
same per-prompt replay-buffer group shape NeMo already samples.

Because the collector is a single actor swap, the rest of the trainer’s async machinery runs unmodified. Before a group of rollouts reaches the replay buffer, the collector checks that every session produced a usable trace and raises if one did not.

if valid_sessions < group_sessions:
    raise RuntimeError(
        "Polar collector refusing invalid group before replay-buffer add: "
        f"valid_sessions={valid_sessions} group_sessions={group_sessions}"
    )

Failing the step on an incomplete group keeps the trainer from silently receiving fewer samples than it expected. A second check warns when reward variance falls below 1e-5: if all rewards are equal, the group-relative advantages are zero and that group supplies no policy-gradient signal.

The trajectory the collector does pass has to survive tokenization without drift, because the trainer has to see the exact token ids the sampler saw. Token ids come from the inference backend, not from re-tokenizing text: the gateway asks vLLM for token ids and per-token logprobs directly, and the assistant tokens in the merged trace are the raw ids the model sampled.

Byte-pair encoding (BPE) is not round-trip stable. [fish, ing] can re-tokenize to [fishing] in the next turn’s prompt, shifting every token after it. Prime Intellect’s renderers library documents the same class of bug and puts a number on it: letting the inference engine re-render full history each turn broke token identity on 32 of 64 rollouts with Qwen3.5-35B-A3B, where a bridge that never re-renders broke none.

To decide which existing chain of turns a new completion extends, Polar compares canonical prompt tokens, and it keeps the raw sampled assistant ids for the trajectory itself. Re-encoded response text never enters that decision.

def _find_extendable_chain(prompt_ids, chain_tips):
    """A completion continues a chain iff its prompt begins with that
    chain's last prompt. The compared prefix is two server-side
    tokenizations of the same text, so BPE re-tokenization of the
    sampled response never enters the decision.
    """
    best_idx, best_len = None, -1
    for idx, tip in enumerate(chain_tips):
        n = len(tip)
        if n > best_len and 0 < n <= len(prompt_ids) and prompt_ids[:n] == tip:
            best_idx, best_len = idx, n
    return best_idx

The merged stream takes assistant bodies from the raw sampled ids, which keep their real logprobs and a loss mask of 1. It takes interstitials (tool results, chat-template glue) from the canonical next-turn tokens, which get a zero loss mask. Nothing the model generated is decoded and re-encoded.

The serving side has to pause, drain, refit, and resume. When the trainer is ready to push new weights, the collector pauses generation through Polar’s own proxy, an asyncio.Condition gate that stops accepting new requests and drains the in-flight ones before the weight update, then resumes. vLLM shipped a native RL pause-and-weight-transfer API in May, but I keep the gate on Polar’s side because the trainer talks to one stable Polar endpoint whether the engine underneath is vLLM or SGLang.

Defining your config

The config defines the training algorithm. The policy-gradient family shares one ClippedPGLossFn, and GRPO, Dr. GRPO, RLOO, CISPO, GSPO, and DAPO are each a set of overrides on that one loss. GRPO adds no overrides, so it is the base config. In this pinned NeMo implementation, CISPO selects a different clipping path: it sets use_cispo=true, token-level loss, per-token importance ratios, and a clip range of (1, 4).

I ran GRPO and CISPO at matched shape on the same prompts to check that switching objectives did not require another integration path. The smoke runs produced the same observed reward distribution, with weight-sync times of 251 and 237 seconds on TCP, within the timing variation I was already seeing. Comparing how the objectives learn would require longer runs.

The recipe

The runnable experiment specs and smoke scripts live in the repo, jbarnes850/ProRL-Agent-Server.

Setting	Value	Why it matters
NRL_REFIT_SKIP_OPT_OFFLOAD	1	skips the optimizer CPU offload when inference is on another machine, the 131x weight-sync fix
NCCL_TRANSPORT	roce	routes the weight-sync broadcast over RDMA. Confirm Using network IB in the log
NRL_REFIT_BUFFER_MEMORY_RATIO	0.01	sizes the weight-broadcast staging buffer with headroom on unified memory
gpu_memory_utilization	0.5	0.35 and 0.5 ran in this setup; 0.65 and 0.9 exhausted memory
max_num_seqs	256	server admission cap; this experiment uses 32 requests
max_num_batched_tokens	16384	batch headroom for the vLLM engine
precision.train	bfloat16	the default full-weight broadcast uses this dtype
precision.rollout	fp8	quantizes the rollout engine’s weights after the bf16 broadcast
kv_cache_dtype	auto	keeps the cache at the backend’s default dtype
async_grpo.lag	1	max trajectory age. NeMo turns on importance-sampling correction for async mode
rollout shape	8 prompts x 4 generations	one group per prompt
IMAGE / NEMO_RL_REF	c236061b	the pinned NeMo RL image with the CISPO config key

I validate the spec through NeMo’s config loader, write the launch script with the CLI’s launch command and the profile that skips the offload, then run that script on the rollout Spark, which starts the training machine’s container over SSH.

# no-GPU gate: apply the compiled overrides through NeMo RL's own config loader
python scripts/experiment/validate_against_nemo_loader.py \
  examples/experiments/qwen3-4b-instruct-basic_arith-grpo-fp8-rollout-8x4.yaml

# emit the two-machine launch script with the offload fix enabled
python -m nemo_polar_bridge.experiment.cli launch \
  examples/experiments/qwen3-4b-instruct-basic_arith-grpo-fp8-rollout-8x4.yaml \
  --profile examples/experiments/profiles/two_spark_roce_buf001_skipoffload_prod.yaml \
  --slug async-rl-recipe --store runs/lineage --out runs/async-rl-recipe

# execute on the rollout Spark
bash runs/async-rl-recipe/launch.sh

Agentic RL is sample-hungry, so inference throughput matters, but the sampler is only one part of the loop. A 2-machine cluster was small enough to profile each phase, which is what made the unnecessary 227-second offload visible. With that out of the way, I could spend the next round of optimization on the work the learner was actually doing.