[trainer] feat: add V1 trainer support for the RemoteBackend plugin abstraction - #7102
[trainer] feat: add V1 trainer support for the RemoteBackend plugin abstraction#7102sfc-gh-kganesan wants to merge 3 commits into
Conversation
|
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a pluggable out-of-process RemoteBackend abstraction to verl, enabling external RL backends to manage their own GPU compute and weight synchronization. It adds the PPOTrainerRemoteBackend trainer, a RemoteBackendCheckpointEngine adapter, associated registries, CPU-only resource pool support, and comprehensive unit tests. The review feedback highlights three key areas for improvement: avoiding blocking ray.get() calls inside async methods by using asyncio.gather, removing redundant getattr checks for guaranteed abstract methods, and optimizing tensor operations in make_njt by converting to CPU lists before looping to prevent synchronous GPU-to-CPU transfers.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| # the transfer to the actor side (naive: verl actor -> rollout server; | ||
| # remote_backend: plugin owns the transfer, typically CUDA IPC). | ||
| if self.backend in ("naive", "remote_backend"): | ||
| ray.get(self.actor_wg.update_weights(global_steps=global_steps, mode=self.backend)) |
There was a problem hiding this comment.
Avoid using ray.get() inside async methods of Ray actors (or classes running within them like CheckpointEngineManager). Calling ray.get() blocks the event loop, which can lead to performance degradation or deadlocks, and typically triggers warnings in Ray.
Since self.actor_wg.update_weights is non-blocking and returns a list of ray.ObjectRefs, you should asynchronously await them using asyncio.gather.
| ray.get(self.actor_wg.update_weights(global_steps=global_steps, mode=self.backend)) | |
| await asyncio.gather(*self.actor_wg.update_weights(global_steps=global_steps, mode=self.backend)) |
References
- Avoid using
ray.get()insideasyncmethods of Ray actors. This blocks the event loop, which can lead to performance degradation or deadlocks, and typically triggers warnings in Ray. Instead, useawaiton remote calls orasyncio.gatherto parallelize multiple remote calls asynchronously.
| destroy = getattr(backend, "destroy", None) | ||
| if destroy is None: | ||
| return | ||
| result = destroy() |
There was a problem hiding this comment.
Avoid over-defensive programming, such as using redundant getattr and None checks on attributes or methods that are guaranteed to exist.
Since destroy is defined as an @abc.abstractmethod on the RemoteBackend base class, any concrete backend subclass is guaranteed to implement it. Therefore, checking for its existence with getattr(backend, "destroy", None) is redundant.
| destroy = getattr(backend, "destroy", None) | |
| if destroy is None: | |
| return | |
| result = destroy() | |
| result = backend.destroy() |
References
- Avoid over-defensive programming, such as wrapping
inspect.signaturein try-except blocks or using redundantgetattrandNonechecks, unless there is a concrete, documented risk of runtime failure (e.g., dealing with uninspectable C-extensions or attributes that are legitimately optional in supported versions).
| cu_seqlens = data["input_ids"].offsets() | ||
| seq_lengths = cu_seqlens.diff() | ||
| starts = data["attention_mask"].long().argmax(dim=1) | ||
| pieces = [tensor[b, starts[b].item() : starts[b].item() + seq_lengths[b].item()] for b in range(tensor.shape[0])] |
There was a problem hiding this comment.
Calling .item() on PyTorch tensors inside a loop over the batch size causes synchronous GPU-to-CPU transfers for every iteration. This introduces significant host-device synchronization overhead and can become a major performance bottleneck.
Instead, convert the entire starts and seq_lengths tensors to CPU lists once using .tolist() before the list comprehension, and then index into those lists.
| pieces = [tensor[b, starts[b].item() : starts[b].item() + seq_lengths[b].item()] for b in range(tensor.shape[0])] | |
| starts_cpu = starts.tolist() | |
| seq_lengths_cpu = seq_lengths.tolist() | |
| pieces = [tensor[b, starts_cpu[b] : starts_cpu[b] + seq_lengths_cpu[b]] for b in range(tensor.shape[0])] |
|
@wuxibin89 can you please help review the V1 port? |
|
|
||
|
|
||
| @register_trainer("remote_backend") | ||
| class PPOTrainerRemoteBackend(PPOTrainerSync): |
There was a problem hiding this comment.
Does remote backend support async training?
There was a problem hiding this comment.
we will add async support in a followup PR shortly, this PR is to land the remote abstract backend + sync colocate support :)
wuxibin89
changed the title
| from omegaconf import DictConfig | ||
|
|
||
|
|
||
| class RemoteBackend(abc.ABC): |
There was a problem hiding this comment.
Please add RemoteBackend design doc to docs/index.rst.
There was a problem hiding this comment.
Added docs/advance/remote_backend.md and hooked it into docs/index.rst under the Advanced Features section (next to agent_loop / reward_loop / data/transfer_queue). Covers the ABC + registry, the V1 trainer subclass, the CheckpointEngineManager short-circuit + set_global_steps fan-out, the Hydra choice hook, and points at the Arctic-Platform reference implementation. Pushed in 3fe9463.
Addresses review comment on verl-project#7102 (wuxibin89): describe the RemoteBackend ABC + registry, the V1 trainer subclass, the CheckpointEngineManager short-circuit, and how a downstream plugin wires itself in via VERL_USE_EXTERNAL_MODULES. Links the reference Arctic-Platform implementation and lists the CPU-only tests. Co-authored-by: Cursor <cursoragent@cursor.com>
|
/gemini review |
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
Fixes the check_license.py assertion in the pre-commit / pre_commit_for_ppo jobs (verl-project#7102). The file was empty; adding the standard Apache 2.0 header.
|
@wuxibin89 fixed in 5cbfeb9 — |
|
@wuxibin89 the failure was |
|
@sfc-gh-kganesan We will release v0.8.1 in next two weeks, hold this PR until release is done. |
|
@wuxibin89 just checking you had more precise release time for v0.8.1? Thanks! |
|
@wuxibin89 congrats on v0.9.0 release. Can we move forward with this PR, which also meets the need of other users? |
Extends the V0 RemoteBackend seam (verl-project#6422) to the V1 trainer stack so out-of-process RL backends (e.g. Arctic-Platform) can plug into V1 without touching verl's V0 or V1 hot paths. verl-core: * verl/remote_backend/{base,worker_utils,__init__}.py: RemoteBackend ABC + registry (name -> backend class + lazy forwarder loader) and backend-agnostic tensor/metric helpers. Ported from V0 unchanged. * verl/checkpoint_engine/remote_backend.py: no-op CheckpointEngine registered under 'remote_backend'. send/receive are never called on the short-circuit path. * verl/checkpoint_engine/{__init__,base}.py: eager import above + CheckpointEngineManager.update_weights short-circuit for backend='remote_backend' with set_global_steps fan-out to replicas that implement it. * verl/trainer/ppo/v1/trainer_base.py: three PPOTrainer extension hooks (_actor_rollout_wg_extra_kwargs, _llm_server_replica_init_kwargs, _checkpoint_engine_backend). Defaults preserve existing V1 behavior. * verl/trainer/ppo/v1/trainer_remote_backend.py: PPOTrainerRemoteBackend (subclass of PPOTrainerSync). Builds the plugin's RemoteBackend on the driver, smuggles reconnect_handle to forwarder workers and rollout replicas, marks the resource pool CPU-only, and asserts single-forwarder when the backend requires it. * verl/trainer/ppo/v1/__init__.py: re-export. * verl/trainer/config/ppo_trainer.yaml: 'optional remote_backend@remote_backend: null' defaults entry. * verl/trainer/main_ppo.py: _resolve_remote_backend_from_hydra_choice mirrors the 'remote_backend=<name>' Hydra choice onto trainer.remote_backend and trainer.v1.trainer_mode. * verl/workers/rollout/{replica,llm_server}.py: RolloutReplica accepts **kwargs; LLMServerManager accepts replica_init_kwargs and forwards to replica constructors. * verl/single_controller/ray/base.py: ResourcePoolManager.use_gpu (default True); CPU-only pools skip the GPU headroom check. * tests/remote_backend/: 9 CPU-only tests (registries, trainer dispatch, Hydra choice hook). Validation: 8B BIRD text-to-SQL, recipe-aligned, 5 steps. update_actor 380.9s -> 166.0s (2.30x from ZoRRo), total step 731.4s -> 438.3s (1.67x), convergence preserved (reward means match step-over-step). Net diff: +119 / -6 lines in existing files, ~460 new lines in the plugin-seam code and tests. Co-authored-by: Cursor <cursoragent@cursor.com>
Addresses review comment on verl-project#7102 (wuxibin89): describe the RemoteBackend ABC + registry, the V1 trainer subclass, the CheckpointEngineManager short-circuit, and how a downstream plugin wires itself in via VERL_USE_EXTERNAL_MODULES. Links the reference Arctic-Platform implementation and lists the CPU-only tests. Co-authored-by: Cursor <cursoragent@cursor.com>
Fixes the check_license.py assertion in the pre-commit / pre_commit_for_ppo jobs (verl-project#7102). The file was empty; adding the standard Apache 2.0 header.
5cbfeb9 to
6329d45
Compare
|
Rebased onto latest
|
4 participants
V1 trainer support for the
RemoteBackendplugin abstractionSummary
Adds V1 trainer support for the
RemoteBackendplugin abstraction. The V0seam for out-of-process RL backends (training + sampling that live outside
verl core, e.g. Arctic-Platform's DeepSpeed/vLLM stack) lands in
#6422; this PR extends the same seam to V1 without
regressing any V1 path and without requiring the V0 forwarder to be
rewritten.
Design goal (from the V0 review): all GPU work happens on the plugin side;
verl-core drives it through a thin CPU-only forwarder actor. This PR
preserves that invariant.
Contract
Users opt in with a Hydra group choice
remote_backend=<name>. Inmain_ppo.main,_resolve_remote_backend_from_hydra_choicestampstrainer.remote_backend=<name>and (unless the user already picked acustom mode)
trainer.v1.trainer_mode=remote_backend. That mode dispatchesto
PPOTrainerRemoteBackend, which:RemoteBackend.from_config(config)on thedriver and captures its
reconnect_handle().PPOTrainer:_actor_rollout_wg_extra_kwargs-> smuggles{main_config, backend_handle}into every actor-rollout worker via
RayClassWithInitArgs._llm_server_replica_init_kwargs-> forwards the same handle intoevery
RolloutReplicaviaLLMServerManager.replica_init_kwargs._checkpoint_engine_backend->"remote_backend"._init_resource_pool_mgrto swap in the plugin-providedworker class and mark the resource pool CPU-only. The forwarder does no
GPU work; the plugin already claims all GPUs internally, so
double-booking through Ray would starve it ("Total available GPUs 0"
at placement time).
reports
requires_single_forwarder()(Arctic does): otherwiseDispatch.ONE_TO_ALLcalls duplicate against the single backend andmesh-dispatched compute fragments the global batch.
CheckpointEngineManager.update_weightsgets one new short-circuit: forbackend='remote_backend'it invokesactor_wg.update_weights(...)(theforwarder relays to the plugin, which does the transfer, typically CUDA
IPC), then fans
set_global_steps(...)out to any replicas thatimplement it so downstream metrics can tag rollouts with their policy
version.
Files changed
verl-core:
RemoteBackendABC + registry (name -> backend class + lazy forwarderloader) + backend-agnostic tensor/metric helpers. Ported from V0
unchanged.
engine stub registered under
"remote_backend".send_weights/receive_weightsare never called on the short-circuit path.CheckpointEngineManager.update_weightsalso short-circuits forbackend='remote_backend'and fans outset_global_stepsto replicasthat implement it.
PPOTrainerextensionhooks. Defaults preserve existing V1 behavior.
PPOTrainerRemoteBackend, ~170 lines.optional remote_backend@remote_backend: nulldefaults entry._resolve_remote_backend_from_hydra_choice.RolloutReplica.__init__accepts**kwargs(forwarded fromLLMServerManager.replica_init_kwargs).LLMServerManageracceptsreplica_init_kwargsand forwards to replica constructors.ResourcePoolManager.use_gpu(default
True, no behavior change for existing paths). CPU-only poolsskip the GPU headroom check in
_check_resource_available.lazy worker loader, checkpoint-engine dispatch, trainer dispatch, and
the Hydra choice hook (9 tests).
Net diff: +119 / -6 lines in existing files, plus ~460 lines of
new plugin-seam code + tests.
Plugin side (out of tree, for context)
The Arctic-Platform migration is confined to a new
arctic_platform/integrations/verl/v1/subpackage(
worker.py~430 LoC,replica.py~100 LoC,server.py~75 LoC). V0modules are untouched. Register-time auto-detection picks V0 vs V1 based
on whether
verl.trainer.ppo.v1.trainer_remote_backendimports.Validation
BIRD text-to-SQL,
data.max_prompt_length=32768,data.max_response_length=4096,train_batch_size=128,rollout.n=16,ppo_max_token_len_per_gpu=98304,recipe-aligned with Arctic-Platform's
recipe/skyrl-integration/recipes/rl/verl/txt2sql/run_qwen3_32b_bird_grpo_arl_zorro_yes.sh.Model: Qwen3-8B. 5 steps per config on 8× H200.
Speedup (per-step, s):
timing_s/update_actortiming_s/old_log_probtiming_s/gentiming_s/step(total)The wall-clock win is dominated by ZoRRo's fused actor update.
genisneutral: with the current Arctic recipe knob
enforce_eager=FalseandvLLM 0.18.0's new
cudagraph_mode=FULL_AND_PIECEWISEdefault, FCA issilently disabled at runtime. Reverting to the pre-plugin
enforce_eager=Truere-enables FCA (all 8 inference workers logForest Cascade Attention ENABLED (... cudagraph_mode=NONE)) and dropsgenfrom ~205 s to ~191 s on 8B/BIRD (~7% inference speedup, 4-stepaverage). This knob is on the Arctic-Platform side and is being tracked
in a separate PR against the plugin.
Convergence: rewards stay in the same band across both paths (step-1
critic/score/mean: stock 0.635 vs Arctic 0.635; step-5 stock 0.63 vsArctic 0.63), with matching response-length trajectories. No divergence.
Test plan
pytest tests/remote_backend/— 9 CPU-only tests (registries, trainerdispatch, Hydra choice hook). All green.
ruff checkandruff format --checkon all modified files: clean.via the
remote_backend=<name>Hydra choice.Notes for reviewers
ResourcePoolManager.use_gpudefaults toTrue; onlyremote_backendcallers flip it. Existing non-plugin paths keep the current behavior
bit-for-bit.
set_global_stepsfan-out inCheckpointEngineManageris gated onhasattr(r, 'set_global_steps'); verl's built-inRolloutReplicadoesn't implement it, so it's a no-op for the naive path.
RemoteBackendandRemoteBackendRegistryare ported from V0 unchangedso the two trainers can share a single ABC surface. No new methods; the
V1-specific plumbing lives entirely in
PPOTrainerRemoteBackendandthe extension hooks in
PPOTrainer.Companion PR (Arctic-Platform plugin V1 shims + from-scratch setup guide): Snowflake-AI-Research/Arctic-Platform#41