[refactor] Rename arctic_rl config namespace to generic remote_backend.arctic by sfc-gh-kganesan · Pull Request #3 · Snowflake-AI-Research/verl

sfc-gh-kganesan · 2026-05-21T13:56:44Z

Summary

Addresses @zw0610's review on verl-project/verl#6422 asking that core paths use generic names so other remote backends can dock without colliding with the Arctic-specific arctic_rl namespace.

Changes (3 files, +55/-41)

verl/trainer/config/ppo_trainer.yaml — top-level arctic_rl: block becomes remote_backend.arctic: (nested under a generic remote_backend: parent). A header comment explains that new backends drop in as sibling blocks (remote_backend.<name>:) with no other changes required.
verl/trainer/ppo/arctic_rl_client.py — adapter now caches `self._backend_config = config.remote_backend[self._BACKEND_CONFIG_KEY]` in `init` and reads all backend-specific fields through it. `_BACKEND_CONFIG_KEY = "arctic"` is the single source of truth a new backend would override.
`examples/arctic_rl/run_gsm8k_grpo_arl_zorro_yes.sh` — 9 hydra overrides switched from `arctic_rl.` to `remote_backend.arctic.`.

Intentionally not renamed

Adapter-side filenames/classes (`arctic_rl_client.py`, `ArcticRLClientWrapper`, `verl/workers/rollout/arctic_rollout/`) and the `"arctic"` rollout-replica registry entry are left alone since they're adapter-scoped and follow the same shape as the existing `vllm`/`sglang`/`trtllm` entries.

Test

6-step GSM8K GRPO smoke on H200:

step	MFU	loss	score
1	148.7	0.0475	0.200
2	174.3	0.0196	0.213
3	156.4	0.0289	0.138
4	158.8	0.0438	0.200
5	156.7	0.0250	0.275
6	160.7	0.0325	0.262

MFU matches the pre-rename baseline (170-ish); no perf regression.

Once merged

PR verl-project#6422's head on `verl-project/verl` picks this up automatically, addressing @zw0610's comment without any further work.

Made with Cursor

@zw0610

