feat: integrate Arctic RL backend (opt-in via [arctic])

integrations/arctic-rl/README.md

@@ -0,0 +1,54 @@
+# Arctic RL backend for PRIME-RL
+
+Opt-in integration that delegates training (forward/backward/optimizer/weight-sync)
+and rollout generation to a remote Arctic RL server. Lives entirely under
+`integrations/arctic-rl/` so core PRIME-RL has zero Arctic-specific code.
+
+## Enabling
+
+Add this to any PRIME-RL `rl.toml`:
+
+```toml
+[trainer]
+backend = "arctic_rl"
+
+[arctic]
+backend = "remote"
+url = "http://your-arctic-server:7000"
+```
+
+The PRIME-RL launcher (`prime_rl.entrypoints.rl:main`) peeks at
+`trainer.backend` before strict config parse and dynamically dispatches to
+`arctic_rl.entrypoint:main` when it's set to `"arctic_rl"`. Otherwise the
+native PRIME-RL path runs unchanged and `arctic_rl` is never imported.
+
+## Install
+
+Install this integration alongside PRIME-RL:
+
+```
+uv pip install -e ./integrations/arctic-rl
+```
+
+Plus the Arctic RL client (private until release — see Arctic RL release
+notes for the install command).
+
+## What gets replaced
+
+In Arctic mode the launcher swaps two of PRIME-RL's three subprocesses:
+
+| Slot | Native PRIME-RL | Arctic mode |
+|---|---|---|
+| Trainer | `torchrun` + FSDP2 | `arctic-trainer` (single CPU process, HTTP client) |
+| Inference | vLLM server | `arctic-shim` (FastAPI OpenAI-compat proxy) |
+| Orchestrator | verifiers + rollouts | unchanged |
+
+The orchestrator points at the shim's `base_url` (rewritten by the
+launcher); rollout files / STABLE markers / weight-sync semantics are
+all preserved.
+
+## Out of scope (rejected by validator)
+
+- Multi-node deployment
+- LoRA, multi-run, `teacher_inference`, `[inference]` (mutually exclusive)
+- Multimodal

integrations/arctic-rl/arctic_rl/__init__.py

@@ -0,0 +1,4 @@
+"""PRIME-RL ↔ Arctic RL backend integration.
+
+Activated by ``trainer.backend = "arctic_rl"`` in ``rl.toml``.
+"""

integrations/arctic-rl/arctic_rl/_trainer_entrypoint.py

