[arctic_rl] backport rl-correctness payload-encoding deltas onto arctic_rl_share_v0.7.1

[sfc-gh-kganesan]

What this PR does

Brings the new RemoteBackend abstraction shape to convergence parity with recipe/rl-correctness on 1.7B GRPO, by porting the correctness / numerical / scheduling fixes onto our Arctic adapter and per-backend yaml — without touching the core abstraction (verl/remote_backend/{base,trainer,worker_utils}.py).

When verl-project#6422 merges and we rebase our internal recipe/rl-correctness fork on top of it, the adapter shape, payload encoding, DS engine config, and training-time guards already match.

Architecture (recap)

• Core abstraction (untouched): verl/remote_backend/{base.py, trainer.py, worker_utils.py}. Defines RemoteBackend, RemoteBackendRegistry, RemoteBackendTrainer.
• Per-backend worker (extension point, also untouched): verl/remote_backend/workers/arctic_rl/worker.py. Driver-side wiring for the Arctic backend.
• Adapter (this PR): verl/workers/remote_client/arctic_rl.py. Implements RemoteBackend for arctic_platform.rl. All wire-format, payload encoding, DS engine config and per-call meta lives here.
• Per-backend yaml (this PR): verl/trainer/config/remote_backend/arctic.yaml. Loaded via Hydra at cfg.remote_backend when trainer.remote_backend=arctic.

Other backends drop in by registering a new class + a new yaml file. Nothing in verl/remote_backend/ ever needs to change.

Commits (6 total, all on top of 2174486b)

Five of the six commits are direct ports from recipe/rl-correctness. Authorship and subject lines preserved exactly so this is one-click for Tunji to audit.