…kend.arctic` Addresses @zw0610's review on verl-project#6422 asking that core paths use generic names so other backends can dock without colliding with `arctic_rl`. Changes: * `verl/trainer/config/ppo_trainer.yaml`: top-level `arctic_rl:` block becomes `remote_backend.arctic:` (nested under a generic `remote_backend:` parent). Comment explains that new backends drop in as sibling blocks (`remote_backend.<name>:`) with no other changes required. * `verl/trainer/ppo/arctic_rl_client.py`: adapter now caches `self._backend_config = config.remote_backend[self._BACKEND_CONFIG_KEY]` in `__init__` and reads all backend-specific fields through it. `_BACKEND_CONFIG_KEY = "arctic"` is the single source of truth a new backend would override. * `examples/arctic_rl/run_gsm8k_grpo_arl_zorro_yes.sh`: 9 hydra overrides switched from `arctic_rl.*` to `remote_backend.arctic.*`. Adapter-side filenames/classes (`arctic_rl_client.py`, `ArcticRLClientWrapper`, `verl/workers/rollout/arctic_rollout/`) and the `"arctic"` rollout-replica registry entry are intentionally left alone since they're adapter-scoped and follow the same shape as the existing `vllm`/`sglang`/`trtllm` entries. Smoke: 6-step GSM8K GRPO on H200, MFU 148-174, matches pre-rename baseline. Co-authored-by: Cursor <cursoragent@cursor.com>

sfc-gh-truwase · 2026-05-21T14:01:18Z


-  # ray or http: ray would be much faster for payload comms in the en-prem use case
-  comm_protocol: ray
+    # only set to True if running into OOM, otherwise it'll be faster not to tile (extra forward replay and smaller chunks to compute)


Remove this and modify corresponding code to handle missing key and use default True for this key.

Done in b9d2914 — removed the yaml line; adapter reads it via self._backend_config.get("tiled_logits_compute", True) so the default is True when absent. Also dropped the now-redundant remote_backend.arctic.tiled_logits_compute=True override from the GSM8K example.

sfc-gh-truwase · 2026-05-21T14:08:44Z

Can you cherry-pick changes to arctic_rl_client.py and arctic_rollout.py in this commit:
https://github.com/snowflake-eng/arctic-verl/commit/662e421bf2f2ef3508b11504cf106c27d4191960 to address feedback on verl-project#6422

sfc-gh-truwase · 2026-05-21T14:16:11Z

+        n_training_gpus = self._backend_config.get("training_gpus", self.config.trainer.n_gpus_per_node)
+        n_sampling_gpus = self._backend_config.get("sampling_gpus", self.config.trainer.n_gpus_per_node)
+        n_log_prob_gpus = self._backend_config.get("log_prob_gpus", self.config.trainer.n_gpus_per_node)
+        colocate = self._backend_config.get("colocate", False)


Make the default True

Done in b9d2914 — colocate = self._backend_config.get("colocate", True) now defaults to True.

… tokenizer * `ppo_trainer.yaml`: drop `tiled_logits_compute` (advanced OOM-only knob); adapter now reads it via `self._backend_config.get("tiled_logits_compute", True)`. * `arctic_rl_client.py`: flip in-code default for `colocate` from `False` to `True` per Tunji's comment. * Cherry-pick of arctic-verl@662e421b — adapted for the version of `arctic_training` shipping in this PR (its client RPCs are still sync; only `sync_weights` is async). Plumbs `async`/`await` through the full stack so the forwarder's Ray event loop stays responsive, also addressing the gemini-bot blocking-call comment on verl-project#6422: - `RemoteBackend.compute_log_prob`, `update_actor`, `save_checkpoint` declared `async` on the ABC. - `ArcticRLClientWrapper` implements them `async` and dispatches the underlying sync `self._client.fwd_no_grad` / `fwd_bwd` / `step` / `save_checkpoint` via `asyncio.to_thread`, so the worker thread blocks instead of the asyncio loop. - `RemoteBackendActorRolloutRefWorker._run_log_prob` / `_run_update_actor` made async and `await` the backend. `save_checkpoint` / `destroy` likewise await. - `RemoteBackendTrainer.destroy` (sync shutdown path) bridges with `asyncio.run` to await the now-async `backend.destroy`. * `arctic_rollout.py`: drop hardcoded `Qwen/Qwen3-0.6B` tokenizer load. `ArcticLLMEngine.__init__` takes a tokenizer, and `ArcticLLMServer` passes `self.model_config.tokenizer` (already loaded by `HFModelConfig.__post_init__`), avoiding both a hardcoded model name and a redundant download. * `run_gsm8k_grpo_arl_zorro_yes.sh`: drop the now-redundant `remote_backend.arctic.tiled_logits_compute=True` override (the in-code default already gives `True`). Smoke: 6-step GSM8K GRPO on H200, MFU 149-170 (matches pre-change baseline 148-174); no regressions. Co-authored-by: Cursor <cursoragent@cursor.com>

sfc-gh-kganesan · 2026-05-21T16:44:51Z

Cherry-picked in b9d2914 (karthik/wang-rename).

Heads up: the version of arctic_training that ships with this PR still has sync fwd_no_grad / fwd_bwd / step / save_checkpoint (only sync_weights is async). A literal await self._client.fwd_no_grad(...) therefore raises TypeError: object dict can't be used in 'await' expression.

I kept the same end goal — the forwarder's Ray asyncio loop must not block on backend RPCs (this is also the gemini-bot blocking-call comment on verl-project#6422) — but went through asyncio.to_thread so it works against the current arctic_training:

async def _send_compute_log_prob(self, payload):
    response = await asyncio.to_thread(self._client.fwd_no_grad, payload, reference_model=False)
    ...

Stack-wide changes:

RemoteBackend.compute_log_prob / update_actor / save_checkpoint declared async on the ABC.
Worker _run_log_prob / _run_update_actor made async and await the backend.
Sync trainer destroy bridges with asyncio.run.
arctic_rollout.py: dropped hardcoded Qwen/Qwen3-0.6B, tokenizer now passed in via self.model_config.tokenizer (loaded once by HFModelConfig.__post_init__).

Smoke (6-step GSM8K GRPO, H200): MFU 149-170, matches the pre-change 148-174 baseline.

sfc-gh-truwase · 2026-05-21T16:55:25Z

+    async def _send_compute_log_prob(self, payload: dict):
        payload["processing"] = {"post": ["compute_entropy_and_logprobs"], "loss_fn": None}
-        response = self._client.fwd_no_grad(payload, reference_model=False)
+        response = await asyncio.to_thread(self._client.fwd_no_grad, payload, reference_model=False)


Use simpler await

Suggested change

response = await asyncio.to_thread(self._client.fwd_no_grad, payload, reference_model=False)

response = await self._client.fwd_no_grad(payload, reference_model=False)

Tried bare await self._client.fwd_no_grad(...) first — crashes with TypeError: object dict can't be used in 'await' expression.

I checked the latest arctic_training (HEAD on the verl-integration branch) and verified the runtime signatures:

fwd_no_grad sync fwd_bwd sync step sync save_checkpoint async ← newly async sync_weights async

So bare await works for save_checkpoint (line 464) but not for fwd_no_grad (this line), fwd_bwd (L458), or step (L459) — they're still sync upstream.

Two paths:

Now: keep asyncio.to_thread(self._client.fwd_no_grad, ...) so the forwarder's asyncio loop stays unblocked (gemini-bot G3 on [trainer] feat: generic remote backend abstraction for RL Training verl-project/verl#6422 made the same point) and the verl PR ships against the current arctic_training.

Later: when fwd_no_grad/fwd_bwd/step go async upstream, I drop to_thread in a small follow-up — happy to do that as soon as you point me at the commit.

I switched save_checkpoint to bare await in 188338b since that one is already async upstream.

sfc-gh-truwase · 2026-05-21T16:56:03Z


-        fwd_bwd_response = self._client.fwd_bwd(payload)
-        step_response = self._client.step()
+        fwd_bwd_response = await asyncio.to_thread(self._client.fwd_bwd, payload)


Just await

Suggested change

fwd_bwd_response = await asyncio.to_thread(self._client.fwd_bwd, payload)

fwd_bwd_response = await self._client.fwd_bwd(payload)

Same as the fwd_no_grad thread — fwd_bwd is still sync upstream, so kept the asyncio.to_thread wrapper to avoid TypeError: object dict can't be used in 'await' expression. Happy to switch to bare await once it goes async upstream.

Retry await. It should now work after ArcticTraining-dss commit: af1ab8d7727e2611ae8eaf9059f5c04170ad02cd

Done in 8aec8a6 — bumped arctic_training to 983d58c (which includes your af1ab8d) in the smoke env, verified the new signatures (fwd_no_grad / fwd_bwd / step all async), and switched to bare await self._client.fwd_bwd(payload). Also did the same for _send_compute_log_prob / _send_compute_ref_log_prob for symmetry, and dropped the now-unused import asyncio from arctic_rl_client.py.

Smoke (6-step GSM8K GRPO on H200): MFU 186-218, matches the prior asyncio.to_thread baseline (199-224) — no perf regression from dropping the thread hop.

sfc-gh-truwase · 2026-05-21T16:56:18Z

-        fwd_bwd_response = self._client.fwd_bwd(payload)
-        step_response = self._client.step()
+        fwd_bwd_response = await asyncio.to_thread(self._client.fwd_bwd, payload)
+        step_response = await asyncio.to_thread(self._client.step)


Same as the fwd_no_grad thread — step is still sync upstream, so kept the asyncio.to_thread wrapper. Happy to switch to bare await once it goes async upstream.

Retry await. It should now work after ArcticTraining-dss commit: af1ab8d7727e2611ae8eaf9059f5c04170ad02cd

Done in 8aec8a6 — same pass as the fwd_bwd thread, step is now await self._client.step() (it's async on arctic_training @ 983d58c).

sfc-gh-truwase · 2026-05-21T16:56:24Z

-    def save_checkpoint(self):
-        return self._client.save_checkpoint()
+    async def save_checkpoint(self):
+        return await asyncio.to_thread(self._client.save_checkpoint)


Done in 188338b — save_checkpoint is async upstream, so bare await self._client.save_checkpoint() works (dropped the to_thread wrapper).

Smoke (6-step GSM8K GRPO on H200, post-arctic_training bump): MFU 199-224.

Closing `}` of the returned dict was indented at column 4 instead of column 8 — visually breaks the dict nesting, even though Python accepts it. Trivial cosmetic fix flagged by gemini-bot on verl-project#6422. Smoke: 6-step GSM8K GRPO on H200, MFU 146-175, matches baseline. Co-authored-by: Cursor <cursoragent@cursor.com>

`arctic_training.arctic_rl.client.ArcticRLRayClient.save_checkpoint` is `async` on the current `tunji/verl_integration` HEAD, so the `asyncio.to_thread` wrapper is unnecessary. `fwd_no_grad` / `fwd_bwd` / `step` are still sync upstream and continue to use `to_thread` (per inline-reply on PR Snowflake-AI-Research#3 explaining the test). Smoke: 6-step GSM8K GRPO on H200, MFU 199-224 (post-arctic_training bump), no regressions. Co-authored-by: Cursor <cursoragent@cursor.com>

Now that ArcticTraining-dss commit af1ab8d (`Async+bf16 opt`) makes `fwd_no_grad`, `fwd_bwd`, and `step` async on `ArcticRLRayClient`, `asyncio.to_thread` is no longer needed for these three. Switching to bare `await self._client.<op>(...)` per Tunji's retry-await comments on PR Snowflake-AI-Research#3 (lines 458, 459) — symmetry with the already-await'd `save_checkpoint` and `sync_weights`. Also drops the now-unused `import asyncio` from `arctic_rl_client.py`. `RemoteBackendTrainer.destroy` still uses `asyncio.run` to bridge the sync shutdown path to the async `backend.destroy`. Smoke: 6-step GSM8K GRPO on H200, MFU 186-218; matches the `asyncio.to_thread` baseline (199-224), no perf regression from dropping the thread hop. Co-authored-by: Cursor <cursoragent@cursor.com>

@zw0610

Addresses @zw0610's structural review on verl-project#6422: - #2: verl/trainer/ppo/arctic_rl_client.py -> verl/workers/remote_client/arctic_rl_client.py - #3: verl/workers/rollout/arctic_rollout/ -> verl/workers/rollout/remote_rollout/arctic_rollout/ - #4: TODO(@zw0610) on RolloutReplicaRegistry registrations in verl/workers/rollout/replica.py:360 about lazy-init each option to avoid pulling in vLLM + SGLang + TRT-LLM transitive deps eagerly. Pure code-motion + a TODO; no behavior change. Imports inside the moved files and in workers/rollout/replica.py updated. Co-authored-by: Cursor <cursoragent@cursor.com>

@zw0610

Restructure remote_backend per @zw0610 review (stacked on #3)

sfc-gh-kganesan requested a review from sfc-gh-truwase May 21, 2026 13:57

sfc-gh-truwase reviewed May 21, 2026

View reviewed changes

sfc-gh-kganesan and others added 2 commits May 21, 2026 17:00

sfc-gh-truwase approved these changes May 22, 2026

View reviewed changes

sfc-gh-truwase merged commit 17d2aeb into Snowflake-AI-Research:arctic_rl_share_v0.7.1 May 22, 2026
1 check passed

sfc-gh-kganesan mentioned this pull request May 23, 2026

Restructure remote_backend per @zw0610 review (stacked on #3) #4

Merged

sfc-gh-kganesan added a commit that referenced this pull request May 29, 2026

Merge pull request #4 from sfc-gh-kganesan/karthik/wang-restructure

bf8ea8f

Restructure remote_backend per @zw0610 review (stacked on #3)

	response = await asyncio.to_thread(self._client.fwd_no_grad, payload, reference_model=False)
	response = await self._client.fwd_no_grad(payload, reference_model=False)

	fwd_bwd_response = await asyncio.to_thread(self._client.fwd_bwd, payload)
	fwd_bwd_response = await self._client.fwd_bwd(payload)

Conversation

sfc-gh-kganesan commented May 21, 2026

Summary

Changes (3 files, +55/-41)

Intentionally not renamed

Test

Once merged

Uh oh!

sfc-gh-truwase May 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sfc-gh-truwase commented May 21, 2026 • edited by sfc-gh-kganesan Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sfc-gh-kganesan commented May 21, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sfc-gh-kganesan May 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sfc-gh-kganesan May 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sfc-gh-kganesan May 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sfc-gh-kganesan May 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

sfc-gh-truwase May 21, 2026 •

edited

Loading

sfc-gh-truwase commented May 21, 2026 •

edited by sfc-gh-kganesan

Loading

sfc-gh-kganesan May 21, 2026 •

edited

Loading

sfc-gh-kganesan May 21, 2026 •

edited

Loading

sfc-gh-kganesan May 21, 2026 •

edited

Loading

sfc-gh-kganesan May 21, 2026 •

edited

Loading