Skip to content

async grpo native weight sync with vllm>=0.22.0 #5892

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

Merged
AmineDiro merged 1 commit into from
Jun 15, 2026
Merged

async grpo native weight sync with vllm>=0.22.0 #5892

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
1 change: 0 additions & 1 deletion trl/experimental/async_grpo/async_grpo_trainer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -442,7 +442,6 @@ def __init__(
"dtype_names": weight_dtype_names,
"shapes": weight_shapes,
"packed": True,
"is_checkpoint_format": True,
},
)
self.rollout_worker = AsyncRolloutWorker(
Expand Down
24 changes: 13 additions & 11 deletions trl/experimental/async_grpo/weight_transfer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@
from trl.import_utils import is_vllm_available


if is_vllm_available(min_version="0.17.1"):
if is_vllm_available(min_version="0.22.0"):
from vllm.distributed.weight_transfer.nccl_engine import NCCLTrainerSendWeightsArgs, NCCLWeightTransferEngine
from vllm.utils.network_utils import get_ip, get_open_port

Expand All @@ -37,9 +37,9 @@ def __init__(
server_timeout: float = 240.0,
init_weight_transfer_timeout: int = 1800,
):
if not is_vllm_available(min_version="0.17.1"):
if not is_vllm_available(min_version="0.22.0"):
raise ImportError(
"vLLM >= 0.17.1 is required to use WeightTransferClient. Install it with: pip install 'vllm>=0.17.1'"
"vLLM >= 0.22.0 is required to use WeightTransferClient. Install it with: pip install 'vllm>=0.22.0'"
)
self.vllm_server_url = vllm_server_url.rstrip("/")
self.server_timeout = server_timeout
Expand Down Expand Up @@ -103,25 +103,27 @@ def send_weights(self, iterator) -> None:
if self.model_update_group is None:
return
t0 = time.time()
# Prepare the workers for the reload; must complete before any weights are sent.
requests.post(
f"{self.vllm_server_url}/start_weight_update",
json={"is_checkpoint_format": True},
timeout=1800,
)
# The /update_weights POST drives the workers' blocking NCCL recv, so it runs on a thread
# concurrently with the trainer-side broadcast.
t_update = threading.Thread(
target=requests.post,
args=(f"{self.vllm_server_url}/update_weights",),
kwargs={"json": {"update_info": self._weight_update_info}, "timeout": 1800},
)
t_update.start()
logger.debug(f"[weight_sync] /update_weights POST sent ({time.time() - t0:.1f}s)")
t_nccl = time.time()
NCCLWeightTransferEngine.trainer_send_weights(
iterator=iterator,
trainer_args=NCCLTrainerSendWeightsArgs(group=self.model_update_group, packed=True),
)
logger.debug(f"[weight_sync] NCCL transfer took {time.time() - t_nccl:.1f}s")
t_join = time.time()
t_update.join()
logger.debug(
f"[weight_sync] /update_weights join took {time.time() - t_join:.1f}s "
f"(total send_weights: {time.time() - t0:.1f}s)"
)
requests.post(f"{self.vllm_server_url}/finish_weight_update", timeout=1800)

@cursor cursor Bot Jun 3, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Finish not called after sync errors

High Severity

send_weights now calls start_weight_update before NCCL and finish_weight_update only on the happy path. If trainer_send_weights, the threaded /update_weights call, or join raises or aborts, finish_weight_update is skipped while vLLM workers may still be in layerwise reload, breaking later rollouts or weight syncs until the server is restarted.

Fix in Cursor Fix in Web

Reviewed by Cursor Bugbot for commit e828e14. Configure here.

@AmineDiro AmineDiro Jun 3, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

mirrors upstream : https://github.com/vllm-project/vllm/tree/main/examples/rl

no example handles this failure mode

logger.debug(f"[weight_sync] send_weights took {time.time() - t0:.1f}s")

def pause(self) -> None:
t0 = time.time()
Expand Down
Loading