feat: integrate Arctic RL backend (opt-in via [arctic]) - #1
Draft
sfc-gh-kganesan wants to merge 2 commits into
Draft
feat: integrate Arctic RL backend (opt-in via [arctic])#1sfc-gh-kganesan wants to merge 2 commits into
sfc-gh-kganesan wants to merge 2 commits into
Conversation
sfc-gh-kganesan
force-pushed
the
feat/arctic-integration
branch
2 times, most recently
from
May 27, 2026 07:53
5b26b03 to
a018632
Compare
Adds an opt-in alternative training backend for prime-rl. All integration
code lives under integrations/arctic-rl/; core prime-rl gains a small
generic dispatch hook (~66 LOC, no integration name).
Enabling: any existing recipe flips to Arctic via a one-line overlay,
no edits to the base recipe:
uv run rl @ configs/gsm8k/rl.toml @ integrations/arctic-rl/examples/arctic_overlay.toml
The launcher peeks at trainer.backend before strict config parse and
dispatches via importlib.import_module(f"{backend}.entrypoint").main().
Default is "native"; when unset the integration package is never imported.
Core changes (5 files, +66/-2):
- packages/prime-rl-configs: backend field on TrainerConfig;
client_type + class_path fields on ClientConfig
- src/prime_rl/entrypoints/rl.py: _peek_backend dispatch
- src/prime_rl/orchestrator/orchestrator.py: honor client_type;
no-op update_weights for custom clients (they own their own
weight sync)
- src/prime_rl/utils/client.py: forward class_path to vf.ClientConfig
Integration (~2.35k LOC, all under integrations/arctic-rl/):
- ArcticClient: a verifiers.Client subclass that batches and proxies
rollouts to the Arctic RL coordinator. No subprocess, no monkey-patch.
- arctic-trainer: single-process HTTP client to the Arctic RL coordinator
(/fwd-bwd, /step, /sync-weights). Reuses prime-rl's DataLoader and
scheduler, writes the STABLE marker after each weight sync so the
orchestrator's polling loop is unchanged.
Validation (4xH200, wandb offline):
- Native dispatch regression check on reverse_text, alphabet_sort,
gsm8k: convergence preserved.
- Arctic-mode E2E convergence on reverse_text: 0.105 -> 0.833 peak
over 30 steps (7.1x), monotonic ramp, plateau by step ~18.
- Arctic-mode E2E on gsm8k via overlay (native recipe untouched):
0.561 -> 0.792 peak over 30 steps. Aggregate: 480+ HTTP 200,
0 5xx, ~7700 rollouts with 0 failures.
Companion: sfc-gh-kganesan/verifiers#1 (~30 LOC)
adds client_type="custom" + class_path field on the verifiers ClientConfig.
Out of scope: multi-node, LoRA, multi-run, teacher_inference, multimodal.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
sfc-gh-kganesan
force-pushed
the
feat/arctic-integration
branch
from
May 27, 2026 08:05
a018632 to
abb08b4
Compare
Targets ArcticTraining-dss `arctic_training.arctic_rl` API (tunji/skyrl_integration),
which carries `ds_worker_config` for ZoRRO and uses the `{batch, meta, processing}`
fwd-bwd wire format.
- client.py: build ds_config (AdamW + grad_clip=1.0 + grad_accum + bf16) and, when
`[arctic].use_zorro`, ds_worker_config (use_zorro, response_len, rollout_n, ...).
- config.py: ZoRRO config knobs (use_zorro, zorro_response_len, zorro_rollout_n, ...).
- context.py: `_zorro_layout()` reformats prime-rl's variable [prompt|response]
rollouts into the patcher's required [left-pad prompt | right-pad response=last
response_len] layout (response-only RL tensors). This is the fix that makes ZoRRO
CONVERGE (reverse-text reward 0.086->0.76, matching the no-ZoRRO baseline).
- trainer.py: send {batch: model_kwargs, meta: context}; include position_ids
(required by ZoRRO PromptDeduplicator); resolve response_len.
- verifiers_backend.py: send a single sampling_params dict (server broadcasts).
- patches.py: no-op the DP pause/resume monkey-patch on vLLM 0.18 (renamed
`_pause_complete`).
Verified E2E on reverse-text (converges with and without ZoRRO) and gsm8k
(trains stably with ZoRRO engaged).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Opt-in Arctic RL training backend for prime-rl (
[arctic]block,trainer.backend = "arctic_rl"), now with ZoRRO (prompt-dedup) support that converges.What works (verified E2E)
Key changes
arctic_training.arctic_rlAPI; builds DeepSpeedds_config(AdamW +gradient_clipping=1.0+ grad-accum + bf16) and, when[arctic].use_zorro,ds_worker_config(use_zorro, response_len, rollout_n, …).use_zorro,zorro_response_len,zorro_rollout_n, …)._zorro_layout()reformats prime-rl's variable[prompt|response]rollouts into the layout the server'sQwen3ModelOncePatcherrequires ([left-pad prompt | right-pad response = last response_len tokens], response-only RL tensors). This is the fix that makes ZoRRO converge rather than diverge.{batch, meta, processing}fwd-bwd wire format; includesposition_ids(required by ZoRRO PromptDeduplicator); resolvesresponse_len.sampling_paramsdict (server broadcasts across the batch)._pause_complete).Sibling-repo dependencies (NOT in this PR — needed server-side)
tunji/skyrl_integration(commit0e26105) — the only branch that fully wires ZoRRO (client sendsds_worker_config; servermodel_patch_in_zorro→Qwen3ModelOncePatcher). Requires two small local patches not yet upstreamed:arctic_rl/processors/pipeline.py:model_outputs["logprobs"] = outputs.log_probs(the ZoRRO patcher emitslog_probs; every loss_fn readslogprobs).arctic_rl/server.py:GenerateRequest.prompts: Any+model_config = {"extra": "allow"}(accept token-id prompts).mert/verl_integration_vllm18(commit4cfc9d2).🤖 Generated with Claude Code