[arctic_rl] repoint to arctic_platform.rl + apply verl PR #6 correctness fixes; address reviewer comments - #5
Merged
sfc-gh-kganesan merged 25 commits intoJun 26, 2026
Conversation
Switches the SkyRL Arctic RL integration from the deprecated
`arctic_training.arctic_rl` client to the new `arctic_platform.rl`
client. Scope is intentionally minimal: only the import paths, client
construction, async/await plumbing, and response/argument shapes that
the new client requires. No design changes to the SkyRL-side dispatch
or to core SkyRL; recipe authors keep using
`python -m skyrl.train.entrypoints.main_base trainer.backend=arctic_rl`
unchanged. All correctness/protocol work targeting verl parity moves to
arctic_platform.rl in a separate PR.
API deltas adapted (5 files, +84/-43):
- import: arctic_training.arctic_rl.{client,config}
-> arctic_platform.rl.{create_arctic_rl_client, ArcticRLClientConfig}
- client construction: ArcticRLClient(cfg)
-> create_arctic_rl_client(cfg, server_state)
- async client: WorkerDispatch is sync, so dispatch wraps coroutines
with a tiny `_run(coro)` helper (asyncio.run with thread-pool
fallback); already-async paths (`save_weights_for_sampler`,
`_ArcticInferenceEngineStub`, fully-async sync hook) use `await`
directly.
- response shapes: fwd_no_grad now `reference_model=True/False` (no
`post_processors=`); outputs land under `result["batch"]` (was
`model_outputs`); `step()` metrics under `result["metrics"]`.
- weight sync kwargs: sync_weights(cuda_ipc=, low_memory=) plumbed
from `cfg.trainer.arctic_rl.{cuda_ipc_weight_sync,low_memory_weight_sync}`
(default False; preserves public behavior).
- entrypoint pre-init: capture `pre_client.get_server_state()` and
forward it through `skyrl_entrypoint` to the reconnecting client
(Ray comm-protocol requires it; None for http). Use
`ray.init(ignore_reinit_error=True)` so the driver's `num_gpus=0`
init can reuse the GPU cluster the pre-init started.
- config: set `comm_protocol="ray"` explicitly (new default is "http")
and pass `checkpoint_path=cfg.trainer.ckpt_path`. Auto-pick
`log_prob_engine="deepspeed"` only when `log_prob_gpus>0`.
Verified: package imports cleanly, `build_rl_config` produces a valid
`ArcticRLClientConfig` (pydantic) end-to-end from the public 4-GPU
GSM8K CLI overrides.
Co-authored-by: Cursor <cursoragent@cursor.com>
sfc-gh-kganesan
force-pushed
the
karthik/skyrl-arctic-rl-refactor-deltas
branch
from
June 17, 2026 20:33
e2ac216 to
1df089b
Compare
Adds a BIRD-SQL training recipe to the Arctic RL integration so that the same model (Qwen3-1.7B), same dataset (BIRD parquet at /data/snowflakesql/xyu/open-source-text2sql/), and the same reward function as the validated verl PR NovaSky-AI#6 run can be exercised through SkyRL. This gives us an apples-to-apples step-1 baseline to compare SkyRL's arctic_rl backend against the verl arctic_rl backend on identical inputs. New files (all under the integration; zero SkyRL public surface change): - arctic_rl/envs/bird.py: single-turn skyrl-gym env that delegates the reward to the existing `arctic_platform.rl.projects.txt2sql.bird_reward.compute_score` function -- the *exact* reward fn the verl PR NovaSky-AI#6 run uses (no reimplementation, no schema rewriting). Reads `extras["reward_model"].ground_truth` + `extras["extra_info"].db_path` straight out of the verl-format parquet, so no data conversion is required. - arctic_rl/envs/__init__.py: registers `bird` with skyrl_gym via the standard `register(id, entry_point)` API. Registration is triggered as a side-effect import from `arctic_rl/__init__.py`, so any recipe that imports the integration sees the env without modifying skyrl-gym upstream. - examples/run_bird_grpo_1.7b_8gpu.sh: launcher that mirrors verl PR NovaSky-AI#6 hyperparameters (8 GPUs colocated, ZeRO-3 + optimizer offload, BSZ=32 prompts x ROLL_N=16, prompt_len 32K, response_len 4K, lr=2e-6, no KL, entropy_coeff=0, single epoch). Toggles backends via `trainer.backend=arctic_rl` (the public discoverability hook -- no PYTHONPATH gymnastics for the recipe author). Verified (no GPU): - `arctic_rl` import -> `bird` appears in `skyrl_gym.envs.registration.registry` - BirdEnv built from a real verl-format parquet row; gold-SQL response scores 1.0, executable-wrong scores 0.1, junk scores 0.0 (matches bird_reward.compute_score semantics). - Launcher CLI overrides pass `ArcticSkyRLConfig.from_cli_overrides`, SkyRL `validate_cfg`, and `build_rl_config` -> emits a valid `arctic_platform.rl.ArcticRLClientConfig` with grad_accum=1 (matches verl PR NovaSky-AI#6's PPO_MINI_BSZ_PER_GPU == BSZ_PER_GPU). Wire-protocol caveat documented in the launcher header: absolute step-1 metric values will still drift from verl until the tied-embeddings weight-sync fix + verl-shape meta dict / verl_grpo loss / post-processors land in `arctic_platform.rl` (a separate PR there, per the strategy pivot). Step-1 invariants (clipfrac==0, ppo_kl==0 on epoch 1) and metric *shape* should already line up. Co-authored-by: Cursor <cursoragent@cursor.com>
`create_arctic_rl_client(rl_config)` now initializes Ray itself during pre-init (it owns a GPU cluster), so the driver-side `ray.init(..., runtime_env=..., ignore_reinit_error=True)` happens *second* and Ray silently drops the runtime_env on a no-op re-init. Workers then deserialize `ArcticSkyRLConfig` without `integrations/arctic-rl` on sys.path and crash with `ModuleNotFoundError: No module named 'arctic_rl'`. Move the runtime_env to task granularity (`skyrl_entrypoint.options( runtime_env=...).remote(...)`) so it applies regardless of whether we or arctic_platform.rl initialized the cluster. Repro before fix: BIRD smoke at /tmp/skyrl_bird_20260617T205957Z.log crashed at `ray.exceptions.RaySystemError: System error: No module named 'arctic_rl'` immediately after the ArcticRL jobs came up. Co-authored-by: Cursor <cursoragent@cursor.com>
The arctic_platform.rl server's payload contract -- shared with the verl adapter at verl/workers/remote_client/arctic_rl.py -- expects each batch to carry: prompts [B, P] prompt-only slice responses [B, A] response-only slice response_mask [B, A] response-only mask input_ids [B, S] full prompt+response (already there) attention_mask[B, S] already there position_ids [B, S] derived from attention_mask Producing those is a per-framework adapter job (not a server change): verl ships them already, SkyRL must too. The data is right there -- `convert_prompts_responses_to_batch_tensors` left-pads prompts to max_prompt_len and right-pads responses to max_response_len, so the split is a uniform `prompt_len = sequences.shape[1] - response_mask.shape[1]` across the batch. Verified via standalone tensor-shape unit test that the rebuilt `_to_batch` survives the server-side compute_packing_info_for_batch path that previously failed with `KeyError: 'prompts'`: prompts.shape == (B, P), response_mask.shape == (B, A), responses.shape == (B, A), input_ids/attention_mask/position_ids shape (B, S); attention_mask[:, P:].sum(dim=1) == per-row actual response lengths. Repro before: /tmp/skyrl_bird_20260617T211633Z.log step-1 generate succeeded (avg_raw_reward=0.4443, response_length=1179, in the same ballpark as verl PR NovaSky-AI#6 step-1's 0.5313/1192.9), then fwd_bwd died at arctic_platform/rl/processors/pipeline.py:236 KeyError('prompts'). Co-authored-by: Cursor <cursoragent@cursor.com>
Translate SkyRL's TrainingInputBatch to the arctic_platform.rl server's verl-shape contract entirely on the client side so step-1 metrics match verl PR NovaSky-AI#6 byte-for-byte without modifying arctic_platform.rl: - Per-sample repack of `[PAD_L | prompt | response]` (SkyRL native) to `[PAD_L*(P-p_i) | prompt | response | PAD_R*(R-r_i)]` (verl shape) so prompts/responses live in fixed [B, P] / [B, R] regions of a [B, P+R] sequence. Server's `compute_packing_info_for_batch` reads `prompts.shape[1]` to derive response_lens; SkyRL's variable-position layout was producing mixed prompt/pad tokens in the response window. - Synthesize `prompts`, `responses`, `position_ids` (cumsum-1) on the wire so the server's deepspeed_worker payload contract is satisfied. - Add full `meta` dict matching verl: `pad_token_id` (was the immediate KeyError blocker), `temperature`, `actor_config` + `policy_loss_config` (GRPO defaults: clip_ratio=0.2, entropy_coeff=0, use_kl_loss=False, loss_mode=vanilla), `dp_size`, `batch_num_tokens`, `global_batch_size`, `max_prompt_len`/`max_response_len`, `max_token_len_per_gpu`, `zorro_train_enable` mirrored from `arl.use_zorro`. - In `forward_backward`: compute old_log_probs server-side via `fwd_no_grad` (mirrors verl's pre-update_actor `compute_log_prob` call) so PPO ratio at step 1 is exp(0)=1 → ppo_kl=0, clipfrac=0. - Left-pad `response_mask`/`advantages`/`old_log_probs` to seq_len, set `loss_mask = response_mask`, drive `verl_grpo` loss with post=["apply_temperature", "compute_entropy_and_logprobs"]. - Wire ZoRRO toggle through honestly: `meta.zorro_train_enable` follows `arl.use_zorro` so the per-call meta and ds_worker_config don't desync. Payload is ZoRRO-compatible regardless (response-only tensors are already left-padded to seq_len). Validated by `/tmp/test_repack.py` (per-sample shape correctness on a 3-row toy batch with varying prompt/response lengths). All server-side changes are still deferred to a separate arctic_platform.rl PR.
After the wire-protocol bridge proved correct (3 consecutive steps emit
clean pg_loss / ppo_kl / clipfrac / grad_norm / reward metrics), two
follow-ups surfaced from the live logs:
- examples/run_bird_grpo_smoke.sh: drop-in shrink of the full BIRD
recipe (TRAIN_BSZ 32→4, N_SAMPLES 16→4, PROMPT_LEN 32768→4096,
RESPONSE_LEN 4096→512). Hits the same arctic_platform.rl wire path
as the full recipe but each training step takes ~15-20s instead of
~7-8min, so wire-shape / meta-dict / verl_grpo loss-config
iterations land in ~1.5min/iter instead of ~10min/iter. WandB logger
disabled by default to keep smoke runs out of the production project.
- trainer.optim_step: server reports grad_norm as a per-DP-rank list
([g]*world_size); SkyRL's reduce_metrics expects a scalar and warned
"Metrics for key grad_norm are not all numbers" each step. Flatten
+ pick first (every rank sees the same ZeRO-3 / DDP-reduced value).
Step-1 invariant note: ppo_kl ≈ -0.04 (vs verl's expected ~0) is bf16 /
flash-attn precision drift between the eval-mode fwd_no_grad (old log
probs) and train-mode fwd_bwd (new log probs), not a wire bug. Will
confirm against verl PR NovaSky-AI#6 baseline values before chasing further.
ArcticGenerator.generate() was scoring rollouts serially after vLLM returned. For BIRD this dominates: 512 samples * ~0.4-1.0s per SQLite query against the gold answer. Fan the post-generation phase out to a ProcessPoolExecutor (8 workers by default, override via ARCTIC_RL_SCORING_WORKERS). Process pool, not threads, so each worker keeps its own BIRD sqlite handle. Measured on BIRD / Qwen3-1.7B / 8xH100 colocated, averaged over the first 10 training steps: generate avg 273.4s -> 124.6s (-54%) step total avg 318.1s -> 170.1s (-47%) train_critic_and_policy 34.7s -> 35.0s (unchanged) Does not yet match verl xid2pl9f (timing_s/gen ~65s) -- verl's agent_loop overlaps generation and scoring; we still block generation before scoring starts. Closing that needs a per-prompt fan-out in ArcticGenerator.generate() instead of one batched arctic_client.generate call, tracked separately. Repro before: /data-fast/skyrl-runs/20260618T220520Z/skyrl_full.log Repro after: /data-fast/skyrl-runs/20260618T231552Z/skyrl_full.log Co-authored-by: Cursor <cursoragent@cursor.com>
Surface the full memory/perf surface that verl's xid2pl9f BIRD-1.7B
converged-reference run (launch_1.7b_newshape.sh in arctic-verl) sets,
so SkyRL+ArcticRL can match its memory headroom and step-time profile
on the same 8xH100 colocated topology. Defaults stay safe-off so
existing recipes are unaffected.
config.py: ArcticRLTrainerConfig gets 1:1 mappings to verl
- use_liger, attn_implementation (flash_attention_3),
enable_gradient_checkpointing, ulysses_sequence_parallel_size,
logits_optimization (memory|none),
logits_optimization_peak_mem_size_in_gib, cuda_ipc_weight_sync
- vllm_enforce_eager, vllm_enable_prefix_caching,
vllm_max_num_batched_tokens (40960 in verl)
- lr_warmup_ratio, optimizer_betas
- server_logs, startup_timeout
Plus gradient_accumulation_steps made ulysses_sp-aware so DeepSpeed's
batch-size assertion passes when SP > 1.
run_bird_grpo_1.7b_8gpu.sh: surface every knob above at the verl
default values, plus turn on use_zorro (server-side prompt dedup
and packing) which verl uses for BIRD.
Verified: step-1 PPO metrics match verl xid2pl9f within tolerance --
ppo_kl=0 exact, grad_norm=0.547 vs verl ~0.5, pg_loss=-0.0019 vs
verl ~0. No OOM on 8xH100 at colocated 0.5 vllm gpu_memory_utilization.
Wandb run: arctic_rl_bird_sql/9xfcr0sr.
Co-authored-by: Cursor <cursoragent@cursor.com>
…arity
Three trainer-side fixes landed together because they were surfaced
end-to-end while reproducing the verl xid2pl9f BIRD-1.7B run. Each is
small but independent.
1. Drop the colocate-mode sleep_inference / wake_training handshake
from train_critic_and_policy. The reference verl arctic_rl client
(arctic-verl/verl/workers/remote_client/arctic_rl.py) does not call
either anywhere -- colocated GPU memory is managed server-side and
fits via gradient-checkpointing + flash-attn + liger + a bounded
ppo_max_token_len_per_gpu, all set at engine-build time. The extra
handshake was both unnecessary and a divergence from the verl
contract.
2. Require update_epochs_per_batch == 1 and remove the per-epoch loop
wrapping fwd_bwd. _ArcticDispatch.forward_backward calls
_compute_old_log_probs every iteration, so >1 epoch would refresh
old_log_probs and collapse PPO ratio to exp(0)=1, defeating
clipping. The verl BIRD recipe uses 1 epoch so we assert rather
than silently misbehave. Multi-epoch support requires hoisting the
old-log-prob call into fwd_logprobs_values_reward (verl's
compute_log_prob placement); tracked separately.
3. Pin _build_meta to known-good engine-build values:
drop_position_ids=False, logits_optimization=none,
logits_optimization_peak_mem_size_in_gib=4,
logits_compute_in_fp32=False. A prior attempt to source these
dynamically per-call (logits_optimization=memory,
drop_position_ids=True) crashed at step 5 with a ZoRRO shape
mismatch:
RuntimeError: shape mismatch: value tensor of shape [8026]
cannot be broadcast to indexing result of shape [7042]
in arctic_platform/rl/zorro_train/seqlen_balancing.py on the first
packed micro-batch where prompts had been deduplicated. The
hardcoded values now match the engine-build config end-to-end.
Verified: 11 consecutive steps clean on the current run, including the
step-5 watershed where the prior _build_meta crashed. Step-1 PPO
metrics match verl reference (ppo_kl=0 exact, grad_norm 0.547 vs ~0.5).
Wandb run: arctic_rl_bird_sql/9xfcr0sr.
Co-authored-by: Cursor <cursoragent@cursor.com>
sfc-gh-kganesan
force-pushed
the
karthik/skyrl-arctic-rl-refactor-deltas
branch
from
June 19, 2026 00:33
d911e08 to
8afddb8
Compare
Qwen3ModelOncePatcher is built once at engine init with a fixed response_len (= cfg.generator.sampling_params.max_generate_length, wired through arctic_rl/config.py:412 -> deepspeed_worker.py:328). On every forward the patched causal-LM splits input_ids via `prompt_len = seq_len - response_len`. The client batch shaped with the dynamic per-batch max_r, so when no rollout in a batch hit the cap the patcher split shifted left -- the last (response_len - max_r) prompt tokens per sample leaked into the "response" region, the model returned sum(p_i_suffix + r_i) logprobs while pipeline.py:214 only had sum(r_i) attention slots, and the unpack crashed with "shape mismatch" at random steps. Fix: in `_repack_to_verl_shape`, when ZoRRO is enabled, pad max_r up to the patcher response_len. Extra positions get attention_mask=0 and pad_token_id so model forward and verl_grpo loss treat them as masked. The patcher seq_len - response_len now always equals the unpacker max_prompt_len. Repro before: BIRD/Qwen3-1.7B run 9xfcr0sr survived steps 1-23 then died at step 24 with "value tensor of shape [13252] cannot be broadcast to indexing result of shape [8556]"; surviving steps all happened to have at least one rollout hit the 4096 cap. Verified: synthetic repack test (4 samples, response_lens [50,30,100,20], patcher_response_len=4096) confirms patcher prompt_len == unpacker max_prompt_len, response-region attention sum still == sum(real response_lens). Co-authored-by: Cursor <cursoragent@cursor.com>
sfc-gh-kganesan
force-pushed
the
karthik/skyrl-arctic-rl-refactor-deltas
branch
from
June 19, 2026 05:03
c71c34b to
fb42f29
Compare
…e baseline Adds the Qwen3-32B 4-node BIRD recipes (ArcticRL + FSDP-native counterpart) used for the E2E comparison, and the client-side fixes uncovered while bringing them up. Validated by a stable 17-step ArcticRL 32B run on the 4-node Lustre cluster; the SkyRL FSDP-native counterpart is the comparison baseline. Client-side changes ------------------- * entrypoint.py / trainer.py (Option B): source `colocate`, `cuda_ipc`, `low_memory` from `cfg.trainer.arctic_rl` instead of `client.config`. `ArcticRLRayClient.reconnect_config()` strips the schema to a minimal serializable subset when shipping the client to Ray workers, so the worker-side `_ArcticDispatch` was seeing `colocate=False` regardless of the launcher flag — which silently disabled both sleep gates and OOM'd the DeepSpeed worker on step 2 of every 32B run. * trainer.py (`_ArcticInferenceEngineStub.sleep`): force `level=2` so vLLM's CuMemAllocator releases bf16 weight pages alongside KV cache. `level=1` keeps ~64 GiB resident and OOMs the DS worker on the first MLP allocation of the backward pass at 32B. * trainer.py (`save_weights_for_sampler`): keep the explicit `empty_training_cache()` / `wake_training()` / `wake_inference()` handshake around `client.sync_weights(cuda_ipc=…, low_memory=…)` — required at 32B because the platform-side orchestration leaves vLLM weights resident and the IPC clone OOMs without the manual drain. * config.py: `sampling_gpus = num_engines * tensor_parallel_size` (was just `num_engines`). At TP=4 / num_engines=8 this previously asked the orchestrator for 2 sampling replicas instead of 8, silently shrinking rollout parallelism 4x and tripping the multi-node FlashInfer workspace collision at init. * entrypoint.py: forward `WANDB_*` env vars to Ray workers (mirror the existing `ARCTIC_*` forwarding). Previously the worker actor was 401'ing against `api.wandb.ai` instead of `snowflake.wandb.io`. SkyRL-core changes (FSDP-native counterpart) -------------------------------------------- * fsdp_worker.py: `SKYRL_USE_LIGER` env-var opt-in for Liger fused linear-CE in `HFModelWrapper`. Needed for 32B FSDP-native runs (vocab 151936 + packed-seq up to 36864 + micro>=4 OOMs the LM head without it). Off by default — flag is opt-in. Launchers --------- * examples/run_bird_grpo_32b_32gpu.sh (new) — ArcticRL recipe * examples/run_bird_grpo_32b_32gpu_fsdp.sh (new) — FSDP-native * examples/run_bird_grpo_1.7b_32gpu.sh (new) — 1.7B 4-node * examples/run_bird_grpo_1.7b_8gpu.sh (touched) * examples/fsdp_bird_entry.py (new) — wrapper that registers the `bird` skyrl-gym env on the driver and forwards `arctic_rl` on PYTHONPATH for Ray workers, so the FSDP-native path can train on BIRD-SQL without depending on the ArcticRL backend. Deps ---- * uv.lock: pick up `arctic-training==0.8.0` (matches the `arctic-rl` extras in pyproject.toml). Co-authored-by: Cursor <cursoragent@cursor.com>
Self-review pass. Removes 120 lines of verbose explanatory comments and one stray `logger.info` diagnostic block in `train_critic_and_policy` that was a debug aid during Option-B bring-up. Behavioral net-zero: all sleep/wake/sync-weights/colocate semantics are unchanged. Only docstrings, in-line comments, and the temporary diagnostic log are trimmed. Co-authored-by: Cursor <cursoragent@cursor.com>
…5+S6) Replace the integration-specific ``trainer.backend`` dispatch + the generic sys.path injector (``_ensure_backend_importable``) in ``skyrl/train/entrypoints/main_base.py`` with a 6-line ``trainer.override_entrypoint`` peek-and-import, per Sumanth's PR #1 review comment: #1 (comment) Closes: - S5 (``_ensure_backend_importable`` is not true for all integrations) - S6 (keep ``main_base.main()`` simple; use ``trainer.override_entrypoint``) - C3 (revert the arctic-specific routing block) Core surface area shrinks: -38 / +12 lines, with no integration named in ``skyrl/`` and no sys.path manipulation. Config: - ``skyrl/train/config/config.py``: replace ``backend: str = "fsdp"`` with ``override_entrypoint: Optional[str] = None``. Dispatch: - ``skyrl/train/entrypoints/main_base.py``: peek ``trainer.override_entrypoint=`` from sys.argv before strict parse; if set, ``importlib.import_module(<path>).main()`` and return. Otherwise the standard FSDP path runs unchanged. Integration alignment with the new core API: - ``integrations/arctic_rl/`` (flattened from ``integrations/arctic-rl/arctic_rl/``; separate Option-A commit): docstrings + config narrative updated to reference ``trainer.override_entrypoint=integrations.arctic_rl.entrypoint``. Migration: launchers swap ``trainer.backend=arctic_rl`` → ``trainer.override_entrypoint=integrations.arctic_rl.entrypoint`` Co-authored-by: Cursor <cursoragent@cursor.com>
Make the two 32B launchers respect ``SKYRL_DIR`` and ``PYBIN`` from the environment so the same script can target alternate envs (e.g. ``PYBIN= /home/.../envs/skyrl_v2/bin/python``) without editing the file. Default values unchanged. Co-authored-by: Cursor <cursoragent@cursor.com>
…override
Removes the boilerplate ``trainer.arctic_rl={}`` from every recipe and makes
the ``arctic-rl`` extra self-contained, so any stock SkyRL recipe can opt
into Arctic by appending a single CLI override.
User experience after this commit (lands on top of 2587f0e + bb43337):
# one-time install
uv sync --extra fsdp --extra arctic-rl
# any recipe → Arctic backend, one flag
bash examples/<recipe>/run.sh \
trainer.override_entrypoint=integrations.arctic_rl.entrypoint
# optional knobs
... trainer.arctic_rl.colocate=true trainer.arctic_rl.zero_stage=3 ...
Changes:
- pyproject.toml: ``arctic-rl`` extra now pins ``arctic-platform`` and
``arctic-inference[vllm]`` (was a stale ``arctic_training`` PyPI
reference). Adds ``[tool.uv.sources]`` git entries pointing both at public
main (arctic-platform isn't on PyPI yet).
- integrations/arctic_rl/entrypoint.py: when ``cfg.trainer.arctic_rl`` is
``None`` (user didn't pass ``trainer.arctic_rl=`` overrides), default to
``ArcticRLTrainerConfig()`` so the single ``override_entrypoint`` flag is
enough.
- integrations/arctic_rl/examples/*.sh: drop ``trainer.arctic_rl={}`` (Hydra
auto-creates the parent dict from sub-key overrides; entrypoint fills the
default when missing entirely).
- integrations/arctic_rl/README.md: rewrite the Quick Start around the
single-flag any-recipe pattern; drop the stale references to the legacy
``arctic-skyrl`` / ``ArcticTraining-dss`` branches.
Co-authored-by: Cursor <cursoragent@cursor.com>
Removes hardcoded Snowflake-internal defaults so the launchers work
unmodified for any user:
- ``HF_HOME`` and ``VLLM_CACHE_ROOT`` default to ``$HOME/.cache/{huggingface,vllm}``
(was: ``/checkpoint/huggingface`` / ``/modeling-checkpoints/vllm``).
- ``DATA_DIR`` defaults to ``$HOME/data/bird`` (was hardcoded internal path).
- ``WANDB_API_KEY`` defaults to empty — user sets in their environment.
- ``WANDB_BASE_URL`` not set — falls back to public ``api.wandb.ai``.
- ``WANDB_PROJECT`` defaults to ``skyrl_arctic_rl``.
- ``ATTN_IMPL`` defaults to ``flash_attention_2`` (broadly available);
set ``ATTN_IMPL=flash_attention_3`` for the Hopper-only build.
- ``SKYRL_DIR`` defaults to script-relative repo root; ``PYBIN`` to ``python``.
Also drops stale README references to the legacy companion repos and
strips remaining internal-name comments in the 32B launcher.
Co-authored-by: Cursor <cursoragent@cursor.com>
sfc-gh-kganesan
force-pushed
the
karthik/skyrl-arctic-rl-refactor-deltas
branch
from
June 25, 2026 01:20
2c1fe23 to
ee13bb2
Compare
…ficient) The ``arctic-rl`` extra already provides everything an Arctic run needs: - ``skyrl[skyrl-train]`` (SkyRL training core) - ``arctic-platform`` (DeepSpeed training workers) - ``arctic-inference[vllm]`` (vLLM server; brings torch + vllm transitively) ``--extra fsdp`` only matters if you also want to run the SkyRL native FSDP baseline side-by-side (the two extras pin different vLLM versions, so arctic-rl alone is the cleaner install for pure Arctic runs). Co-authored-by: Cursor <cursoragent@cursor.com>
The launcher hardcoded ``trainer.logger=wandb``, which forced ``WANDB_API_KEY`` to be set even for smoke tests / console-only runs (SkyRL ``validate_generator_cfg`` asserts the key when logger is wandb). Thread ``LOGGER`` env var through (default ``wandb`` for parity with prior behavior; ``LOGGER=console`` for no-creds runs). Co-authored-by: Cursor <cursoragent@cursor.com>
Add explicit Prerequisites / Clone+install / Ray bootstrap / Data+model prep / Run steps so a fresh user can go from empty directory to a running GRPO loop without leaving the README. Also fixes a stale claim that the launchers default to ``LOGGER=console`` (they actually default to ``wandb``; ``LOGGER=console`` is the no-creds opt-out). Co-authored-by: Cursor <cursoragent@cursor.com>
Removes manual prep steps a fresh user previously needed:
- All launchers: ``HF_HUB_OFFLINE`` / ``TRANSFORMERS_OFFLINE`` now default
to ``0`` (auto-download). Set to ``1`` on isolated clusters where the
model is pre-staged in ``HF_HOME``.
- 32B + 32B FSDP launchers: drop the manual ``$HF_HOME/hub/.../refs/main``
snapshot dance. Pass ``Qwen/Qwen3-32B`` as the HF id; transformers/vLLM
auto-download to ``HF_HOME`` on first use. Multi-node users with a
shared pre-staged cache can still ``MODEL=<absolute path>`` to skip
the hub lookup.
- GSM8K launcher: auto-run ``examples/train/gsm8k/gsm8k_dataset.py`` when
``$DATA_DIR/{train,validation}.parquet`` doesn't exist.
- BIRD-SQL launcher: clear error pointing to ``$DATA_DIR`` when parquets
are missing (no public prep script — BYO data).
- 32B launchers: ``CHECKPOINT_DIR`` defaults to ``$HOME/skyrl-runs/ckpts/<run>``
(was a hardcoded ``/data/skyrl-runs/...`` cluster path).
- README: drop the manual prep section; document the auto-prep behavior.
Net result: ``bash integrations/arctic_rl/examples/run_gsm8k_grpo_4gpu.sh``
works on a fresh machine after ``uv sync --extra arctic-rl`` + ``ray start``
with no other prep.
Co-authored-by: Cursor <cursoragent@cursor.com>
The launchers set ``trainer.arctic_rl.use_liger=true`` (fused linear-CE + MLP/RMSNorm kernels, critical for 32B memory). arctic-platform's DeepSpeedWorker imports ``liger_kernel.transformers.monkey_patch`` at init time when use_liger is on, so a fresh ``uv sync --extra arctic-rl`` without liger-kernel crashed at worker initialize with ``ModuleNotFoundError: No module named 'liger_kernel'``. Pinning ``liger-kernel`` (PyPI, pip-installable) directly in the ``arctic-rl`` extra so a fresh install has everything the recipes need. Co-authored-by: Cursor <cursoragent@cursor.com>
…ic-platform branch The BIRD env's reward fn (`arctic_platform.rl.projects.txt2sql.bird_reward`) only exists on Arctic-Platform's private recipe/rl-correctness branch — it isn't shipped from public main. When users install arctic-platform from public main, every Ray actor silently fails the import and `ArcticGenerator` returns score=0 for every sample; the run looks healthy but never converges. Vendor `bird_reward.py` (~10 KB, stdlib-only) into the integration so it's self-contained and matches the validated verl PR NovaSky-AI#6 reward function bit-for-bit. Keep the upstream copy as the source of truth — re-sync if it changes. Co-authored-by: Cursor <cursoragent@cursor.com>
sfc-gh-kganesan
force-pushed
the
karthik/skyrl-arctic-rl-refactor-deltas
branch
from
June 25, 2026 14:46
8d6a1ba to
356cb30
Compare
This makes the BIRD-SQL pipeline runnable on the public main branches of
arctic-inference and arctic-platform with no private-branch dependencies:
- integrations/arctic_rl/envs/preprocess_bird.py (vendored verbatim from
arctic_platform.rl.projects.txt2sql.preprocess_bird; that module lives
on a private recipes branch and is not on public arctic-platform main).
- README: walkthrough for running the preprocessor + FlashAttention-3
install from the official PyTorch wheel index (needed for the 2x
speedup on Hopper).
Together with the earlier vendoring of bird_reward.py, the SkyRL client
now ships everything needed to reproduce the BIRD GRPO 32B benchmark.
Re-sync these vendored files if their upstream copies change.
Co-authored-by: Cursor <cursoragent@cursor.com>
vLLM's ``AsyncEngineArgs.__post_init__`` converts nested overrides only when ``isinstance(value, dict)`` is true, so an ``OmegaConf.DictConfig`` (what Hydra hands us) gets silently dropped — including the ``compilation_config`` / ``speculative_config`` / ``forest_cascade_attn_configs`` needed for the 2x Arctic-RL speedup. The previous ``dict(...)`` cast was shallow, leaving nested ``DictConfig`` values intact. Round-trip through ``OmegaConf.create() -> to_container(resolve=True)`` deep-coerces the whole tree to plain Python and is idempotent for both ``DictConfig`` and plain ``dict`` inputs. Matches the ``OmegaConf.to_container`` idiom used in arctic-verl's ``workers/remote_client/arctic_rl.py`` (tunji/remote_backend). Also drop the unused 1.7B smoke launcher scripts (kept locally during debugging; not part of the public recipe set) and trim the README's post-install ``transformers<5`` fixup note to just the one-liner users need. Co-authored-by: Cursor <cursoragent@cursor.com>
Re-pin optimization_level: 1 in the 32B launcher and ship a sibling 8B launcher. Today's TP=4 experiment confirmed that the OmegaConf round-trip in integrations/arctic_rl/config.py is *not* sufficient on its own: vLLM's engine init still resolves cudagraph_mode=FULL_AND_PIECEWISE and fuse_allreduce_rms=True even with an explicit compilation_config override on the CLI — i.e., the nested override is being dropped somewhere between ArcticRLClientConfig and AsyncEngineArgs. Until that plumbing is traced end-to-end, optimization_level=1 (which hard-codes fuse_allreduce_rms=false inside vLLM) is the reliable speedup config and reproduces the Jun 24 (skyrl_v1) 2x baseline on Hopper TP>1. Also add TORCHINDUCTOR_FORCE_DISABLE_CACHES=1 so a prior compiled graph that baked in flashinfer_trtllm_fused_allreduce_norm can't be reused across config flips (VLLM_DISABLE_COMPILE_CACHE only covers vLLM's own cache, not inductor's). The 8B launcher mirrors the 32B recipe (TP=4, FCA, CUDA-IPC weight sync, ZoRRo, Liger) for fast iteration on the same TP>1 code path; spec-dec is off by default since the published 32B 3-head checkpoint is architecturally tied to Qwen3-32B. Co-authored-by: Cursor <cursoragent@cursor.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
2 participants
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.
What this PR does
Brings the SkyRL Arctic RL integration to parity with the validated Snowflake-AI-Research/verl PR #6 (1.7B Qwen3 + BIRD, val/exec 0.913 at step 269) and addresses CharlieFRuan's footprint-reduction comments from the parent PR #1. SumanthRH's core-side comments (Jun 2 / Jun 3 / Jun 15) are tracked but not yet applied — see status table below.
Latest update (commit
c54d400) extends the integration to Qwen3-32B on 4 nodes and adds a SkyRL FSDP-native counterpart for the head-to-head comparison — Arctic-RL is 2.09× faster end-to-end per step at 32B. Numbers below.Commit map
fad7d70arctic_training.arctic_rl→arctic_platform.rl; moveArcticRLTrainerConfigtointegrations/arctic-rl/arctic_rl/config.py; drop_propagate_arctic_env_varsfrom core; remove arctic-specific routing block frommain_base.py; rewrite trainer / generator / entrypoint / config against the new async cliente2ac216checkpoint_pathtoArcticRLClientConfig;ray.init(ignore_reinit_error=True);server_stateactor handle plumbed throughskyrl_entrypointforcomm_protocol=rayreconnect;gradient_accumulation_stepsformula matched to verl (one DS step per PPO mini-batch); explicittrain_batch_sizeinds_config; fullvllm_configset (max_model_len,max_num_seqs,enforce_eager,enable_chunked_prefill)eb9fcdaProcessPoolExecutor— generation phase dropped from ~296s → ~157s on BIRD/Qwen3-1.7B947f8e2use_liger,attn_implementation=flash_attention_3,enable_gradient_checkpointing,ulysses_sequence_parallel_size,logits_optimization=memory,cuda_ipc_weight_sync,vllm_enforce_eager=false,vllm_enable_prefix_caching,vllm_max_num_batched_tokens=40960,lr_warmup_ratio,optimizer_betas,server_logs,startup_timeout) onArcticRLTrainerConfig8afddb8update_epochs_per_batch=1(otherwise per-epochold_log_probsrecompute collapses PPO ratio to 1); drop explicitsleep_inference / wake_traininghandshake (server manages its own memory); pin per-callmetato known-good values (drop_position_ids=False,logits_optimization="none",logits_optimization_peak_mem_size_in_gib=4,logits_compute_in_fp32=False) to avoid ZoRRO shape-mismatch crash from per-call meta overridefb42f29_repack_to_verl_shape, pad batch responses up tocfg.generator.sampling_params.max_generate_lengthwhen ZoRRO is on.Qwen3ModelOncePatcherderivesprompt_len = seq_len - response_lenat every forward using the fixedresponse_lenbaked in at engine build; the client was using a dynamic per-batchmax_r, so when no rollout in a batch hit the cap the patcher's split shifted left and pipeline.py:214 crashed withshape mismatch. Padding the response region to the fixedresponse_lenkeeps the patcher's split aligned with the server-side unpacker'smeta["max_prompt_len"]c54d400colocate/cuda_ipc/low_memoryfromcfg.trainer.arctic_rl(worker-sideArcticRLRayClient.reconnect_config()stripsclient.configto a minimal subset, so trusting it silently disables both sleep gates and OOMs step 2);_ArcticInferenceEngineStub.sleepforceslevel=2so vLLM releases bf16 weight pages with KV (level=1 keeps ~64 GiB resident and OOMs the DS worker's first MLP allocation at 32B); keep the explicitempty_training_cache/wake_training/wake_inferencehandshake aroundclient.sync_weights(cuda_ipc=…, low_memory=…)at 32B (platform orchestration leaves vLLM weights resident and the IPC clone OOMs without it);sampling_gpus = num_engines * tensor_parallel_size(was justnum_engines, silently shrinking rollout replicas 4x at TP=4); forwardWANDB_*env vars to Ray workers;SKYRL_USE_LIGERenv-var opt-in infsdp_worker.pyto enable Liger fused linear-CE inHFModelWrapperfor the FSDP-native baseline57431ddlogger.infodiagnostic block intrain_critic_and_policyleft over from the Option-B bring-up. Behavioral net-zero; only docstrings, in-line comments, and the temporary diagnostic log are changedReview comments status (parent PR #1)
pyproject.toml: minimalarctic-rlextras, reverttool.setuptools.packages.findto upstreamarctic-rl = ["skyrl[skyrl-train]", "arctic_training"];include = ["skyrl*"]restoredArcticRLTrainerConfigout ofskyrl/train/config/config.pyinto the integrationintegrations/arctic-rl/arctic_rl/config.py:29+ArcticSkyRLConfig = make_config(trainer_cls=ArcticTrainerConfig)at line 223 (mirrorsHarborSkyRLConfig)skyrl/train/entrypoints/main_base.pymain_base.py_propagate_arctic_env_varsfromskyrl/train/utils/utils.pyARCTIC_*forwarding now inline inarctic_rl/entrypoint.py:148-150uv run --extra arctic-rl -m arctic_rl.entrypointrun_gsm8k_grpo_4gpu.sh/run_bird_grpo_1.7b_8gpu.shstill usespython -m skyrl.train.entrypoints.main_base trainer.backend=arctic_rl .... Will flip once Sumanth's NovaSky-AI#11 lands andtrainer.override_entrypointexists.megatronthe same asfsdpmain_base.py:512is stillif backend != "fsdp", sotrainer.backend=megatronwould mis-trigger the integration import path. Will land with NovaSky-AI#11 below._ensure_backend_importablePYTHONPATH magicmain_base.py:475, called at line 516. Will land with NovaSky-AI#11 below.import arctic_rlrequires running fromintegrations/arctic-rl/trainer.backendfield rather than reusingTrainerConfig.strategy?backend: str = "fsdp"still inskyrl/train/config/config.py:637. Resolved by deletion when NovaSky-AI#11 lands._ensure_backend_importableassumption isn't true for all integrationstrainer.override_entrypointfieldoverride_entrypointdoes not exist in code; the oldtrainer.backendmachinery is still inmain_base.py. Will land as a follow-up commit on this PR (see "Follow-ups" below).Follow-ups (planned for this PR)
trainer.override_entrypoint: Optional[str] = NonetoTrainerConfig(resolves Sandbox timeout after several minutes into training, and "422 Unprocessable Entity" NovaSky-AI/SkyRL#9, Fix links for RL example scripts NovaSky-AI/SkyRL#11)_ensure_backend_importable+ thetrainer.backend != "fsdp"dispatch inmain_base.pywith a 5-lineoverride_entrypointdispatch (resolves Update readme NovaSky-AI/SkyRL#6, Which commit of verl is this repo based on? NovaSky-AI/SkyRL#7, NameError: name 'full_response_texts' is not defined NovaSky-AI/SkyRL#8, Add more concrete instruction on how to download dataset NovaSky-AI/SkyRL#10, Fix links for RL example scripts NovaSky-AI/SkyRL#11 — Sumanth's exact suggested shape)backend: str = "fsdp"field fromTrainerConfig(Sandbox timeout after several minutes into training, and "422 Unprocessable Entity" NovaSky-AI/SkyRL#9)PYTHONPATH=integrations/arctic-rl trainer.override_entrypoint=arctic_rl.entrypoint([arctic_rl] repoint to arctic_platform.rl + apply verl PR #6 correctness fixes; address reviewer comments #5)Invocation (today)
After the override_entrypoint follow-up lands:
PYTHONPATH=integrations/arctic-rl python -m skyrl.train.entrypoints.main_base \ trainer.override_entrypoint=arctic_rl.entrypoint <flags>Correctness fixes ported from verl PR NovaSky-AI#6
Failure mode without them was entropy collapse around step 8-16.
training_horizon = total_training_steps * update_epochs_per_batch * num_minibatches. The verl bug was passingepochs=1which collapsed warmup to 0 steps.trainer.policy.optimizer_config.max_grad_norminto both DS optimizer + engine.arctic_platform.rlshape:torch_autocast,communication_data_type,data_types,zero_optimization,log_level.zorro_train_enablenaming + exposelogits_compute_in_fp32,logits_compute_from_fp32_inputs,logits_optimization*._shifted(roll(-1)) convention documented inarctic_platform.rl.ray_client.fwd_bwd.Validation — 1.7B BIRD
arctic_platform.rlbackend, different driver. 269 steps, val/exec 0.913. (wandb xid2pl9f)eval/bird/avg_score = 0.5828at step 200 — already +0.038 above the VERL reference's all-time peak (0.5445 @ step 60). Zero crashes across 200+ ZoRRO steps after thec71c34b/fb42f29patcher-response-len fix. PPO invariants pristine throughout (ppo_kl=0,pg_clipfrac=0). (wandb st3ue30x)eval/bird/avg_scoreval-core/bird/reward/mean@132B BIRD-SQL E2E: ArcticRL vs SkyRL-Native FSDP
Hardware: 4 × 8 × H200 (32 GPUs, 140 GiB HBM each).
Model: Qwen3-32B. Task: BIRD-SQL GRPO,
train_batch=128prompts ×n_samples=16= 2048 trajectories/step.Sequence:
max_prompt=32768,max_response=4096, packed-seq (sample_packing=true).Same dataset, batch sizes, DP=32, vLLM TP=4, 8 engines,
gpu_mem_util=0.5for both runs.W&B project:
new_verl_bird_sqlonhttps://snowflake.wandb.io/karthik-ganesan/.E2E per-step (steady-state):
Where the deltas come from:
forest_cascade_attn_configs), speculative decoding (3-head Arctic draft,num_speculative_tokens=3), and piecewise CUDA graphs (cudagraph_mode: PIECEWISE).n_samples=16, grad_accum=4), Liger fused linear-CE (critical for Qwen3-32B with vocab=151,936 + packed-seq up to 36,864 — naive logits would be ~10 GiB per sample), Liger fused MLP/RMSNorm, flash_attention_3,logits_optimization=memory,offload_optimizer=true(Adam state on CPU). SkyRL native FSDP is forced tomicro_train_batch_size_per_gpu=2(grad_accum=32) — Liger is available viaSKYRL_USE_LIGER=1in this PR but the Qwen3 Liger kernel hits a Triton illegal-memory-access on packed-seq inputs (cu_seqlensvariable +attention_mask=None+ explicitposition_ids). Without Liger, micro=16 OOMs the LM-head logits (~174 GiB). 16× more grad-accum steps → ~3.8× slower per-step train.gather_cuda_ipc_handles_low_memory()(parameter-at-a-time on the training side, IPC handles passed to vLLM, vLLM reconstructs from shared memory; includes vLLMwake_inference(level=2)reallocation overhead). SkyRL native uses NCCL collectivebroadcastfrom DP rank-0 of FSDP into the vLLM ranks via the colocated PG — direct GPU-to-GPU. The FSDP win on sync_weights is real and material, but the trainer + generate wins dominate the total step.Reproduce:
Open items for the FSDP path (tracked, not blocking this PR):
lce_forward. Either patch Liger or disable sample_packing on the FSDP path; both have tradeoffs.PYTORCH_CUDA_ALLOC_CONF=expandable_segments:Trueonly on the FSDP worker process (not vLLM, which is incompatible) may help.n_samplesreplication into a single micro-batch would unlock similar grad-accum efficiency.File structure