Skip to content

Delta weight sync for AsyncGRPO (sparse patches over HF Bucket) #5937

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
AmineDiro wants to merge 6 commits into
base: main
Choose a base branch
from
Open

Delta weight sync for AsyncGRPO (sparse patches over HF Bucket) #5937

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
0f13f37
Add sparse delta codec and AdamW-inversion change detector
AmineDiro Aug 14, 2026
72cf851
Add NCCL and HF-bucket weight transfer backends
AmineDiro Aug 14, 2026
d2e5671
Wire delta weight sync into the AsyncGRPO trainer
AmineDiro Aug 14, 2026
798d4fe
Add delta weight sync examples and docs
AmineDiro Aug 14, 2026
6a77b46
Import the bucket engine lazily so the NCCL path needs no Hub bucket …
AmineDiro Aug 14, 2026
8bbc35f
Address review nits on docs and example headers
AmineDiro Aug 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
62 changes: 59 additions & 3 deletions docs/source/async_grpo_trainer.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,12 +1,12 @@
# Asynchronous GRPO

> [!IMPORTANT]
> This trainer requires `vllm>=0.22.0` and `transformers>=5.2.0`. For distributed training, only FSDP2 is supported (DeepSpeed ZeRO is not).
> This trainer requires `vllm>=0.23.0` and `transformers>=5.2.0`. For distributed training, only FSDP2 is supported (DeepSpeed ZeRO is not).
>
> Currently, `vllm` and `transformers` have conflicting dependency constraints. To work around this, install vLLM first and then force-install transformers:
>
> ```bash
> pip install 'vllm>=0.22.0'
> pip install 'vllm>=0.23.0'
> pip install 'transformers>=5.2.0' --no-deps
> ```

Expand Down Expand Up @@ -34,7 +34,7 @@ The rollout worker runs in a separate process spawned from the trainer, so rewar
>
> If you do need a GPU reward model, the recommended approach is to **serve it behind its own inference engine** (vLLM, TGI, …) on separate GPUs and have a lightweight, picklable reward function call it over HTTP. This keeps the reward model on its own device while the rollout process stays CPU-only, and it scales independently of the trainer.

