verl abstract remote backend

examples/arctic_rl/run_gsm8k_grpo_arl_zorro_yes.sh

@@ -70,10 +70,10 @@ python3 -m verl.trainer.main_ppo \
     remote_backend.training_gpus=1 \
     remote_backend.sampling_gpus=1 \
     remote_backend.log_prob_gpus=0 \
-    remote_backend.zero_optimization.stage=2 \
-    remote_backend.zero_optimization.offload_optimizer.device=none \
-    remote_backend.zero_optimization.offload_param.device=none \
-    remote_backend.zorro_train.enable=True \
+    remote_backend.train.deepspeed.zero_optimization.stage=2 \
+    remote_backend.train.deepspeed.zero_optimization.offload_optimizer.device=none \
+    remote_backend.train.deepspeed.zero_optimization.offload_param.device=none \
+    remote_backend.train.zorro_train.enable=True \
     trainer.critic_warmup=0 \
     trainer.logger="['console']" \
     trainer.experiment_name=gsm8k_grpo_qwen3_0p6b_ngpu1_gbs16_rolln5_zorroTrue \

verl/remote_backend/worker_utils.py

@@ -56,6 +56,22 @@ def make_njt(data: TensorDict, tensor: Tensor) -> Tensor:
     return torch.nested.nested_tensor_from_jagged(flat, cu_seqlens)
 
 
