feat: Arctic RL training backend integration

integrations/arctic-rl/README.md

@@ -0,0 +1,179 @@
+# Arctic RL Integration for SkyRL
+
+Routes SkyRL's GRPO training loop through the Arctic RL server — **all GPU operations** (training, generation, log-probs, weight sync) happen on the server. The SkyRL client is CPU-only.
+
+## Architecture
+
+```
+SkyRL Client (CPU-only, ray num_gpus=0)
+  - Data loading, reward scoring (skyrl-gym)
+  - Orchestration: generate → score → train
+  - HTTP calls to Arctic RL server
+        |
+        | HTTP (torch-serialized batches)
+        v
+Arctic RL Server (own Ray cluster, all GPUs)
+  - DeepSpeed Workers: forward/backward, GRPO loss, optimizer step
+  - ArcticInference (vLLM) Replicas: generation, log-probs
+  - NCCL weight sync between training and inference
+```
+
+## Validated Results
+
+### Arctic RL Server Backend (this integration)
+
+**Setup**: Qwen2.5-1.5B-Instruct, 4 DeepSpeed training GPUs + 2 ArcticInference (vLLM) sampling GPUs + 1 log-prob GPU (7x H200), GRPO with 5 samples/prompt.
+
+| Step | GSM8K Eval (pass@1) | Training Reward |
+|------|-------------------|-----------------|
+| 0 (base) | 7.8% | — |
+| 5 | 33% | 0.22 |
+| 10 | 63% | 0.60 |
+| 30 | 70% | 0.65 |
+| 45 | 73% | 0.81 |
+| 75 | **75.4%** | 0.82 |
+
+### SkyRL Default (FSDP2 baseline)
+
+| Model | GSM8K Accuracy (1,319 test) |
+|---|---|
+| Base (Qwen2.5-1.5B-Instruct) | 7.43% |
+| Trained (step 59) | **79.00%** |
+
+## Quick Start
+
+### Prerequisites
+
+- 7+ GPUs (H200/A100 recommended)
+- `arctic-skyrl` repo (this repo, `arctic-rl-integration` branch)
+- `ArcticTraining-dss` repo (`arctic-rl-grpo-loss` branch)
+- `arctic-inference` package installed
+- GSM8K dataset prepared
+
+### Step 1: Clone repos
+
+```bash
+# Clone arctic-skyrl (client)
+git clone https://github.com/snowflake-eng/arctic-skyrl.git
+cd arctic-skyrl
+git checkout arctic-rl-integration
+pip install -e ".[arctic-rl]"
+
+# Clone ArcticTraining (server) — in a separate directory
+cd ..
+git clone https://github.com/snowflakedb/ArcticTraining.git
+cd ArcticTraining
+git checkout arctic-rl-grpo-loss
+pip install --no-deps -e .
+```
+
+### Step 2: Prepare GSM8K dataset
+
+```bash
+cd arctic-skyrl
+python examples/train/gsm8k/gsm8k_dataset.py --output_dir $HOME/data/gsm8k
+```
+
+### Step 3: Run training
+
+```bash
+bash examples/train_integrations/arctic_rl/run_gsm8k_grpo_arctic.sh
+```
+
+This will:
+1. Start an Arctic RL server (DeepSpeed + ArcticInference) on localhost
+2. Initialize a CPU-only SkyRL client via Ray
+3. Run GRPO training on GSM8K with eval every 5 steps
+4. Log to console (set `LOGGER=wandb` for W&B)
+
+### Step 4: Monitor
+
+```bash
+# Watch live metrics
+tail -f /tmp/arctic_rl_training.log | grep -E "avg_raw_reward|global_step|pass_at_1"
+
+# Check GPU usage
+nvidia-smi
+```
+
+## Configuration
+
+### GPU Allocation (environment variables)
+
+| Variable | Default | Description |
+|----------|---------|-------------|
+| `ARCTIC_TRAINING_GPUS` | 4 | DeepSpeed training workers (DP) |
+| `ARCTIC_SAMPLE_GPUS` | 2 | ArcticInference (vLLM) sampling replicas |
+| `ARCTIC_LOG_PROB_GPUS` | 1 | ArcticInference (vLLM) log-prob engine |
+| `ARCTIC_SERVER_PORT` | 7000 | Server HTTP port |
+| `ARCTIC_SERVER_LOGS` | 0 | Set to 1 for verbose server output |
+| `ARCTIC_STARTUP_TIMEOUT` | 600 | Server startup timeout (seconds) |
+
+Total GPUs needed: `TRAINING + SAMPLE + LOG_PROB` (default: 7).
+
+### Key Training Parameters
+
+The launch script passes these to SkyRL via Hydra overrides:
+
+| Parameter | Value | Notes |
+|-----------|-------|-------|
+| `trainer.train_batch_size` | 256 | Prompts per step |
+| `trainer.policy_mini_batch_size` | 2 | Prompts per mini-batch |
+| `generator.n_samples_per_prompt` | 5 | Completions per prompt |
+| `trainer.policy.optimizer_config.lr` | 1e-6 | Learning rate |
+| `trainer.epochs` | 20 | Training epochs |
+| `trainer.eval_interval` | 5 | Eval every N steps |
+
+DeepSpeed config is set automatically:
+- `gradient_accumulation_steps` = `train_batch_size * n_samples / (policy_mini_batch_size * n_samples)` = 128
+- `gradient_clipping` = 1.0
+- `optimizer` = AdamW with lr from config
+
+## File Structure
+
+```
+integrations/arctic-rl/                # under integrations/, sibling of skyrl/
+├── README.md                         # This file
+├── arctic_rl/                        # importable Python package
+│   ├── __init__.py                   # Exports ArcticPPOTrainer, ArcticGenerator
+│   ├── trainer.py                    # ArcticPPOTrainer: routes training to server
+│   ├── generator.py                  # ArcticGenerator: routes generation to server vLLM
+│   ├── config.py                     # ArcticRLClientConfig builder
+│   └── entrypoint.py                 # Entrypoint: sets up client + server
+└── examples/
+    ├── setup_arctic_rl.sh            # One-command env setup
+    └── run_gsm8k_grpo_4gpu.sh        # Launch script for GSM8K GRPO
+
+skyrl/train/entrypoints/
+└── main_base.py                      # 5-line shim that routes to arctic_rl
+                                      # when trainer.arctic_rl is set in config
+```
+
+The outer folder is at `integrations/arctic-rl/` (top-level `integrations/` namespace)
+of `skyrl/`, matching the legacy `skyrl-tx/` placement. The inner Python
+package uses underscore (`arctic_rl`, the standard Python module name
+convention). It is distinct from the upstream `arctic_training` package's
+`arctic_training.arctic_rl` sub-namespace — both coexist at import time
+without collision.
+
+## How It Works
+
+1. **`arctic_rl.entrypoint`** creates an `ArcticRLClient` which spawns the server as a subprocess with a clean environment (stripped `CUDA_VISIBLE_DEVICES` and `RAY_*` vars so the server gets its own GPU access)
+
+2. **`ArcticPPOTrainer`** overrides the standard SkyRL training loop:
+   - `fwd_logprobs_values_reward` → no-op (server computes old log-probs internally)
+   - `compute_advantages_and_returns` → no-op (server computes GRPO advantages from rewards)
+   - `train_critic_and_policy` → sends batches to server via HTTP, server runs GRPO loss + backward
+
+3. **`ArcticGenerator`** routes generation to server vLLM and scores completions via `skyrl-gym`
+
+4. **Server-side `grpo_loss`** (in `processors.py`) is self-contained:
+   - Computes per-token log-probs with causal shift
+   - Derives old log-probs by detaching (correct for `update_epochs_per_batch=1`)
+   - Computes group-relative advantages from per-sequence rewards
+   - Applies PPO clipped surrogate (eps_clip=0.2)
+
+## Companion PRs
+
+- **Client (this repo)**: [`arctic-rl-integration`](https://github.com/snowflake-eng/arctic-skyrl/compare/arctic-rl-integration) branch
+- **Server (ArcticTraining-dss)**: [`arctic-rl-grpo-loss`](https://github.com/snowflakedb/ArcticTraining/compare/arctic-rl-grpo-loss) branch — PR [#20](https://github.com/snowflakedb/ArcticTraining/pull/20)

integrations/arctic-rl/arctic_rl/__init__.py

@@ -0,0 +1,28 @@
+"""Arctic RL backend for SkyRL.
+
+Provides ``ArcticPPOTrainer`` and ``ArcticGenerator`` that route all GPU
+work to an Arctic RL server, allowing any SkyRL recipe to switch backends
+by setting ``trainer.arctic_rl={}`` (and ``colocate_all: false``).
+
+Install::
+
+    uv sync --extra arctic-rl
+
+Usage in a recipe::
+
+    from arctic_rl import ArcticPPOTrainer, ArcticGenerator
+
+The folder on disk is ``integrations/arctic-rl/arctic_rl/`` — top-level sibling of
+``skyrl/`` (matching the legacy ``skyrl-tx/`` placement). The Python
+package is ``arctic_rl`` (top-level); it is distinct from the upstream
+``arctic_training`` package (which has its own ``arctic_training.arctic_rl``
+sub-namespace) — both coexist at import time without collision.
+
+Dependencies:
+    arctic_training — pip package providing ArcticRLClient/Server
+"""
+
+from arctic_rl.trainer import ArcticPPOTrainer
+from arctic_rl.generator import ArcticGenerator
+
+__all__ = ["ArcticPPOTrainer", "ArcticGenerator"]

integrations/arctic-rl/arctic_rl/config.py

@@ -0,0 +1,189 @@
+"""Arctic RL configuration types.
+
+Defines:
+- ``ArcticRLTrainerConfig``: backend-specific knobs (colocate, zero_stage, ...)
+- ``ArcticTrainerConfig``: extends core ``TrainerConfig`` with ``arctic_rl`` field
+- ``ArcticSkyRLConfig``: top-level config used by the integration's entrypoint
+- ``build_rl_config(cfg)``: translates ``SkyRLTrainConfig`` → ``ArcticRLClientConfig``
+
+These live in the integration to keep core SkyRL integration-agnostic — core only
+knows about a generic ``trainer.backend: str`` field that lazily dispatches here.
+All shared knobs (GPU counts, vLLM settings, colocation) are derived from existing
+SkyRL config fields by ``build_rl_config``.
+"""
+
+from dataclasses import dataclass
+from typing import Optional
+
+from arctic_training.arctic_rl.config import ArcticRLClientConfig
+from skyrl.train.config import SkyRLTrainConfig
+from skyrl.train.config.config import BaseConfig, TrainerConfig, make_config
+
+
+# ---------------------------------------------------------------------------
+# Arctic RL backend configuration
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class ArcticRLTrainerConfig(BaseConfig):
+    """Arctic RL (DeepSpeed) backend settings.
+
+    Only contains params unique to the Arctic RL server with no equivalent in
+    the standard SkyRL config.  Shared knobs are derived in ``build_rl_config``.
+    """
+
+    colocate: bool = False
+    """Share GPUs between training and inference on the ARL server.
+
+    Distinct from ``trainer.placement.colocate_all`` which controls Ray placement
+    groups.  ARL colocation is server-side GPU sharing — ``colocate_all`` must
+    stay ``false`` when using the ARL backend.
+    """
+    use_zorro: bool = False
+    """Enable ZoRRO (prompt deduplication) on the training server."""
+    zero_stage: int = 0
+    """DeepSpeed ZeRO stage (0, 2, or 3)."""
+    log_prob_gpus: int = 0
+    """Number of GPUs for log-prob computation (0 = skip separate log-prob)."""
+    offload_optimizer: bool = False
+    """Offload optimizer state to CPU when ``zero_stage >= 2``."""
+    host: str = "localhost"
+    """Server host for HTTP comm protocol; ignored for Ray."""
+    port: int = 7000
+    """Server port for HTTP comm protocol; ignored for Ray."""
+    startup_timeout: float = 300.0
+    """Seconds to wait for server jobs to come up."""
+    server_logs: bool = False
+    """Forward server logs to stdout for debugging."""
+
+
+@dataclass
+class ArcticTrainerConfig(TrainerConfig):
+    """``TrainerConfig`` extended with the Arctic RL field.  Used only when
+    ``trainer.backend == "arctic_rl"`` is set in the recipe."""
+
+    arctic_rl: Optional[ArcticRLTrainerConfig] = None
+    """Arctic RL backend settings.  ``None`` falls back to defaults."""
+
+
+# Top-level config for arctic_rl recipes.  Used by the integration's entrypoint
+# either directly (``uv run -m integrations.arctic_rl.entrypoint``) or via core
+# dispatch (``trainer.backend=arctic_rl`` from ``main_base``).
+ArcticSkyRLConfig = make_config(trainer_cls=ArcticTrainerConfig)
+
+
+# ---------------------------------------------------------------------------
+# Translation: SkyRLTrainConfig → ArcticRLClientConfig
+# ---------------------------------------------------------------------------
+
+
+def build_rl_config(cfg: SkyRLTrainConfig) -> ArcticRLClientConfig:
+    """Build ``ArcticRLClientConfig`` from ``SkyRLTrainConfig``.
+
+    Raises ``ValueError`` if ``cfg.trainer.arctic_rl`` is not set.
+    """
+    arl = cfg.trainer.arctic_rl
+    if arl is None:
+        raise ValueError(
+            "trainer.arctic_rl must be set when using the Arctic RL entrypoint. "
+            "Add 'trainer.arctic_rl={}' to your config overrides to enable it "
+            "with defaults, or set individual fields like "
+            "'trainer.arctic_rl.zero_stage=2'."
+        )
+
+    # -- Derived from existing SkyRL configs ---------------------------------
+    training_gpus = (
+        cfg.trainer.placement.policy_num_gpus_per_node
+        * cfg.trainer.placement.policy_num_nodes
+    )
+    sampling_gpus = cfg.generator.inference_engine.num_engines
+    colocate = arl.colocate
+    vllm_gpu_mem = cfg.generator.inference_engine.gpu_memory_utilization
+    tp_size = cfg.generator.inference_engine.tensor_parallel_size
+
+    # -- From ARL-specific config --------------------------------------------
+    lr = cfg.trainer.policy.optimizer_config.lr
+    n_samples = cfg.generator.n_samples_per_prompt
+    mini_batch_size = cfg.trainer.policy_mini_batch_size * n_samples
+    train_batch_size = cfg.trainer.train_batch_size * n_samples
+    grad_accum_steps = max(1, train_batch_size // mini_batch_size // training_gpus)
+
+    # -- vLLM config (only for colocated or TP > 1) --------------------------
+    vllm_cfg: dict | None = None
+    if colocate:
+        vllm_cfg = {
+            "gpu_memory_utilization": vllm_gpu_mem,
+            "enforce_eager": True,
+        }
+    if tp_size > 1:
+        vllm_cfg = vllm_cfg or {}
+        vllm_cfg["tensor_parallel_size"] = tp_size
+
+    # -- DeepSpeed config ----------------------------------------------------
+    ds_config = {
+        "train_micro_batch_size_per_gpu": 1,
+        "gradient_accumulation_steps": grad_accum_steps,
+        "optimizer": {
+            "type": "AdamW",
+            "params": {
+                "lr": lr,
+                "betas": list(cfg.trainer.policy.optimizer_config.betas)
+                    if hasattr(cfg.trainer.policy.optimizer_config, "betas")
+                    else [0.9, 0.999],
+                "eps": getattr(cfg.trainer.policy.optimizer_config, "eps", 1e-8),
+                "weight_decay": getattr(cfg.trainer.policy.optimizer_config, "weight_decay", 0.0),
+            },
+        },
+        "gradient_clipping": cfg.trainer.policy.optimizer_config.max_grad_norm,
+        "bf16": {"enabled": True},
+    }
+
+    zero_cfg: dict = {"stage": arl.zero_stage}
+    if arl.zero_stage >= 2 and arl.offload_optimizer:
+        zero_cfg["offload_optimizer"] = {"device": "cpu", "pin_memory": True}
+    ds_config["zero_optimization"] = zero_cfg
+
+    # -- ZoRRO worker config -------------------------------------------------
+    ds_worker_config = None
+    if arl.use_zorro:
+        ds_worker_config = {
+            "use_zorro": True,
+            "response_len": cfg.generator.sampling_params.max_generate_length,
+            "max_token_len": (cfg.trainer.max_prompt_length + cfg.generator.sampling_params.max_generate_length)
+                * cfg.trainer.policy_mini_batch_size * n_samples,
+            "rollout_n": n_samples,
+            "temperature": getattr(cfg.generator.sampling_params, "temperature", 1.0),
+            "use_unpad": True,
+        }
+
+    return ArcticRLClientConfig(
+        model_name=cfg.trainer.policy.model.path,
+        backend="local",
+        host=arl.host,
+        port=arl.port,
+        training_gpus=training_gpus,
+        sampling_gpus=sampling_gpus,
+        log_prob_gpus=arl.log_prob_gpus,
+        log_prob_engine="vllm",
+        colocate=colocate,
+        vllm_config=vllm_cfg,
+        ds_config=ds_config,
+        ds_worker_config=ds_worker_config,
+        training_config={
+            "dtype": "bfloat16",
+            "gradient_checkpointing": True,
+            "gradient_accumulation_steps": grad_accum_steps,
+            "optimizer": {
+                "lr": lr,
+                "weight_decay": 0.0,
+                "beta1": 0.9,
+                "beta2": 0.999,
+                "lr_scheduler_type": "constant",
+                "gradient_clipping": cfg.trainer.policy.optimizer_config.max_grad_norm,
+                "warmup_steps_proportion": 0.0,
+            },
+        },
+        startup_timeout=arl.startup_timeout,
+        server_logs=arl.server_logs,
+    )