After every `weight_sync_steps` training steps, the updated weights are transferred to the vLLM server via NCCL so that subsequent generations reflect the latest policy.
After every `weight_sync_steps` training steps, the updated weights are transferred to the vLLM server so that subsequent generations reflect the latest policy. How they are transferred is configurable, see [Weight synchronization](#weight-synchronization) below.

Because generation and training run concurrently, the training samples may have been generated by a slightly older version of the model. The `max_staleness` parameter controls how many weight updates a sample can lag behind before being discarded.

Expand Down Expand Up @@ -76,6 +76,62 @@ CUDA_VISIBLE_DEVICES=0 VLLM_SERVER_DEV_MODE=1 vllm serve Qwen/Qwen3-4B \
CUDA_VISIBLE_DEVICES=1 accelerate launch train_async_grpo.py
```

## Weight synchronization

After each weight sync the trainer pushes the updated policy to the vLLM server. Two choices control how:

- **`weight_sync_mode`**: `"sparse"` (default) or `"full"`.
- `"sparse"` sends **only the bf16 weights that changed** in the step. The changed set is recovered by _inverting_
the AdamW update from the optimizer's resident moments (`θ_old = (θ_t + lr·m̂/(√v̂+ε)) / (1−lr·wd)`) and diffing
against the live weights, so **no pre-step snapshot is kept**. It requires a `torch.optim.AdamW` optimizer (the
trainer raises otherwise) and a vLLM with sparse weight transfer ([vllm-project/vllm#40096](https://github.com/vllm-project/vllm/pull/40096)), served with `--model-impl transformers` and `VLLM_USE_V2_MODEL_RUNNER=0`. A full **anchor** is sent every `weight_sync_anchor_interval` syncs to bound drift.
- `"full"` broadcasts the entire policy every sync. Use it when the optimizer is not AdamW. It is always sent over NCCL.

> [!WARNING]
> `"sparse"` is **not supported for MoE models** (Mixture-of-Experts). vLLM's transformers backend stores the
> experts as a fused buffer (`w13_weight`/`w2_weight`), so the in-place sparse apply cannot address them by their
> Hugging Face parameter names. The trainer raises at initialization for MoE models under `"sparse"` — use
> `weight_sync_mode="full"` for MoE (full sync loads the experts via `load_weights`, which handles the fused
> mapping, including under expert parallelism).
- **`weight_sync_backend`**: the transport for sparse patches: `"nccl"` (default) or `"bucket"`.

| backend | data plane | when to use |
| ---------- | ----------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------- |
| `"nccl"` | NCCL broadcast over a group shared with vLLM | trainer and vLLM **co-located** (same node / NVLink). ~100× faster per sync. |
| `"bucket"` | sparse patches uploaded to an **HF Storage Bucket**, applied in place on vLLM | trainer and vLLM on **different hosts** (e.g. a remote vLLM HF Space). Object-storage latency (~seconds/sync). |

Serve the vLLM side to match the backend:

```bash
# nccl backend (default): co-located trainer + vLLM
VLLM_USE_V2_MODEL_RUNNER=0 vllm serve Qwen/Qwen3-4B --model-impl transformers \
--weight-transfer-config '{"backend":"nccl"}'

# bucket backend: register the engine via the worker extension
VLLM_USE_V2_MODEL_RUNNER=0 vllm serve Qwen/Qwen3-4B --model-impl transformers \
--worker-extension-cls trl.experimental.async_grpo.delta_engine.HFBucketWorkerExtension \
--weight-transfer-config '{"backend":"hf_bucket"}'
```

A runnable end-to-end script for the `"bucket"` backend is in
[`examples/scripts/async_grpo_delta.py`](https://github.com/huggingface/trl/blob/main/examples/scripts/async_grpo_delta.py).

### Disaggregating training and inference

The `"bucket"` backend decouples _where training runs_ from _where generation runs_: the control plane is plain HTTP
(`vllm_server_base_url`) and the data plane is the HF Hub (a bucket reachable from anywhere). Nothing requires the
trainer and the vLLM server to share a machine, a network, or a NCCL group. So you can keep the trainer on your local
training GPUs and serve generation from a **remote vLLM HF Space** (or several, scaled independently), syncing only the
~1% of weights that change each step.

The end-to-end example (local trainer + remote vLLM Space + remote environment Space + bucket) is in
[`examples/scripts/async_grpo_buckets/async_grpo_buckets.py`](https://github.com/huggingface/trl/blob/main/examples/scripts/async_grpo_buckets/async_grpo_buckets.py) (see its `README.md` for the deploy + run guide).

> [!TIP]
> Sparse sync cost is roughly flat in model size (only the changed elements move), while a full broadcast grows with
> the model, so the sparse advantage widens for larger policies. On a single node, `"nccl"` is the fast default; reach
> for `"bucket"` specifically for the cross-host / remote-Space setup.

## Design philosophy

This trainer is intentionally kept minimal and is not meant to grow into a general-purpose solution. If you need a feature that is not supported, we recommend cloning the repository and adapting the trainer to your needs directly. New features will only be considered when there is significant community demand.
Expand Down
3 changes: 3 additions & 0 deletions docs/source/example_overview.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,9 @@ Scripts are maintained in the [`trl/scripts`](https://github.com/huggingface/trl

| File | Description |
| --- | --- |
| [`examples/scripts/async_grpo.py`](https://github.com/huggingface/trl/blob/main/examples/scripts/async_grpo.py) | This script shows how to use the [`experimental.async_grpo.AsyncGRPOTrainer`] to fine-tune a model with fully asynchronous rollouts served by a vLLM server, with sparse delta weight sync over NCCL. |
| [`examples/scripts/async_grpo_delta.py`](https://github.com/huggingface/trl/blob/main/examples/scripts/async_grpo_delta.py) | This script shows how to use the [`experimental.async_grpo.AsyncGRPOTrainer`] with sparse delta weight sync routed through an HF Storage Bucket, for when the trainer and the vLLM server are not in the same NCCL world. |
| [`examples/scripts/async_grpo_buckets/async_grpo_buckets.py`](https://github.com/huggingface/trl/blob/main/examples/scripts/async_grpo_buckets/async_grpo_buckets.py) | This script shows how to run the [`experimental.async_grpo.AsyncGRPOTrainer`] fully disaggregated: a local trainer with a remote vLLM Space and a remote environment Space, syncing weights through an HF Storage Bucket. |
| [`examples/scripts/bco.py`](https://github.com/huggingface/trl/blob/main/examples/scripts/bco.py) | This script shows how to use the [`KTOTrainer`] with the BCO loss to fine-tune a model to increase instruction-following, truthfulness, honesty, and helpfulness using the [openbmb/UltraFeedback](https://huggingface.co/datasets/openbmb/UltraFeedback) dataset. |
| [`examples/scripts/cpo.py`](https://github.com/huggingface/trl/blob/main/examples/scripts/cpo.py) | This script shows how to use the [`experimental.cpo.CPOTrainer`] to fine-tune a model to increase helpfulness and harmlessness using the [Anthropic/hh-rlhf](https://huggingface.co/datasets/Anthropic/hh-rlhf) dataset. |
| [`examples/scripts/distillation.py`](https://github.com/huggingface/trl/blob/main/examples/scripts/distillation.py) | This script shows how to use the [`DistillationTrainer`] to distill a teacher model into a student on-policy, supporting full training and LoRA. |
Expand Down
13 changes: 12 additions & 1 deletion examples/scripts/async_grpo.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,12 +24,23 @@
"""
pip install math_verify

CUDA_VISIBLE_DEVICES=1 VLLM_SERVER_DEV_MODE=1 vllm serve Qwen/Qwen3-0.6B \
AsyncGRPO defaults to *sparse* weight sync over NCCL: only the bf16 weights changed by each optimizer step are
broadcast and applied in place on vLLM (the changed set is recovered by inverting the AdamW step from the resident
optimizer moments — no snapshot kept). This needs a vLLM with sparse weight transfer (vllm-project/vllm#40096), the
`transformers` model impl (so vLLM's runtime param names match the trainer's HF names), and the V1 model runner
(`apply_sparse_weight_patches` is V1-only):

CUDA_VISIBLE_DEVICES=1 VLLM_SERVER_DEV_MODE=1 VLLM_USE_V2_MODEL_RUNNER=0 vllm serve Qwen/Qwen3-0.6B \
--model-impl transformers \
--max-model-len 2048 \
--logprobs-mode processed_logprobs \
--weight-transfer-config '{"backend":"nccl"}'

CUDA_VISIBLE_DEVICES=0 accelerate launch examples/scripts/async_grpo.py

To fall back to broadcasting the full policy every sync (e.g. a non-AdamW optimizer), set
`weight_sync_mode="full"` in the config and serve without the sparse-only flags (a plain
`--weight-transfer-config '{"backend":"nccl"}'` is enough).
"""

from datasets import load_dataset
Expand Down
91 changes: 91 additions & 0 deletions examples/scripts/async_grpo_buckets/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,91 @@
# Disaggregated async GRPO with bucket weight sync (`async_grpo_buckets.py`)

Train a policy on your **local GPU** while a **remote vLLM HF Space** does generation — the two never share a NCCL
group. They stay in sync through an **HF Storage Bucket**: after each optimizer step the trainer uploads only the bf16
weights that changed (a sparse patch, recovered by inverting the AdamW step — no snapshot), and the remote vLLM applies
it in place. A full **anchor** is sent every N syncs to bound drift.

This is the power of disaggregation: **training and inference scale and live independently.** Put the trainer wherever
your training GPUs are, serve generation from an autoscaling Space (or many), and connect any environment server — all
glued together by a bucket and plain HTTP.

```
┌──────────────────────────┐ sparse patches / anchors ┌───────────────────────────┐
│ Local trainer (1 GPU) │ ───────────────────────────────▶ │ HF Storage Bucket │
│ AsyncGRPOTrainer │ │ anchors/ + deltas/ │
│ + rollout worker │ ◀─────────────────────────────── └───────────────────────────┘
└──────────┬───────────────┘ apply in place ▲
│ /v1/completions (HTTP) │ fetch
▼ │
┌──────────────────────────┐ ┌───────────────────────────┐
│ vLLM HF Space (GPU) │ ◀────────────────────────────── │ HFBucketWorkerExtension │
│ serves generation │ │ (hf_bucket backend) │
└──────────┬───────────────┘ └───────────────────────────┘
│ tool calls (HTTP)
▼
┌──────────────────────────┐
│ Wordle env HF Space │ (no GPU; public one at openenv-wordle.hf.space)
└──────────────────────────┘
```

Files in this directory:

- `async_grpo_buckets.py` — the local trainer (AsyncGRPO + Wordle env, `weight_sync_backend="bucket"`).
- `vllm_space/` — Dockerfile + README to deploy the **vLLM inference Space** (GPU).
- `wordle_space/` — Dockerfile + README to deploy your own **Wordle environment Space** (optional; a public one exists).

## Prerequisites

```sh
pip install "trl @ git+https://github.com/huggingface/trl.git@delta-weight-sync-v3"
pip install "openenv-textarena @ git+https://huggingface.co/spaces/openenv/wordle" # the Wordle env client
hf auth login # needs write access to create the bucket + (for Option 1) deploy Spaces
```

The vLLM side needs sparse weight transfer (vllm-project/vllm#40096), which shipped in **vLLM 0.23.0** — both the
Space Dockerfile and a local install just need `vllm>=0.23.0`.

### Step 1 — deploy the vLLM inference Space (GPU)

```sh
# Create the Space (l4 GPU, Docker SDK). HF_TOKEN lets the Space read the bucket.
hf repos create <your-username>/vllm-wordle-inference \
--repo-type space --space-sdk docker

hf upload <your-username>/vllm-wordle-inference \
examples/scripts/async_grpo_buckets/vllm_space/ . --repo-type space

# Set the GPU + secrets/vars in the Space settings (or via the CLI / web UI):
# hardware: l4x1 ; secret HF_TOKEN=<token> ; the Dockerfile already sets VLLM_SERVER_DEV_MODE=1
```

Wait until `https://<your-username>-vllm-wordle-inference.hf.space/health` returns 200 (first build pulls the image and
loads the model — a few minutes).