#	commit	author	subject	source (recipe/rl-correctness)
1	fbab32f3	Karthik Ganesan	arctic_rl: adapter for arctic_platform.rl on the new RemoteBackend shape	— (new wiring)
2	7eb21212	Xiaodong Yu	Fix zorro log-prob off-by-one in no_padding_2_padding (verl-project#25)	d597a2fd
3	bc510b12	Mert Hidayetoglu	RL correctness: grad clipping, LR schedule horizon, fp32 grads, zorro load-balancer flag	d8d9ea6f
4	1f41c88b	Olatunji Ruwase	ds_config format + ds_config passthru	303eb275 + 8356f215
5	0126bb79	Stas Bekman	enable_gradient_checkpointing is configurable now	85cdbedd
6	e6899904	Stas Bekman	add logits_compute_from_fp32_inputs + logits_compute_in_fp32 configs	da16e078

Each ported commit's body cites its recipe/rl-correctness source SHA and calls out any newshape-specific adaptation (e.g. config-path tweaks for remote_backend.zorro_train.enable vs upstream's arctic_rl.use_zorro, or applying the same change to our adapter file instead of upstream's arctic_rl_client.py). Tunji: please read the commit bodies — they're meant to make the audit easy.

Convergence (1.7B BIRD GRPO, 8×H100, ZeRO-3, ZoRRo)

Full run on this branch, hparams identical to the validated PR verl-project#37 launcher (see launch_1.7b_newshape.sh).

Validation (test_freq=10) — beats PR verl-project#37 baseline at every checkpoint

step	val reward	val format	val exec	vs PR verl-project#37
10	0.512	0.975	0.774	reward +6.9%, exec +7.4%
20	0.507	0.976	0.787	reward +5.0%, exec +6.8%
30	0.513	0.977	0.807	reward +4.8%, exec +6.6%
40	0.518	0.975	0.819	—
50	0.535	0.967	0.816	+9.4% reward / +7.8% exec vs PR verl-project#37@30

Training (rolling 20-step windows)

window	entropy (mean / range)	train reward (mean / range)
steps 3-20	0.175 / [0.166, 0.191]	0.507 / [0.374, 0.615]
steps 21-40	0.188 / [0.171, 0.199]	0.508 / [0.337, 0.662]
steps 41-58	0.199 / [0.180, 0.214]	0.560 / [0.466, 0.660]

Entropy stable in the 0.17-0.21 band (target 0.17-0.20). For comparison, prior newshape runs (this branch before the ported commits) had entropy at 0.4-0.6 by step 16 and never recovered.

Root causes that were killing convergence

Two changes from the upstream port are load-bearing:

1. Log-prob off-by-one in no_padding_2_padding (zorro path) — commit 7eb21212 / upstream d597a2fd. Adapter requested zorro/response-aligned log-probs but the trainer applied the legacy "predict-next" -1 shift before computing PPO loss → systematic one-token misalignment → biased policy gradient. Fixed by _model_output_shift() returning shift=0 when remote_backend.zorro_train.enable=True.
2. LR-scheduler horizon collapsed to 1 — part of commit bc510b12 / upstream d8d9ea6f. Adapter passed training_horizon=trainer.total_epochs=1, so a 5% warmup ratio evaluated to 0.05 * 1 = 0 warmup steps → LR jumped from 0 to peak in a single step, blew out the policy in the first ~10 updates. Fixed by computing training_horizon = global_steps * ppo_epochs * num_minibatches (matches recipe/rl-correctness).

Key surfaces

verl/workers/remote_client/arctic_rl.py

• Implements RemoteBackend for arctic_platform.rl (was arctic_training.rl).
• _create_ds_config(n_gpus, training=True): passes the whole remote_backend.deepspeed block through (OmegaConf.to_container + deepcopy) and merges in train_micro_batch_size_per_gpu / train_batch_size / gradient_accumulation_steps / sequence_parallel_size. Drops data_types when training=False.
• _create_ds_worker_config(): wires enable_gradient_checkpointing, logits_compute_from_fp32_inputs, logits_compute_in_fp32, logits_optimization, logits_optimization_peak_mem_size_in_gib. max_token_len = actor.ppo_max_token_len_per_gpu (matches upstream).
• _initialize_client: computes training_horizon = global_steps * ppo_epochs * num_minibatches, forwards optim.clip_grad to the DS engine, threads lr_scheduler config (type / warmup_ratio / min_lr_ratio).
• compute_log_prob.meta + update_actor.meta: keys byte-identical to recipe/rl-correctness (arctic_remote_adapter / arctic_workers).

verl/trainer/ppo/ray_trainer.py

• _model_output_shift() returns 0 for the zorro path (remote_backend.zorro_train.enable=True), -1 otherwise. Threaded into no_padding_2_padding in _compute_values and _compute_old_log_prob.

verl/workers/utils/padding.py

• no_padding_2_padding(...) now accepts shift.

verl/trainer/config/remote_backend/arctic.yaml

• zorro_train.{enable, load_balancer, max_rollouts}
• drop_position_ids (flat)
• logits_optimization{,_peak_mem_size_in_gib}, logits_compute_in_fp32, logits_compute_from_fp32_inputs
• cuda_ipc_weight_sync, low_memory_weight_sync
• deepspeed.{torch_autocast, communication_data_type, data_types, zero_optimization, log_level} (passthrough block)

Dependency notes

• arctic_platform, arctic_inference, deepspeed are pinned to specific commits / PyPI versions and installed non-editable from local clones (no reliance on anyone's working tree). Process codified in skill pin-vendor-deps.
• mixed_precision.fp32_gradients (recipe/rl-correctness uses fp32 grad accum) requires deepspeed sfc-gh-truwase/fp32_grads (commit 9acf2392). Defaulting to grad_accum_dtype: bf16 on PyPI DS 0.19.2 since convergence holds without it (val/exec 0.816 at step 50). Easy flip if needed: switch DS to that branch, change data_types.grad_accum_dtype: fp32 in arctic.yaml.

Out of scope (deferrable follow-ups)

• mixed_precision.fp32_gradients (requires DS branch swap; not on the convergence path).
• The PPO mini-batch loop in update_actor (PR#37 does it in the adapter; the target shape lets the verl-side worker drive it). No semantic change needed for current configs (ppo_epochs=1, train_batch_size == ppo_mini_batch_size).
• use_fca, spec_model, full_determinism, seed — orthogonal features, not required for parity.

Related

• snowflake-eng/arctic-verl#37 — predecessor PR (Arctic adapter on recipe/rl-correctness, older abstraction shape; converging baseline).
• recipe/rl-correctness — source branch for ported commits (d597a2fd, d8d9ea6f, 303eb275, 8356f215, 85cdbedd, da16e078).
• verl-project/verl#6422 — upstream abstraction PR this branch is the head of.

[sfc-gh-kganesan]

E2E smoke result (deferred from PR description) — PASS

Stopped the 1.7B convergence run on PR verl-project#37 branch and ran a 4-step Qwen3-0.6B BIRD GRPO smoke against this PR (commit 96249486).

Setup: 1 node × 8 H200, ZoRRo train enabled, colocate=True, cuda_ipc_weight_sync=True, gpu_memory_utilization=0.6, zero.stage=3, optimizer.offload=cpu, param.offload=none, logits_optimization=memory, hparams matched validated PR verl-project#37 0.6B smoke.

Issue caught & fixed by the smoke (committed as 96249486):

TypeError: ArcticRLRayClient.generate() got an unexpected keyword argument 'routing_key'

arctic_platform.rl.ray_client.ArcticRLRayClient.generate has signature (prompts, sampling_params=None) and handles routing internally — adapter must not forward routing_key. Wrapper signature unchanged; only the inner call was modified. Vendor-only change (verl/workers/remote_client/arctic_rl.py); verl/remote_backend/{base,trainer,worker_utils,workers/arctic_rl/worker}.py are diff-clean vs public/arctic_rl_share_v0.7.1.

Training (4 steps):

step	score	actor/loss	grad_norm	lr	step time
1	0.290	0.153	0.373	0	120.2 s
2	0.293	0.051	0.230	1e-6	106.1 s
3	0.276	0.190	0.415	2e-6	96.1 s
4	0.306	0.128	0.328	2e-6	101.4 s

Final validation @ step 4 (BIRD full val set):

• val-core/bird/reward/mean@1 = 0.2941
• val-aux/bird/execution_success/mean@1 = 0.5222
• val-aux/bird/format_correct/mean@1 = 0.9426

Numerically matches PR verl-project#37 reference smoke (PR verl-project#37: score 0.293 → 0.308, val exec 0.510; this PR: 0.290 → 0.306, val exec 0.522) — within noise, exec_success marginally better.

Clean shutdown, GPUs released, no other errors. Inner library confirmed live as arctic_platform.rl (no arctic_training traces in process / sys.modules).

[sfc-gh-kganesan]

Self-review fix-up — 5ab85d58

Did a self-review pass before requesting review and surfaced three real findings; pushed a fix-up commit (5ab85d58). Scope remains vendor-only (still exactly 2 files vs public/arctic_rl_share_v0.7.1; verl/remote_backend/{base,trainer,worker_utils,workers/arctic_rl/worker}.py are still diff-clean).

Finding 1 (important — fixed): logits knobs were silently no-op'd

The server's read path for logits_optimization, logits_optimization_peak_mem_size_in_gib, and logits_compute_in_fp32 is arctic_platform/rl/deepspeed_worker.py:319-322, which reads from ds_worker_config (server config), NOT per-call request meta. Before this fix-up the adapter forwarded them only in the per-call meta dicts, so the server fell back to its hard-coded defaults ("none", 4, False).

The shipped yaml values happen to match those defaults, so this was masked at the validated config (and explains why the 0.6B 4-step smoke and the 1.7B step-10 val both passed even with the bug present). But it would have silently no-op'd the moment someone tried logits_optimization=memory for OOM mitigation — exactly the kind of debugging round-trip this PR is trying to prevent.

Fix: thread the 3 keys into _create_ds_worker_config() inside the if self.zorro_train_enable block, alongside the existing zorro-only fields (response_len, max_token_len, rollout_n, temperature). Per-call meta forwarding kept for forward-compat with a future server release that might add per-batch overrides.

Finding 2 (important — documented): zorro_train_load_balancer has no server consumer

grep -r zorro_train_load_balancer arctic_platform/rl/ → zero hits. The current server release has no consumer for this key. We keep forwarding it (matches the PR verl-project#37 wire format so the server-side handler — likely tied to ArcticInference-internal PR verl-project#92's strict group-balanced routing — can land independently without an adapter change), but added a TODO comment so the next reader knows it's intentionally a no-op today, not a wiring bug.

Nit 3 (documented): sibling logits_compute_from_fp32_inputs deliberately not exposed

Server has both logits_compute_in_fp32 (output-side upcast, exposed) and logits_compute_from_fp32_inputs (input-side upcast, not exposed). Matches PR verl-project#37; added a yaml + __init__ comment so a future reviewer doesn't think the missing sibling is an oversight.

Validation

• Import smoke: OK
• Yaml loads via OmegaConf, all 5 new fields parse correctly
• _create_ds_worker_config() source now contains all 3 logits keys (verified)
• No new lints
• Live 1.7B convergence run (wandb xqj9pqbt, currently at step ~12/269 with val@step-10 reward=0.497 / exec=0.768 / format=0.969 — ahead of the killed PR [distro] refactor: cleanup dependencies in setup script verl-project/verl#37 baseline) is pinned to commit 96249486 via editable import and is NOT affected by this push; will pick up the fix on its next launch.

PR diff stats after fix-up: 2 files, +138 / −14 vs public/arctic_rl_share_v0.7.1.

[sfc-gh-sbekman]

Question: Is this how the code now in the original? That super narrow comment column and a huge slop is so hard to read.

Do we need to fix it in the origin first? or have you adapted this and it's not matching the source? Asking so that we don't do the work twice.

Let's try to keep the slop away, this over-explaining makes the code very difficult to read. I found I need to instruct claude to rewrite its comments to 119 chars (not sure what verl's width is) and make them terse and only when necessary, instead using code that is self-explanatory.

[sfc-gh-kganesan]

Thanks @sfc-gh-sbekman — fair, the slop was all mine, not inherited. Pushed 59d61579.

Answer to your question: the upstream arctic_rl_share_v0.7.1 adapter is already terse (~10% comment density, mostly 1-line docstrings). The over-explained blocks were comments I added on top of upstream while wiring in the arctic_platform.rl / ZoRRO config knobs. Nothing to fix at the source — just here.

What I changed in this commit (only touches lines I added; pre-existing upstream comments untouched):

• verl/workers/remote_client/arctic_rl.py: 5 multi-paragraph blocks → 1-2 lines each. Dropped PR refs and forward-compat history (belongs in commit messages). Doctring on _prepare_padded_arctic_batch_dict collapsed to 3 lines.
• verl/trainer/config/remote_backend/arctic.yaml: each knob now has a single-line comment (only logits_optimization keeps a 3-line enum because the menu is genuinely useful at the config site).

Width: verl's pyproject.toml sets line-length = 120. Confirmed every added/edited line is now ≤ 119 chars.

Diff: +24 / -109, no runtime code touched. Convergence run unaffected (currently at step 30/269, validating).

Will keep comments terse going forward — appreciate the explicit "rewrite to 120 / only when necessary" guidance, going to carry that into the rest of the work too.

[sfc-gh-kganesan]

Quick follow-up self-audit on b74f8c1f — caught three residuals against your guidance:

1. PR ref still in source (ArcticInference-internal PR #92) → dropped. Belongs in commit messages, not code.
2. Fragile line numbers (deepspeed_worker.py:319-322) → dropped; kept the stable file path.
3. Redundant block at __init__ → dropped; consolidated to the single comment at _create_ds_worker_config.

Final state, 4 essential Python comments (vs ~35 before your review), all ≤ 119 chars:

# Set False for non-arange position_ids (mrope / 3D rope).
# No server consumer yet; forwarded for forward-compat with group-balanced routing.
# Server reads these from ds_worker_config (arctic_platform/rl/deepspeed_worker.py), not per-call meta.
# `routing_key` kept for caller-API compat; arctic_platform.rl handles routing internally.

YAML similarly trimmed; comment density now matches upstream arctic_rl_share_v0.7.1.

[sfc-gh-sbekman]

so when we do this adaption can we avoid changing anything at all other than the layout of the files, otherwise how will we be merging this back into our branch?

is there a reason for changing things?

If something is broken we need to fix this first in the original branch and this should be only movement of code with only required by the verl team layout changes only pertaining to new code location (e.g. imports)