Async GRPO with LoRA across HF Jobs: a bucket, a proxy, and no NCCL

Update on GitHub

Upvote 22

TL;DR

AsyncGRPOTrainer

can now train a LoRA adapter and sync only that adapter to vLLM (TRL v1.14).

A rank-1 adapter is a few megabytes, so it can travel through a Storage Bucket mounted in every Job instead of over NCCL. The trainer and the vLLM replicas run as separate Hugging Face Jobs on separate machines.

A small proxy in front of the replicas adds the auth header, routes each rollout to the replica that already holds its KV prefix, and broadcasts adapter loads to every replica.

The AsyncGRPO metrics show where the bottleneck sits. Five runs take the same recipe from 3 h 27 min to 53 min for 500 steps.

LoRA support recently landed in TRL's

AsyncGRPOTrainer

LoRA training is particularly suited for RL, as shown in Thinking Machines's blog LoRA Without Regret. They show that LoRA can match full fine-tuning for policy-gradient RL, even with rank 1. This stems from the fact that the advantage function only gives

~O(1)

There is also a systems consequence of LoRA training. A rank-1 adapter for a 1.5B model is a few megabytes, while the full model is around 3 GB. Instead of sending the full policy to the inference workers after every update, we can just send the adapter. vLLM can also keep several adapters loaded at once. Old rollouts finish with the policy they started with, while new rollouts use the latest one.

TRL's

AsyncGRPOTrainer

What we want is to run the same setup with Hugging Face Jobs. Essentially, an HF Job is one container running on one VM. This means that one Job cannot spawn multiple nodes (at least for now) to hold a trainer and a fleet of vLLM servers (we are limited to 8xH200 at most per node). The

AsyncGRPOTrainer

Well, with a full-weight sync, the answer would be "not far". Every update would have to move gigabytes between machines, which is what NCCL is for in a dense cluster, but Jobs can't communicate across nodes. There is no shared local disk and obviously no shared

localhost

The setup ended up being quite small:

• a trainer Job running

  AsyncGRPOTrainer

  with LoRA (and FSDP, more on that later),
• two vLLM Jobs, each serving the base model plus whatever adapter the trainer last published,
• a Storage Bucket mounted in all three at the same path, which is how the adapter gets from the trainer to the servers,
• a proxy server. We'll dive deeper into why we need one, but at a high level we need a proxy that routes each rollout to the replica most likely to hold its KV cache, and broadcasts every adapter update to all vLLM replicas.

The architecture: leveraging Hugging Face Jobs and Storage Buckets 🪣

The new adapter-only sync path in

AsyncGRPOTrainer

<output_dir>/.vllm_lora/trl-policy-v{N}

/v1/load_lora_adapter

model="trl-policy-v{N}"

This is how runtime adapter loading already works in vLLM. The endpoint takes a path, not tensors, so the trainer and the server are expected to share a filesystem. On a Slurm cluster, that is the network filesystem. On Jobs, we get the same thing by mounting a Storage Bucket as a volume at the same path in every Job, as we mentioned earlier. Under the hood, it uses

hf-mount

# every Job gets the same bucket at the same absolute path
hf jobs run ... -v hf://buckets/aminediroHF/asyncgrpo-lora-buckets:/lora ...

Nothing in TRL or vLLM had to change for this. The trainer writes to

/lora/<run>/.vllm_lora/

Note that we also store the checkpoints and the final adapter in the bucket. The HF Jobs are ephemeral, but a preempted trainer can resume training, as the final adapter is always persisted to the bucket and is never lost when the Job stops.

The three Jobs

The vLLM replicas

Each replica uses one GPU and the stock

vllm/vllm-openai

The number of adapter slots follows from

max_staleness

AsyncGRPOTrainer

max_staleness

max_staleness=4

trl-policy-v3

v7

v3

v3

max_staleness + 1

--max-loras 6