@@ -0,0 +1,52 @@
+"""arctic-trainer entrypoint.
+
+Mirror of prime_rl.trainer.rl.train:main, but single-process (no torchrun)
+and running the Arctic HTTP adapter instead of a local FSDP2 model.
+
+The launcher writes trainer.toml and arctic.toml side-by-side. We load the
+trainer config via the standard `cli(TrainerConfig)` mechanism and pick up
+arctic.toml via the ARCTIC_CONFIG_TOML env var (populated by the launcher
+because pydantic-config's cli() takes one positional @-path arg at a time).
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+from pathlib import Path
+
+import tomli
+from loguru import logger
+
+from arctic_rl.config import ArcticConfig
+from arctic_rl.trainer import ArcticTrainerAdapter
+from prime_rl.configs.trainer import TrainerConfig
+from prime_rl.utils.config import cli
+from prime_rl.utils.process import set_proc_title
+
+
+def _load_arctic_config() -> ArcticConfig:
+    toml_path = os.environ.get("ARCTIC_CONFIG_TOML")
+    if not toml_path:
+        raise RuntimeError(
+            "ARCTIC_CONFIG_TOML env var not set. The arctic-trainer is meant to be "
+            "launched by rl_arctic_local which sets this variable pointing to arctic.toml."
+        )
+    path = Path(toml_path)
+    if not path.exists():
+        raise RuntimeError(f"ARCTIC_CONFIG_TOML points at non-existent path: {path}")
+    with open(path, "rb") as f:
+        data = tomli.load(f)
+    return ArcticConfig(**data)
+
+
+def main():
+    set_proc_title("ArcticTrainer")
+    trainer_cfg: TrainerConfig = cli(TrainerConfig)
+    arctic_cfg = _load_arctic_config()
+    logger.info("ArcticTrainer starting (backend={})", arctic_cfg.backend)
+    ArcticTrainerAdapter(trainer_cfg, arctic_cfg).run()
+
+
+if __name__ == "__main__":
+    sys.exit(main())

integrations/arctic-rl/arctic_rl/client.py

@@ -0,0 +1,68 @@
+"""Build the ArcticRLClient that the trainer uses to talk to the Arctic RL server."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from loguru import logger
+
+from arctic_rl.config import ArcticConfig
+
+if TYPE_CHECKING:
+    from prime_rl.configs.trainer import TrainerConfig
+
+
+def _build_training_config(trainer_cfg: TrainerConfig, arctic_cfg: ArcticConfig) -> dict:
+    """Build the dict the Arctic RL server expects at /initialize time."""
+    optim = trainer_cfg.optim
+    return {
+        "dtype": "bfloat16",
+        "gradient_checkpointing": True,
+        "max_seq_len": trainer_cfg.model.seq_len,
+        "n_gpus": arctic_cfg.training_gpus,
+        "optimizer": {
+            "lr": getattr(optim, "lr", 1e-5),
+            "weight_decay": getattr(optim, "weight_decay", 0.0),
+            "beta1": getattr(optim, "betas", (0.9, 0.999))[0] if hasattr(optim, "betas") else 0.9,
+            "beta2": getattr(optim, "betas", (0.9, 0.999))[1] if hasattr(optim, "betas") else 0.999,
+            "eps": getattr(optim, "eps", 1e-8),
+            "gradient_clipping": 1.0,
+            "lr_scheduler_type": "constant",
+            "warmup_steps_proportion": 0.0,
+        },
+    }
+
+
+def _build_vllm_config(arctic_cfg: ArcticConfig, trainer_cfg: TrainerConfig) -> dict:
+    vllm_config = dict(arctic_cfg.vllm_config or {})
+    vllm_config.setdefault("max_model_len", trainer_cfg.model.seq_len)
+    vllm_config.setdefault("tensor_parallel_size", arctic_cfg.sampling_tensor_parallel_size)
+    return vllm_config
+
+
+def build_arctic_client(arctic_cfg: ArcticConfig, trainer_cfg: TrainerConfig):
+    """Build and return an ArcticRLClient. Blocks until all jobs are RUNNING."""
+    # Lazy-imported so the integration package doesn't pull arctic_training
+    # at module-load time.
+    from arctic_training.arctic_rl.client import ArcticRLClient
+    from arctic_training.arctic_rl.config import ArcticRLClientConfig
+
+    model_name = trainer_cfg.model.name
+    client_config = ArcticRLClientConfig(
+        backend="local",
+        model_name=model_name,
+        training_config=_build_training_config(trainer_cfg, arctic_cfg),
+        vllm_config=_build_vllm_config(arctic_cfg, trainer_cfg),
+        training_gpus=arctic_cfg.training_gpus,
+        sampling_gpus=arctic_cfg.sampling_tensor_parallel_size,
+        log_prob_gpus=arctic_cfg.log_prob_gpus,
+    )
+    logger.info("Initializing ArcticRLClient (model={})", model_name)
+    client = ArcticRLClient(client_config)
+    logger.info(
+        "ArcticRLClient ready: training={} sampling={} log_prob={}",
+        client.training_job_id,
+        client.sampling_job_id,
+        getattr(client, "log_prob_job_id", None),
+    )
+    return client

integrations/arctic-rl/arctic_rl/config.py