### Step 2 — (optional) deploy your own Wordle env Space

A public env runs at `https://openenv-wordle.hf.space`. To run your own (higher concurrency), deploy `wordle_space/`
the same way and pass its URL via `--env-url`.

### Step 3 — train locally (1 GPU)

```sh
CUDA_VISIBLE_DEVICES=0 python examples/scripts/async_grpo_buckets/async_grpo_buckets.py \
--model Qwen/Qwen3-1.7B \
--vllm-server-url https://<your-username>-vllm-wordle-inference.hf.space \
--env-url https://openenv-wordle.hf.space \
--weight-sync-bucket-id <your-username>/wordle-deltas
```

The bucket (`<your-username>/wordle-deltas`) is created automatically on the first sync.

## Notes

- **Bucket vs NCCL.** Bucket sync works across hosts/Spaces (data plane = HF Hub, control plane = HTTP), at the cost of
object-storage latency (~seconds/sync). On a single node where the trainer and vLLM share NVLink, the default
`weight_sync_backend="nccl"` is ~100× faster — use the bucket backend specifically for the disaggregated/cross-host
case this example demonstrates.
- **Anchors.** `--weight-sync-anchor-interval N` uploads a full checkpoint every N syncs (sparse deltas in between) to
bound drift from any missed bits. Lower N = more robust, larger uploads.
- **Key flags must match** between the example and the Space: `Qwen/Qwen3-1.7B`, the `hf_bucket` backend, and the
`HFBucketWorkerExtension`.
Loading
Loading