+def shift_nested_response_aligned_to_predict_next(njt: Tensor) -> Tensor:
+    """Convert a response-aligned nested tensor back to the legacy
+    "predict-next" convention by left-shifting each sequence by one.
+
+    The Arctic zorro fast path emits model output already aligned with the
+    response token at each slot (slot i holds the log prob of token[i]). Shift it
+    left so slot i holds log P(token[i+1]); ``no_padding_2_padding`` can then apply
+    a single uniform ``shift=-1`` slice for both the zorro and standard paths,
+    instead of branching the slice on a per-call ``shift`` argument."""
+    v = njt.values()
+    v_new = torch.empty_like(v)
+    v_new[:-1] = v[1:]
+    v_new[-1] = 0  # trailing slot of last seq; never read by the shift=-1 slice
+    return torch.nested.nested_tensor_from_jagged(v_new, njt.offsets())
+
+
 def normalize_backend_metrics(metrics: dict[str, Any]) -> dict[str, Any]:
     """Normalise backend metrics into verl's :class:`Metric` type.
 

verl/remote_backend/workers/arctic_rl/worker.py

@@ -33,11 +33,15 @@
 
 import torch
 from codetiming import Timer
-from omegaconf import DictConfig
+from omegaconf import DictConfig, OmegaConf
 from tensordict import TensorDict
 
 from verl.remote_backend.base import RemoteBackend, RemoteBackendRegistry
-from verl.remote_backend.worker_utils import make_njt, normalize_backend_metrics
+from verl.remote_backend.worker_utils import (
+    make_njt,
+    normalize_backend_metrics,
+    shift_nested_response_aligned_to_predict_next,
+)
 from verl.single_controller.base import Worker
 from verl.single_controller.base.decorator import (
     Dispatch,
@@ -95,13 +99,12 @@ def __init__(self, config: DictConfig, role: str, **kwargs):
         if backend_name is None:
             raise ValueError("ArcticRLActorRolloutRefWorker requires main_config.trainer.remote_backend to be set.")
 
-        # Registry has no lazy `MODULES` table any more (per @zw0610): the
-        # adapter module is imported explicitly by `main_ppo.py` when the
-        # corresponding backend is selected, which decorates the class
-        # with `@RemoteBackendRegistry.register(...)`. Here we only look
-        # it up. `from_config(handle=...)` is the sole public constructor;
-        # the handle is what makes this a re-attach rather than a fresh
-        # init.
+        # Zorro fast path emits response-aligned model output; the forwarder shifts it back to "predict-next" so `no_padding_2_padding` uses one uniform slice.
+        self.zorro_train_enable = bool(
+            OmegaConf.select(self.main_config, "remote_backend.train.zorro_train.enable", default=False)
+        )
+
+        # Registry has no lazy `MODULES` table any more (per @zw0610): the adapter module is imported explicitly by `main_ppo.py` when the backend is selected, which decorates the class with `@RemoteBackendRegistry.register(...)`. Here we only look it up. `from_config(handle=...)` is the sole public constructor; the handle is what makes this a re-attach rather than a fresh init.
         backend_cls = RemoteBackendRegistry.get(backend_name)
         self.backend: RemoteBackend = backend_cls.from_config(self.main_config, handle=backend_handle)
 
@@ -201,9 +204,15 @@ async def _run_log_prob(self, data: TensorDict, *, ref: bool, calculate_entropy:
         model_output_raw = result.get("model_output", {})
         model_output: dict[str, Any] = {}
         if "log_probs" in model_output_raw:
-            model_output["log_probs"] = make_njt(data, model_output_raw["log_probs"])
+            log_probs = make_njt(data, model_output_raw["log_probs"])
+            if self.zorro_train_enable:
+                log_probs = shift_nested_response_aligned_to_predict_next(log_probs)
+            model_output["log_probs"] = log_probs
         if calculate_entropy and "entropy" in model_output_raw:
-            model_output["entropy"] = make_njt(data, model_output_raw["entropy"])
+            entropy = make_njt(data, model_output_raw["entropy"])
+            if self.zorro_train_enable:
+                entropy = shift_nested_response_aligned_to_predict_next(entropy)
+            model_output["entropy"] = entropy
 
         metrics = dict(result.get("metrics", {}))
         metrics.setdefault("mfu", 0.0)

verl/trainer/config/remote_backend/arctic.yaml

@@ -7,7 +7,7 @@
 # Per-backend block loaded into `config.remote_backend` via Hydra defaults
 # (`+remote_backend@remote_backend=arctic`). Selected at runtime by `trainer.remote_backend=arctic`. Companion adapter:
 # `verl.workers.remote_client.arctic_rl.ArcticRLClientWrapper`.
-
+#
 # whether to use colocate mode
 colocate: True
 
@@ -23,85 +23,97 @@ log_prob_gpus: 1
 # tensor-parallel size for the sampling engine
 sampling_tp_size: 1
 
-# ZoRRO training path. Server reads these flat (`zorro_train_<key>`) from ds_worker_config and per-call meta.
-zorro_train:
+# Weight-sync options between the train and rollout engines.
+weight_sync:
+
+  # Zero-copy CUDA IPC weight sync (bypasses NCCL). Only with colocate=True.
+  cuda_ipc: False
 
-  # whether to use zorro packing for off-policy log-probs
-  enable: False
+  # Stream one param at a time during CUDA-IPC sync (cuda_ipc=True + colocate=True) to bound peak extra GPU mem to one param/GPU instead of the whole model.
+  low_memory: False
 
-  # load-balancer over zorro groups (no server consumer yet; forward-compat with group-balanced routing)
+# Training-side settings: ZoRRO, logits, determinism, and the DeepSpeed training/log-prob engine config.
+train:
+
+  # Zorro-group load balancer: when True the server reorganizes the global batch across workers (reorg_global_batch) for balanced routing; read unconditionally from per-call meta (`load_balancer`), independent of zorro_train.enable.
   load_balancer: True
 
-  # should equal actor_rollout_ref.rollout.n for best perf
-  max_rollouts: ${actor_rollout_ref.rollout.n}
+  # ZoRRO training path. Server reads these flat (`zorro_train_<key>`) from ds_worker_config and per-call meta.
+  zorro_train:
+
+    # whether to use zorro packing for off-policy log-probs
+    enable: False
+
+    # should equal actor_rollout_ref.rollout.n for best perf
+    max_rollouts: ${actor_rollout_ref.rollout.n}
+
+  # Logprob/entropy compute strategy over vocab logits; only move off `none` on OOM.
+  logits:
+
+    # none = fused single-call (fastest, ~2x peak logits mem); memory = never manifests full logits, recompute in backward (zorro_train.enable=True only); compute = manifest once then chunk softmax/entropy.
+    optimization: none
+
+    # Per-chunk peak memory budget (GiB) for `memory` / `compute` modes. Smaller -> more chunks, less peak, more compute.
+    optimization_peak_mem_size_in_gib: 4
 
-# ray or http: ray would be much faster for payload comms in the on-prem use case
-comm_protocol: ray
+    # Upcast produced logits to fp32 before logprob/entropy/loss math; no-op if already fp32.
+    compute_in_fp32: False
 
-# Drop position_ids on the wire; server reconstructs from attention_mask. Set False for mrope / 3D rope.
-drop_position_ids: True
+    # Upcast the LM-head input to fp32 so the logits projection runs in fp32 (better precision, extra head compute/mem); ZoRRO path only.
+    compute_from_fp32_inputs: False
 
-# Logprob/entropy compute strategy over vocab logits; only move off `none` on OOM.
-#   none    - fused single-call, fastest, ~2x peak logits memory.
-#   memory  - never manifests full logits; recompute in backward (zorro_train.enable=True only).
-#   compute - manifest logits once, chunk the softmax/entropy; in-between on memory/compute.
-logits_optimization: none
+  # Full-determinism controls for the RL server's DeepSpeed worker.
+  determinism:
 
-# Per-chunk peak memory budget (GiB) for `memory` / `compute` modes. Smaller -> more chunks, less peak, more compute.
-logits_optimization_peak_mem_size_in_gib: 4
+    # If True, the server calls enable_full_determinism (seeds RNGs + deterministic cuBLAS/cuDNN/NCCL/flash-attn).
+    full: False
 
-# Upcast produced logits to fp32 before logprob/entropy/loss math; no-op if already fp32.
-logits_compute_in_fp32: False
+    # seed used by enable_full_determinism when full=True
+    seed: 42
 
-# Upcast the LM-head input to fp32 so the logits projection runs in fp32. Improves
-# numerical precision at the cost of extra compute/memory for the head. ZoRRO path only.
-logits_compute_from_fp32_inputs: False
+  # DeepSpeed engine config: passed through (OmegaConf.to_container) and merged with train_micro_batch_size_per_gpu/train_batch_size/grad_accum/sequence_parallel_size. Mirrors recipe/rl-correctness arctic_rl.train.deepspeed.
+  deepspeed:
 
-# Zero-copy CUDA IPC weight sync between train and rollout engines (bypasses NCCL). Only with colocate=True.
-cuda_ipc_weight_sync: False
 
-# Stream one param at a time during CUDA-IPC sync to bound peak extra GPU mem to one full param/GPU instead of the
-# whole model. Only with cuda_ipc_weight_sync=True and colocate=True.
-low_memory_weight_sync: False
+    # ZeRO sharding stage + offload knobs
+    zero_optimization:
 
-# DeepSpeed engine config: passed through (OmegaConf.to_container) into the DS engine,
-# merged with train_micro_batch_size_per_gpu / train_batch_size / grad_accum / sequence_parallel_size.
-# Mirrors recipe/rl-correctness arctic_rl.deepspeed.
-deepspeed:
+      # 2 for optim + grad sharding; 3 also shards weights
+      stage: 0
 
-  # Torch Autocast (bfloat16 forward). Disabled by default: PyPI DeepSpeed 0.19.2 does not
-  # respect autocast cleanly; enable only on a DS build that supports it (see HANDOFF.md for the pinned clone).
-  torch_autocast:
+      offload_optimizer:
 
-    enabled: False
+        # `none` keeps the optimizer state on GPU; `cpu` offloads it
+        device: none
 
-    dtype: bfloat16
+      offload_param:
 
-  # Communication data type for cross-GPU all_reduce / all_gather
-  communication_data_type: fp32
+        # `none` keeps params on GPU; `cpu` offloads them
+        device: none
 
-  # Gradient accumulation dtype; `fp32` keeps grads in float32 but requires a DS build that
-  # recognises the `data_types` block (see HANDOFF.md for the pinned clone). Use `bf16` on
-  # stock PyPI DeepSpeed 0.19.2 to avoid the unrecognized-config branch.
-  data_types:
+    # DeepSpeed engine log level
+    log_level: INFO
 
-    grad_accum_dtype: bf16
+# Rollout-side Arctic Inference optimizations: passed through as `arctic_inference_config` (parse_arctic_inference_rollout); an all-disabled block is treated as "off". Requires a matching vllm + ArcticInference plugin build (https://github.com/snowflakedb/ArcticInference).
+rollout:
 
-  # ZeRO sharding stage + offload knobs
-  zero_optimization:
+  # Forest Cascade Attention (FCA) for efficient rollout
+  zorro_inference:
 
-    # 2 for optim + grad sharding; 3 also shards weights
-    stage: 0
+    # whether to enable FCA on the sampling engine
+    enable: False
 
-    offload_optimizer:
+  # Arctic speculative decoding draft model (empty disables speculative decoding)
+  speculative_decoding:
 
-      # `none` keeps the optimizer state on GPU; `cpu` offloads it
-      device: none
+    # path or name of the draft model
+    model: ""
 
-    offload_param:
+# Trainer <-> RL server communication settings.
+comms:
 
-      # `none` keeps params on GPU; `cpu` offloads them
-      device: none
+  # ray or http: ray would be much faster for payload comms in the on-prem use case
+  protocol: ray
 
-  # DeepSpeed engine log level
-  log_level: INFO
+  # Drop position_ids on the wire; server reconstructs from attention_mask. Set False for mrope / 3D rope.
+  drop_position_ids: True

verl/trainer/ppo/ray_trainer.py

@@ -1090,16 +1090,6 @@ def _balance_batch(self, batch: DataProto, metrics, logging_prefix="global_seqle
         )
         metrics.update(global_balance_stats)
 
-    def _model_output_shift(self) -> int:
-        """Slice offset for `no_padding_2_padding`.
-
-        * shift=-1 (legacy "predict-next"): model emits log P(token[i+1]) at slot i.
-        * shift= 0 (response-aligned): model output already lines up with the response
-          token at each slot (Arctic zorro fast path).
-        """
-        use_zorro = OmegaConf.select(self.config, "remote_backend.zorro_train.enable", default=False)
-        return 0 if use_zorro else -1
-
     def _compute_values(self, batch: DataProto) -> DataProto:
         if self.use_legacy_worker_impl == "disable":
             batch_td = batch.to_tensordict()
@@ -1135,7 +1125,7 @@ def _compute_ref_log_prob(self, batch: DataProto) -> DataProto:
             # gather output
             log_probs = tu.get(output, "log_probs")
             # step 4. No padding to padding
-            log_probs = no_padding_2_padding(log_probs, batch_td, shift=self._model_output_shift())
+            log_probs = no_padding_2_padding(log_probs, batch_td)
             # step 5: rebuild a tensordict and convert to dataproto
             ref_log_prob = tu.get_tensordict({"ref_log_prob": log_probs.float()})
             ref_log_prob = DataProto.from_tensordict(ref_log_prob)
@@ -1160,9 +1150,8 @@ def _compute_old_log_prob(self, batch: DataProto):
             routed_experts = tu.get(output, "routed_experts")
             old_log_prob_mfu = tu.get(output, "metrics")["mfu"]
             # step 4. No padding to padding
-            shift = self._model_output_shift()
-            entropy = no_padding_2_padding(entropy, batch_td, shift=shift)
-            log_probs = no_padding_2_padding(log_probs, batch_td, shift=shift)
+            entropy = no_padding_2_padding(entropy, batch_td)
+            log_probs = no_padding_2_padding(log_probs, batch_td)
             # step 5: rebuild a tensordict and convert to dataproto
             if routed_experts is not None:
                 old_log_prob = tu.get_tensordict(