@@ -0,0 +1,85 @@
+"""Arctic adapter config.
+
+`ArcticConfig` is attached to `RLConfig.arctic`. When `backend` is None, the
+entire adapter is inactive and Prime-RL runs natively.
+"""
+
+from typing import Annotated, Any, Literal
+
+from pydantic import Field
+
+from prime_rl.utils.config import BaseConfig
+
+
+class ArcticConfig(BaseConfig):
+    """Arctic RL integration config.
+
+    When `backend = "local"`, the launcher dispatches to the Arctic
+    entrypoint, which spawns the Arctic RL server in-process via
+    `ArcticRLClient` and replaces prime-rl's native trainer + vLLM with
+    HTTP calls to it.
+    """
+
+    backend: Annotated[
+        Literal["local"] | None,
+        Field(
+            description=(
+                "Arctic backend mode. 'local' spawns the Arctic RL server in-process. "
+                "None (default) disables Arctic and uses the native Prime-RL path."
+            )
+        ),
+    ] = None
+
+    training_gpus: Annotated[
+        int,
+        Field(description="GPUs requested for the training job.", gt=0),
+    ] = 2
+    sampling_tensor_parallel_size: Annotated[
+        int,
+        Field(description="Tensor-parallel size for the sampling engine (forwarded as vLLM tensor_parallel_size).", gt=0),
+    ] = 1
+    log_prob_gpus: Annotated[
+        int,
+        Field(description="GPUs for the log-prob job. 0 disables it."),
+    ] = 0
+
+    vllm_config: Annotated[
+        dict[str, Any] | None,
+        Field(
+            description=(
+                "Optional vLLM overrides forwarded to Arctic RL sampling/log-prob engines. "
+                "max_model_len defaults to trainer.model.seq_len when omitted."
+            )
+        ),
+    ] = None
+
+    enable_thinking: Annotated[
+        bool,
+        Field(description="Pass enable_thinking to tokenizer.apply_chat_template."),
+    ] = False
+
+    use_cispo_loss: Annotated[
+        bool,
+        Field(
+            description=(
+                "Use CISPO (Clipped IS-weight Policy Optimization) instead of vanilla "
+                "PPO-CLIP. CISPO clips the importance-sampling ratio with a stop-gradient, "
+                "so every token including clipped ones contributes a non-zero gradient. "
+                "Recommended for async-pipeline off-policy training (ScaleRL §3.2). "
+                "Uses asymmetric clips eps=0.2 / eps_higher=0.28 per the paper."
+            )
+        ),
+    ] = False
+
+    loss_agg_mode: Annotated[
+        str,
+        Field(
+            description=(
+                "Loss aggregation mode forwarded to the server's grpo_loss. "
+                "'token-mean' (default): average over all unmasked tokens. "
+                "'prompt-mean': average per rollout first, then across rollouts — "
+                "requires prompt_group_ids (automatically built when example_ids are tracked). "
+                "Other valid values: 'seq-mean-token-sum', 'seq-mean-token-sum-norm', 'seq-mean-token-mean'."
+            )
+        ),
+    ] = "token-mean"

integrations/arctic-rl/arctic_rl/context.py