# --expose 8000                       reachable at https://<job_id>--8000.hf.jobs
# -v ...:/lora:ro                     read-only: the server only reads adapters
# VLLM_ALLOW_RUNTIME_LORA_UPDATING=1  enables /v1/load_lora_adapter
# VLLM_SERVER_DEV_MODE=1              enables /pause, /resume, /server_info (TRL needs all three)
# --max-loras 6                       max_staleness=4 -> 4+2 adapter slots
for replica in 1 2; do
hf jobs run --detach --flavor h200 --timeout 8h --secrets HF_TOKEN \
    --expose 8000 \
    -v "hf://buckets/${BUCKET}:/lora:ro" \
    -e VLLM_ALLOW_RUNTIME_LORA_UPDATING=1 \
    -e VLLM_SERVER_DEV_MODE=1 \
    -- vllm/vllm-openai:v0.27.1 \
    vllm serve Qwen/Qwen2.5-Math-1.5B --host 0.0.0.0 --port 8000 \
        --max-model-len 4096 --logprobs-mode processed_logprobs --generation-config vllm \
        --enable-lora --max-lora-rank 1 --max-loras 6
done

We pin vLLM to

v0.27.1

There is another possible design where the trainer keeps only the latest adapter and always publishes it under the same name. We did not go that way, because vLLM keys its prefix cache by adapter name. With a single name, KV blocks computed under the previous weights would still match after the swap, so the prefill would not be redone and a rollout could get its prefix from one policy version and its decode from the next. The trainer would have no way to tell, and it would show up as

ratio

The dataset choice: the Sanity set

We chose

sail/Sanity-Test-R1D-1.5B

sail-sg/Precision-RL

The authors generated 40 answers for each MATH problem with DeepSeek-R1-Distill-Qwen-1.5B. They kept problems with a success rate between 20% and 80%, yielding 1,460 questions. This dataset is really good for RL validation because the questions are neither already solved nor completely hopeless for that model, meaning the model can get a good early signal to train on and improve.

This is awesome as a robust end-to-end test: if one vLLM replica silently serves the base model under an adapter name, we want to see that in the curve within a few dozen steps. Also, this dataset is small enough to cycle through in less than two hours.

We also take the hyperparameters from the paper's LoRA scripts in

oat/scripts/lora

Qwen/Qwen2.5-Math-1.5B

The trainer

The trainer uses the same

vllm/vllm-openai:v0.27.1

AsyncGRPOTrainer

from peft import LoraConfig
from trl.experimental.async_grpo import AsyncGRPOConfig, AsyncGRPOTrainer

config = AsyncGRPOConfig(
    output_dir="/lora/sanity-lora-r1",       # on the bucket: adapters, checkpoints and the final adapter all land here
    vllm_server_base_url="http://localhost:8000",   # the proxy, not a vLLM Job; TRL never sees the Jobs URLs
    max_staleness=4,
    weight_sync_steps=4,                     # publish an adapter every 4 optimizer steps
    save_strategy="steps", save_steps=50,    # checkpoints go to the same bucket -> resume after preemption
    ...
)
trainer = AsyncGRPOTrainer(
    model="Qwen/Qwen2.5-Math-1.5B",
    args=config,
    peft_config=LoraConfig(r=1, lora_alpha=2, target_modules="all-linear"),  # plain LoRA vLLM can serve as-is
    ...
)

During initialization, TRL calls

/server_info

. If it finds a

lora_config

, it uses adapter-only sync. Configurations vLLM cannot serve directly, such as DoRA,

modules_to_save

, or a rank above

--max-lora-rank

, fall back to merged-weight sync with a warning. The log should contain

Adapter-only vLLM sync enabled

.

The proxy

Now onto the fun stuff. We need a proxy between the trainer and the vLLM Jobs for two reasons:

• Exposed Job ports require an

  Authorization: Bearer <HF token>

  header on every request. The proxy is where that header gets added, so TRL does not need to know about it.

• We want more than one GPU generating. On a single vLLM server, the usual way to get that is

  --data-parallel-size > 1

  , but TRL refuses adapter-only sync in that mode, for a good reason: a call to

  /v1/load_lora_adapter

  only reaches the DP rank that answers it, so the other ranks would keep serving the base model under the new policy name. On Jobs the question does not even arise, since each replica is its own machine. So the data parallelism has to live one level up, in something that fans the adapter load out to every replica.

We therefore run a small proxy at

127.0.0.1:8000

• It sends each completion request to one replica, chosen so that the eight rollouts of a prompt land where their prefix is already cached (details on this below).
• It broadcasts every state-changing request, such as adapter loads, pause and resume, to all replicas, so that a policy name means the same thing everywhere.

Routing rollouts by KV prefix

A quick reminder of why this matters. Generating a completion has two phases with very different workload profiles:</