@@ -0,0 +1,118 @@
+"""Translate PRIME-RL packed microbatches to Arctic RL's batch format.
+
+PRIME-RL packs multiple rollouts into ``[1, T]`` microbatches; Arctic RL
+expects ``[B, max_S]`` with one rollout per row. We unpack, apply a
+per-rollout left-shift (so labels match Arctic's shifted-index convention)
+and zero the loss mask at the wrap-around position.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from arctic_rl.unpack import iter_rollout_slices, unpack_packed_microbatch
+from prime_rl.trainer.rl.data import TensorMicroBatch
+
+# Pad values mirror PRIME-RL's pad_micro_batch so the padded batch is
+# numerically indistinguishable from a natively-padded one.
+_PAD_VALUES: dict[str, float | int | bool] = {
+    "input_ids": 1,
+    "position_ids": 0,
+    "old_log_probs_shifted": 0.0,
+    "advantages": 0.0,
+    "teacher_log_probs_shifted": 0.0,
+    "loss_mask": False,
+}
+
+
+def _extract_and_roll_rollout(mb: TensorMicroBatch, start: int, end: int) -> dict:
+    """Extract one rollout's slice from a packed mb, applying per-rollout `torch.roll(-1)`.
+
+    Returns 1-D tensors (no batch dim). The roll is scoped to this rollout so
+    cross-rollout contamination at boundary positions is impossible.
+    """
+    out: dict = {
+        "input_ids": mb["input_ids"][0, start:end].clone(),
+        "position_ids": mb["position_ids"][0, start:end].clone(),
+        "old_log_probs_shifted": torch.roll(mb["inference_logprobs"][0, start:end], shifts=-1, dims=-1),
+        "advantages": torch.roll(mb["advantages"][0, start:end], shifts=-1, dims=-1),
+    }
+    loss_mask = torch.roll(mb["loss_mask"][0, start:end], shifts=-1, dims=-1).clone()
+    # The last position wraps around within this rollout; not a valid target.
+    loss_mask[-1] = False
+    out["loss_mask"] = loss_mask
+
+    teacher = mb.get("teacher_logprobs")
+    if teacher is not None:
+        out["teacher_log_probs_shifted"] = torch.roll(teacher[0, start:end], shifts=-1, dims=-1)
+    return out
+
+
+def microbatches_to_arctic_context(mbs: list[TensorMicroBatch]) -> dict:
+    """Consolidate all rollouts from a list of packed microbatches into one `[B, max_S]` batch.
+
+    Output:
+        A dict of `[B_total, max_S]` tensors where `B_total = sum of rollouts
+        across all mbs` (trailing pad per mb is dropped) and `max_S` is the
+        longest rollout across the batch. Includes a real `attention_mask`.
+    """
+    assert mbs, "Expected at least one microbatch"
+
+    rollouts: list[dict] = []
+    raw_example_ids: list[int] = []
+    for mb in mbs:
+        slices = iter_rollout_slices(mb)
+        mb_ids = mb.get("example_ids") or []
+        for i, (start, end) in enumerate(slices):
+            rollouts.append(_extract_and_roll_rollout(mb, start, end))
+            if i < len(mb_ids):
+                raw_example_ids.append(mb_ids[i])
+
+    b = len(rollouts)
+    max_s = max(r["input_ids"].shape[0] for r in rollouts)
+    ref = rollouts[0]
+
+    out: dict = {}
+    for key, ref_tensor in ref.items():
+        pad_val = _PAD_VALUES.get(key, 0)
+        tensor = torch.full((b, max_s), pad_val, dtype=ref_tensor.dtype, device=ref_tensor.device)
+        for i, r in enumerate(rollouts):
+            length = r[key].shape[0]
+            tensor[i, :length] = r[key]
+        out[key] = tensor
+
+    attention_mask = torch.zeros((b, max_s), dtype=ref["input_ids"].dtype, device=ref["input_ids"].device)
+    for i, r in enumerate(rollouts):
+        attention_mask[i, : r["input_ids"].shape[0]] = 1
+    out["attention_mask"] = attention_mask
+
+    # Build prompt_group_ids for prompt-mean loss aggregation. Maps each rollout row to its
+    # example group so agg_loss("prompt-mean") averages token losses per rollout before
+    # averaging across rollouts. Only present when all rollouts carried an example_id.
+    if len(raw_example_ids) == b:
+        out["prompt_group_ids"] = torch.tensor(raw_example_ids, dtype=torch.long)
+
+    return out
+
+
+def microbatch_to_arctic_context(mb: TensorMicroBatch) -> dict:
+    """Legacy single-microbatch path kept for unit tests.
+
+    Prefer `microbatches_to_arctic_context([mb, ...])` in the training loop
+    so that a single-rollout bin doesn't trip the server's DP-chunk assertion.
+    """
+    loss_mask = torch.roll(mb["loss_mask"], shifts=-1, dims=-1).clone()
+    loss_mask[..., -1] = False
+
+    rolled: dict = {
+        "input_ids": mb["input_ids"],
+        "position_ids": mb["position_ids"],
+        "old_log_probs_shifted": torch.roll(mb["inference_logprobs"], shifts=-1, dims=-1),
+        "advantages": torch.roll(mb["advantages"], shifts=-1, dims=-1),
+        "loss_mask": loss_mask,
+    }
+    teacher = mb.get("teacher_logprobs")
+    if teacher is not None:
+        rolled["teacher_log_probs_shifted"] = torch.roll(teacher, shifts=-1, dims=-1)
+
+    return unpack_packed_microbatch(rolled)