diff --git a/nemo_rl/algorithms/async_utils/staleness_sampler.py b/nemo_rl/algorithms/async_utils/staleness_sampler.py new file mode 100644 index 0000000000..c729d2f4c9 --- /dev/null +++ b/nemo_rl/algorithms/async_utils/staleness_sampler.py @@ -0,0 +1,455 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Prompt-group staleness policies over a TQReplayBuffer. + +A *staleness policy* owns the whole off-policyness contract for the SC async +loop in one object, injected into both pumps: + + - ``admit`` (rollout-pump side): block until the next prompt batch may + dispatch, then return the ``target_step`` stamp for that batch (``None`` + when the policy doesn't stamp). Owning admission here is what lets the + rollout pump follow whichever sampling algorithm is selected without a + second, hand-kept copy of the gating logic. + - ``select`` / ``evict`` (train-pump side): pick / drop prompt groups. + - ``is_on_policy`` / ``required_buffer_capacity`` (derived facts): what the + weight-sync and capacity-validation paths need *without* re-reading raw + knobs, so those consumers can't drift out of sync with the sampler. + +``PromptGroupSampler`` is the interface; ``WindowedSampler`` / +``WeightFifoSampler`` / ``InOrderSampler`` are the built-in policies, one per +behavior, each owning only the args that apply to it. ``create_sampler`` builds +one from a discriminated-union config (or a ``module:ClassName`` FQN for a +policy defined outside this repo) — the config's ``name`` is the single source +of truth for which behavior runs, so there are no cross-field knob combinations +to validate. +""" + +from __future__ import annotations + +import abc +import asyncio +import importlib +from typing import ( + Annotated, + Callable, + Literal, + Optional, + Protocol, + Union, + runtime_checkable, +) + +from pydantic import BaseModel, Field + +from nemo_rl.algorithms.async_utils.replay_buffer import TQReplayBuffer +from nemo_rl.data_plane import KVBatchMeta + +# Poll interval for the rollout-pump admission gate. +_GATE_POLL_SECONDS = 0.005 + + +@runtime_checkable +class PromptGroupSampler(Protocol): + """Staleness policy shared by the SC rollout and train pumps. + + Implement this (or subclass ``BaseSampler``) to add a custom sampling + algorithm; point ``async_rl.sampler`` at ``module:ClassName`` to load it. + """ + + async def admit(self, *, trainer_version_fn: Callable[[], int]) -> Optional[int]: + """Block until the next prompt batch may dispatch. + + Args: + trainer_version_fn: Zero-arg accessor for the live trainer version + (polled, so a blocking policy sees updates while it waits). + + Returns: + The ``target_step`` to stamp on this batch's slots, or ``None`` when + the policy does not stamp target steps. + """ + ... + + async def select( + self, + *, + current_train_weight: int, + min_prompt_groups: int, + max_prompt_groups: int, + ) -> tuple[Optional[KVBatchMeta], int]: + """Pick up to ``max_prompt_groups`` eligible groups; drop them locally.""" + ... + + async def evict(self, *, current_train_weight: int) -> int: + """Drop groups that can no longer be selected; clear their DP rows.""" + ... + + @property + def is_on_policy(self) -> bool: + """True when the policy admits zero staleness (sync mode).""" + ... + + def required_buffer_capacity(self, groups_per_step: int) -> Optional[int]: + """Buffer-capacity the policy needs, or ``None`` if unconstrained.""" + ... + + +class BaseSampler(abc.ABC): + """Shared machinery for the built-in policies. + + Owns the monotonic dispatch counter (the batch index formerly tracked as + ``SingleControllerActor._max_rollout_version``) and the common + select-finalize / weight-window-evict helpers. + """ + + def __init__(self, buffer: TQReplayBuffer) -> None: + self._buffer = buffer + # Pre-incremented before each admitted batch, so -1 lets the first + # batch through a zero-staleness gate. + self._dispatch_index: int = -1 + + # ── rollout-pump side ──────────────────────────────────────────────── + @abc.abstractmethod + async def admit( + self, *, trainer_version_fn: Callable[[], int] + ) -> Optional[int]: ... + + # ── train-pump side ────────────────────────────────────────────────── + @abc.abstractmethod + async def select( + self, + *, + current_train_weight: int, + min_prompt_groups: int, + max_prompt_groups: int, + ) -> tuple[Optional[KVBatchMeta], int]: ... + + async def evict(self, *, current_train_weight: int) -> int: + """Default: drop *ready* groups below the weight window. + + Skips unready (reserved-but-uncommitted) slots so eviction can't race a + concurrent ``commit`` that re-looks-up the slot after its ``await``. + Policies whose ``select`` key isn't the start weight (e.g. + ``InOrderSampler``) override this so evict and select agree. + """ + min_valid_version = max(0, current_train_weight - self._eviction_window()) + stale_idxs = [ + i + for i, weight in enumerate(self._buffer.start_weight_list) + if weight < min_valid_version and self._buffer.ready_list[i] + ] + if not stale_idxs: + return 0 + return await self._buffer.remove(stale_idxs, remove_in_dp=True) + + # ── derived facts ──────────────────────────────────────────────────── + @property + def is_on_policy(self) -> bool: + return self._eviction_window() == 0 + + def required_buffer_capacity(self, groups_per_step: int) -> Optional[int]: + return None + + # ── shared helpers ─────────────────────────────────────────────────── + def _eviction_window(self) -> int: + """Weight-version span kept selectable; drives the default ``evict``.""" + return 0 + + @staticmethod + def _validate_group_bounds(min_prompt_groups: int, max_prompt_groups: int) -> None: + if min_prompt_groups < 1: + raise ValueError(f"min_prompt_groups must be >= 1, got {min_prompt_groups}") + if max_prompt_groups < min_prompt_groups: + raise ValueError( + f"max_prompt_groups ({max_prompt_groups}) must be >= " + f"min_prompt_groups ({min_prompt_groups})" + ) + + async def _finalize_selection( + self, + valid_idxs: list[int], + min_prompt_groups: int, + max_prompt_groups: int, + ) -> tuple[Optional[KVBatchMeta], int]: + """Cap, drop from the buffer, and concat the chosen groups. + + Greedy without waiting: returns all currently-eligible groups up to + ``max_prompt_groups`` (never fewer on purpose, never waits to fill it), + or ``(None, 0)`` below ``min_prompt_groups``. + """ + if len(valid_idxs) < min_prompt_groups: + return None, 0 + requested_groups = min(len(valid_idxs), max_prompt_groups) + selected_idxs = valid_idxs[:requested_groups] + selected_metas = [self._buffer.meta_list[i] for i in selected_idxs] + await self._buffer.remove(selected_idxs, remove_in_dp=False) + return ( + selected_metas[0].concat(*selected_metas[1:]), # type: ignore + len(selected_idxs), + ) + + +class WindowedSampler(BaseSampler): + """Over-sampled windowed selection. + + Rollout never gates on the trainer version — the pump keeps producing and + samples aged past the window are evicted. ``select`` takes any ready group + within ``[train_weight - max_staleness_versions, train_weight]``, optionally + freshest-first. + """ + + def __init__( + self, + buffer: TQReplayBuffer, + *, + max_staleness_versions: int, + sample_freshest_first: bool = False, + ) -> None: + super().__init__(buffer) + if max_staleness_versions < 0: + raise ValueError( + f"max_staleness_versions must be non-negative, got " + f"{max_staleness_versions}" + ) + self.max_staleness_versions = max_staleness_versions + self.sample_freshest_first = sample_freshest_first + + def _eviction_window(self) -> int: + return self.max_staleness_versions + + async def admit(self, *, trainer_version_fn: Callable[[], int]) -> Optional[int]: + # Over-sampled: dispatch is bounded by buffer capacity, not by version. + return None + + async def select( + self, + *, + current_train_weight: int, + min_prompt_groups: int, + max_prompt_groups: int, + ) -> tuple[Optional[KVBatchMeta], int]: + self._validate_group_bounds(min_prompt_groups, max_prompt_groups) + min_valid_version = max(0, current_train_weight - self.max_staleness_versions) + valid_idxs = [ + i + for i, weight in enumerate(self._buffer.start_weight_list) + if min_valid_version <= weight <= current_train_weight + and self._buffer.ready_list[i] + ] + if self.sample_freshest_first: + valid_idxs.sort( + key=lambda i: ( + current_train_weight - self._buffer.start_weight_list[i], + i, + ) + ) + return await self._finalize_selection( + valid_idxs, min_prompt_groups, max_prompt_groups + ) + + +class _GatedSampler(BaseSampler): + """Base for policies that admit exactly one dispatch batch per trainer step. + + The gate bounds how far generation may run ahead of the trainer + (``gate_window`` versions of lookahead). + """ + + def __init__(self, buffer: TQReplayBuffer, *, gate_window: int) -> None: + super().__init__(buffer) + if gate_window < 0: + raise ValueError(f"gate_window must be non-negative, got {gate_window}") + self._gate_window = gate_window + + def _eviction_window(self) -> int: + return self._gate_window + + def required_buffer_capacity(self, groups_per_step: int) -> Optional[int]: + # One batch of lookahead per version in the window, plus the live batch. + return groups_per_step * (self._gate_window + 1) + + async def admit(self, *, trainer_version_fn: Callable[[], int]) -> Optional[int]: + while self._dispatch_index >= trainer_version_fn() + self._gate_window: + await asyncio.sleep(_GATE_POLL_SECONDS) + self._dispatch_index += 1 + return self._stamp() + + def _stamp(self) -> Optional[int]: + return None + + +class WeightFifoSampler(_GatedSampler): + """Gated, strict weight-version FIFO. + + ``select`` drains the oldest in-window ``start_weight`` first and waits for + that weight's batch to fill. Evict uses the weight window (default). + """ + + def __init__(self, buffer: TQReplayBuffer, *, max_staleness_versions: int) -> None: + super().__init__(buffer, gate_window=max_staleness_versions) + self.max_staleness_versions = max_staleness_versions + + async def select( + self, + *, + current_train_weight: int, + min_prompt_groups: int, + max_prompt_groups: int, + ) -> tuple[Optional[KVBatchMeta], int]: + self._validate_group_bounds(min_prompt_groups, max_prompt_groups) + min_valid_version = max(0, current_train_weight - self.max_staleness_versions) + in_window = [ + weight + for weight in self._buffer.start_weight_list + if min_valid_version <= weight <= current_train_weight + ] + if not in_window: + return None, 0 + target_version = min(in_window) + valid_idxs = [ + i + for i, weight in enumerate(self._buffer.start_weight_list) + if weight == target_version and self._buffer.ready_list[i] + ] + return await self._finalize_selection( + valid_idxs, min_prompt_groups, max_prompt_groups + ) + + +class InOrderSampler(_GatedSampler): + """Gated, exact batch->step matching. + + Each dispatched batch is stamped with its dispatch index as ``target_step``; + ``select`` takes the batch whose ``target_step`` equals the trainer version + (the staleness window is not used for selection). ``evict`` is keyed on + ``target_step`` — not the start weight — so a slot whose target step is still + upcoming is never dropped early, and evict/select can't disagree. + """ + + def __init__(self, buffer: TQReplayBuffer, *, max_lookahead_versions: int) -> None: + super().__init__(buffer, gate_window=max_lookahead_versions) + self.max_lookahead_versions = max_lookahead_versions + + def _stamp(self) -> Optional[int]: + return self._dispatch_index + + async def select( + self, + *, + current_train_weight: int, + min_prompt_groups: int, + max_prompt_groups: int, + ) -> tuple[Optional[KVBatchMeta], int]: + self._validate_group_bounds(min_prompt_groups, max_prompt_groups) + valid_idxs = [ + i + for i, target in enumerate(self._buffer.target_step_list) + if target == current_train_weight and self._buffer.ready_list[i] + ] + return await self._finalize_selection( + valid_idxs, min_prompt_groups, max_prompt_groups + ) + + async def evict(self, *, current_train_weight: int) -> int: + # Keyed on target_step (matches `select`): a ready group whose target + # step has already passed can never be selected, so it is stale. Unready + # slots are skipped to avoid racing a concurrent commit. + stale_idxs = [ + i + for i, target in enumerate(self._buffer.target_step_list) + if target is not None + and target < current_train_weight + and self._buffer.ready_list[i] + ] + if not stale_idxs: + return 0 + return await self._buffer.remove(stale_idxs, remove_in_dp=True) + + +# ── config + factory ──────────────────────────────────────────────────────── + + +class WindowedSamplerConfig(BaseModel, extra="allow"): + name: Literal["windowed"] = "windowed" + # Max weight-version gap a selected group may have from the trainer. + max_staleness_versions: int = 1 + # Prefer smallest lag when picking from the in-window set. + sample_freshest_first: bool = False + + +class WeightFifoSamplerConfig(BaseModel, extra="allow"): + name: Literal["weight_fifo"] = "weight_fifo" + # Lookahead + selectable weight window, in trainer versions. + max_staleness_versions: int = 1 + + +class InOrderSamplerConfig(BaseModel, extra="allow"): + name: Literal["in_order"] = "in_order" + # How far generation may run ahead of the trainer, in dispatch batches. + max_lookahead_versions: int = 1 + + +class CustomSamplerConfig(BaseModel, extra="allow"): + name: Literal["custom"] = "custom" + # "module:ClassName" of a PromptGroupSampler defined outside this repo. + # Extra keys are forwarded to the constructor (after ``buffer``). + target: str + + +# Discriminated on ``name`` so each variant carries only its own typed args and +# pydantic narrows the type at construction — invalid arg combinations are +# unrepresentable rather than caught by a runtime assert. +SamplerConfig = Annotated[ + Union[ + WindowedSamplerConfig, + WeightFifoSamplerConfig, + InOrderSamplerConfig, + CustomSamplerConfig, + ], + Field(discriminator="name"), +] + + +def create_sampler( + buffer: TQReplayBuffer, + cfg: SamplerConfig, +) -> PromptGroupSampler: + """Build a sampler from its config (or import one by FQN).""" + if isinstance(cfg, WindowedSamplerConfig): + return WindowedSampler( + buffer, + max_staleness_versions=cfg.max_staleness_versions, + sample_freshest_first=cfg.sample_freshest_first, + ) + if isinstance(cfg, WeightFifoSamplerConfig): + return WeightFifoSampler( + buffer, max_staleness_versions=cfg.max_staleness_versions + ) + if isinstance(cfg, InOrderSamplerConfig): + return InOrderSampler(buffer, max_lookahead_versions=cfg.max_lookahead_versions) + if isinstance(cfg, CustomSamplerConfig): + module_name, sep, class_name = cfg.target.partition(":") + if not sep: + raise ValueError( + f"custom sampler target must be 'module:ClassName', got {cfg.target!r}" + ) + sampler_cls = getattr(importlib.import_module(module_name), class_name) + sampler = sampler_cls(buffer, **(cfg.model_extra or {})) + if not isinstance(sampler, PromptGroupSampler): + raise TypeError( + f"{cfg.target} does not implement the PromptGroupSampler " + f"interface (needs admit/select/evict)" + ) + return sampler + raise ValueError(f"unknown sampler config {type(cfg).__name__}") diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 16a0461f31..da63413e3d 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -28,11 +28,11 @@ Data flow: _rollout_pump → gen.generate_and_push(prompt, dp_client) ← RPC to GenWorker GenWorker → dp_client.put_samples(...) - _train_pump → dp_client.claim_meta(...) → StalenessSampler + _train_pump → sampler.evict/select against TQReplayBuffer → _advantage_stage(meta) → dp_client.get_samples(...) → adv_estimator.compute_advantage(...) → dp_client.put_samples(...) - → trainer.begin/train_microbatch/finish_train_step (split API, + → trainer.begin/train_microbatches/finish_train_step (split API, driver-side TQPolicy via asyncio.to_thread) Trainer → dp_client.get_samples(...) (via its own client) → dp_client.clear_samples(...) ← SC clears after train @@ -43,20 +43,25 @@ import asyncio import time -from contextlib import nullcontext from functools import partial from typing import Any, Literal, Optional import ray import torch from pydantic import BaseModel, Field -from tensordict import TensorDict +from nemo_rl.algorithms.async_utils.staleness_sampler import ( + InOrderSamplerConfig, + SamplerConfig, + create_sampler, +) from nemo_rl.algorithms.grpo import GRPOLoggerConfig -from nemo_rl.algorithms.staleness_sampler import ( - StalenessSampler, - count_groups, - min_weight_version, +from nemo_rl.algorithms.single_controller_utils.utils import ( + aggregate_step_metrics, + fields_for_put, + reduce_advantage_pump_metrics, + squeeze_trailing_unit_dim, + tensor_field, ) from nemo_rl.data.interfaces import DatumSpec from nemo_rl.data_plane import KVBatchMeta @@ -99,52 +104,41 @@ class SingleControllerConfig(BaseModel, extra="allow"): are validated at construction by pydantic — no runtime assert needed. """ - # Staleness. A "group" is the atomic training unit: group_size + # Group geometry. A "group" is the atomic training unit: group_size # samples sharing one source prompt. GRPO consumers set group_size = # num generations per prompt; SFT/OPD-style consumers degenerate # cleanly to group_size=1 (every sample its own group). - max_weight_staleness_versions: int = 1 min_groups_per_batch: int = 2 target_groups_per_step: Optional[int] = None group_size: int = 4 - batch_selection_strategy: Literal[ - "strict_on_policy", - "staleness_window", - ] = "strict_on_policy" # Concurrency limits max_inflight_prompts: int = 8 max_buffered_rollouts: int = 8 # _buffer_capacity semaphore size - # Rollout dispatch gating (read by _rollout_pump). over_sampling=False - # gates each batch on max_rollout_version vs trainer_version; force_in_order - # stamps target_step on each dispatch so downstream consumers can match - # rollout batches to trainer steps exactly. - over_sampling: bool = False - force_in_order: bool = False + # Staleness policy — the single source of truth for how rollouts are gated + # (rollout pump) and selected/evicted (train pump). Discriminated on + # ``name`` (windowed | weight_fifo | in_order | custom); each variant carries + # only its own args, so invalid combinations are unrepresentable and there + # are no cross-field knobs to validate at runtime. ``custom`` takes a + # ``module:ClassName`` target to load a PromptGroupSampler from outside this + # repo. Default mirrors the exemplar (exact batch->step matching). + sampler: SamplerConfig = Field(default_factory=InOrderSamplerConfig) # Training max_train_steps: int = 10 - max_rollout_prompts: int = 32 # Bounded dataset passes, mirroring grpo.py's max_num_epochs loop. One - # epoch = one pass over the prompt list. None preserves the unbounded - # behavior: max_rollout_prompts alone caps dispatch (cycling through - # the prompt list) and no epoch accounting is performed. + # epoch = one pass over the dataloader. None = unbounded: the dataloader + # is iterated indefinitely until _train_pump reaches max_train_steps. max_num_epochs: Optional[int] = None - # Worker-side optimizer mini-batch size, in samples. SC opens one - # begin / microbatch×K / finish cycle (= one opt.step) per - # ``train_global_batch_size`` worth of samples. Number of opt.steps per - # outer SC step is - # ``target_groups_per_step * group_size // train_global_batch_size``. - # If None, coerced at __init__ to one mini-batch per outer step - # (samples_per_step), preserving single-opt.step-per-outer-step behavior. + # Worker-side global batch size, in samples. SC currently supports exactly + # one optimizer step per outer RL step, so an explicit value must equal + # ``target_groups_per_step * group_size``. If None, it is set to that value + # during initialization. train_global_batch_size: Optional[int] = None # DataPlane partition partition_id: str = "rollout_data" - consumer_task_name: str = "train" - claim_required_fields: list[str] = Field(default_factory=lambda: ["input_ids"]) - max_claim_groups: int = 8 # Advantage calculation. The TQ partition column names for prompt_ids / # reward / token_mask / sample_mask / advantages are fixed conventions @@ -168,58 +162,6 @@ class SingleControllerConfig(BaseModel, extra="allow"): logger: GRPOLoggerConfig -def warn_if_staleness_window_below_minibatches( - cfg: SingleControllerConfig, -) -> None: - """Warn when the staleness window is too small to cover one outer step. - - ``_train_pump`` runs ``num_minibatches`` begin/finish cycles per outer - step, bumping ``_trainer_version`` after each ``finish_train_step`` - (one opt.step) but refreshing generation weights (``_sync_weights``) - only once, at the end of the outer step. A group produced at the version - the step started on therefore ages by up to ``num_minibatches - 1`` - versions before the next sync. When the effective staleness window is - smaller than that, such groups become unselectable mid-step — the - sampler skips them and ``_evict_stale_claimed`` may drop them as stale — - so the pump can spin or thrash instead of consuming them. - - Warns rather than raises: a run may intentionally accept the resulting - eviction churn. Self-contained (does not assume ``cfg`` has been coerced - by ``__init__``) so it is unit-testable without constructing the actor. - - Args: - cfg: SingleController config. ``strict_on_policy`` pins the effective - window to 0 regardless of ``max_weight_staleness_versions``. - """ - effective_window = ( - 0 - if cfg.batch_selection_strategy == "strict_on_policy" - else cfg.max_weight_staleness_versions - ) - target_groups = cfg.target_groups_per_step or cfg.min_groups_per_batch - samples_per_step = target_groups * cfg.group_size - train_global_batch_size = cfg.train_global_batch_size or samples_per_step - groups_per_minibatch = train_global_batch_size // cfg.group_size - if groups_per_minibatch <= 0: - return - num_minibatches = target_groups // groups_per_minibatch - if num_minibatches > 1 and effective_window < num_minibatches - 1: - print( - f"WARNING: max_weight_staleness_versions (effective window " - f"{effective_window}) is smaller than num_minibatches - 1 " - f"({num_minibatches - 1}): each outer step runs {num_minibatches} " - f"optimizer steps but syncs generation weights only once, at the " - f"end, so groups produced at the version the step began on age by " - f"up to {num_minibatches - 1} versions before the next sync and " - f"become unselectable (skipped, then evicted as stale) mid-step " - f"— the train pump may spin or thrash. Remedy: raise " - f"max_weight_staleness_versions to at least {num_minibatches - 1}, " - f"or lower target_groups_per_step * group_size / " - f"train_global_batch_size to reduce num_minibatches.", - flush=True, - ) - - @ray.remote(num_cpus=1, num_gpus=0) # pragma: no cover class SingleControllerActor: """CPU-only Ray actor that orchestrates the RL training loop. @@ -236,26 +178,29 @@ class SingleControllerActor: def __init__( self, cfg: SingleControllerConfig, - prompts: list[str], dp_client_handle: Any, gen_handle: Any, trainer_handle: Any, weight_synchronizer: Any, loss_fn: Any, + rollout_manager: Any, advantage_estimator: Any | None = None, - rollout_manager: Any = None, dataloader: Any = None, + tq_buffer: Any = None, ) -> None: self._cfg = cfg - self._prompts = prompts self._dp_client = dp_client_handle self._gen = gen_handle self._trainer = trainer_handle self._weight_synchronizer = weight_synchronizer self._loss_fn = loss_fn self._advantage_estimator = advantage_estimator - self._rollout_manager = rollout_manager self._dataloader = dataloader + self._buffer = tq_buffer + self._rollout_manager = rollout_manager + # Rebind so writer and sampler share one buffer instance even + # when Ray deserializes rollout_manager and tq_buffer separately. + self._rollout_manager._tq_buffer = self._buffer # Built here, not on the driver: Logger backends (wandb/tb/...) hold # _thread.lock that Ray can't cloudpickle into the actor. @@ -267,27 +212,10 @@ def __init__( "advantage_enabled=True requires an advantage_estimator instance" ) - # batch_selection_strategy is Literal-typed; pydantic rejects unknown - # values at config construction, so no runtime assert is needed here. - if cfg.batch_selection_strategy == "strict_on_policy": - cfg.max_weight_staleness_versions = 0 - cfg.over_sampling = False - print( - "Using strict_on_policy, auto setting max_weight_staleness_versions " - "to 0 and over_sampling to False.", - flush=True, - ) - if cfg.max_weight_staleness_versions == 0 and cfg.over_sampling: - raise ValueError( - "max_weight_staleness_versions=0 requires over_sampling=False: " - "with zero staleness the dispatch gate needs to advance one batch " - "per trainer_version, which over_sampling=True bypasses." - ) - if cfg.force_in_order and cfg.over_sampling: - raise ValueError( - "force_in_order=True requires over_sampling=False so that each " - "dispatched batch corresponds to exactly one target training step." - ) + # Staleness-mode validation that used to live here (strict_on_policy + # coercion, over_sampling/force_in_order mutual-exclusion raises) is gone: + # the sampler config's discriminated union makes those states + # unrepresentable — see ``cfg.sampler`` / ``create_sampler``. if cfg.target_groups_per_step is None: cfg.target_groups_per_step = cfg.min_groups_per_batch @@ -297,11 +225,10 @@ def __init__( f"must be >= min_groups_per_batch ({cfg.min_groups_per_batch})" ) - # Mini-batching contract: SC opens one begin / microbatch×N / finish - # cycle (= one opt.step) per train_global_batch_size worth of samples. - # Matches the sync path's gb_idx loop in the worker (one opt.step - # per gbs slice). When unset, coerce to a single mini-batch per - # outer step so behavior matches the pre-mini-batch design. + # SC opens exactly one begin / microbatch×N / finish cycle + # (= one optimizer step) per outer RL step. Until multiple optimizer + # steps per RL step are supported, the global batch must contain every + # sample selected for that step. samples_per_step = cfg.target_groups_per_step * cfg.group_size if cfg.train_global_batch_size is None: cfg.train_global_batch_size = samples_per_step @@ -310,37 +237,52 @@ def __init__( f"train_global_batch_size must be > 0, " f"got {cfg.train_global_batch_size}" ) - if samples_per_step % cfg.train_global_batch_size != 0: + if cfg.train_global_batch_size != samples_per_step: raise ValueError( - f"target_groups_per_step ({cfg.target_groups_per_step}) " - f"* group_size ({cfg.group_size}) " - f"= {samples_per_step} samples must be divisible by " - f"train_global_batch_size ({cfg.train_global_batch_size})" + f"train_global_batch_size ({cfg.train_global_batch_size}) must " + f"equal target_groups_per_step ({cfg.target_groups_per_step}) " + f"* group_size ({cfg.group_size}) = {samples_per_step}; multiple " + f"optimizer steps per RL step are not supported" ) - if cfg.train_global_batch_size % cfg.group_size != 0: + # Staleness policy owns admit (rollout side) + select/evict (train side) + # + is_on_policy / required_buffer_capacity (derived facts). ``name`` in + # the config is the single switch for which behavior runs. + self._sampler = create_sampler(self._buffer, cfg.sampler) + + # Capacity requirement is a derived fact the sampler owns, so this check + # can't drift from the admission logic. None means the policy self-bounds + # via backpressure (e.g. the over-sampled windowed policy). + required_capacity = self._sampler.required_buffer_capacity( + cfg.target_groups_per_step + ) + if ( + required_capacity is not None + and cfg.max_buffered_rollouts < required_capacity + ): raise ValueError( - f"train_global_batch_size ({cfg.train_global_batch_size}) must " - f"be divisible by group_size " - f"({cfg.group_size}) so a mini-batch contains " - f"whole prompt groups" + f"max_buffered_rollouts ({cfg.max_buffered_rollouts}) is below " + f"the {type(self._sampler).__name__}'s required capacity " + f"({required_capacity} = target_groups_per_step " + f"{cfg.target_groups_per_step} * (lookahead + 1)); the rollout " + f"pump would deadlock waiting for buffer slots." ) - # num_minibatches > 1 runs multiple opt.steps (each bumps - # trainer_version) between weight syncs; warn if the staleness window - # cannot keep same-step groups selectable across those bumps. - warn_if_staleness_window_below_minibatches(cfg) - self._sampler = StalenessSampler(cfg.max_weight_staleness_versions) # ── asyncio state ────────────────────────────────────────────────── # Gate: cleared during _sync_weights, set when generation may proceed self._rollout_permitted: asyncio.Event = asyncio.Event() self._rollout_permitted.set() + # Set only after _rollout_pump exhausts its configured epochs and all + # dispatched tasks finish successfully. Rollout failures propagate + # through run() instead of being reported as normal exhaustion. + self._rollout_exhausted: asyncio.Event = asyncio.Event() + # Count of in-flight generate_and_push calls self._inflight_rollouts: int = 0 - # Rollout batch counter — pre-incremented before each dispatch, so start - # at -1 to allow the first batch through the strict_on_policy gate. - self._max_rollout_version: int = -1 + # The rollout batch counter (dispatch index) is now owned by the sampler + # (``BaseSampler._dispatch_index``), which gates admission and stamps + # target_step — see ``_rollout_pump`` / ``PromptGroupSampler.admit``. # Active rollout tasks used by downstream synchronization/drain paths. # TaskGroup remains responsible for task ownership and cancellation. @@ -354,29 +296,24 @@ def __init__( self._trainer_version: int = 0 self._train_steps: int = 0 + self._step_log_dict: dict[str, list] = { + "rewards": [], + "masked_advantages": [], + "sequence_lengths": [], + } + # Completed prompt-list passes; only advances when # cfg.max_num_epochs is set (see _rollout_pump). self._current_epoch: int = 0 - self._claimed_meta: KVBatchMeta | None = None - self._step_consumed_sample_ids: list[str] = [] print( - f"SingleControllerActor: staleness_cap=" - f"{cfg.max_weight_staleness_versions} " + f"SingleControllerActor: sampler={cfg.sampler.name} " f"buffer={cfg.max_buffered_rollouts} " f"inflight={cfg.max_inflight_prompts} " f"transport={cfg.refit_cfg.impl}", flush=True, ) - def _timed(self, label: str) -> Any: - """Time a code block with the SC Timer when a logger is attached. - - No-op when no logger is configured. Use as - ``with self._timed("phase"): ...``. - """ - return self._timer.time(label) if self._timer is not None else nullcontext() - # ── public API ───────────────────────────────────────────────────────── async def run(self) -> dict[str, Any]: @@ -436,9 +373,10 @@ async def _call_dp(self, method_name: str, **kwargs) -> Any: async def _rollout_pump(self) -> None: """Continuously dispatch rollout tasks until cancellation. - Per batch (over_sampling=False): - 0. Wait while _max_rollout_version >= trainer_version + max_staleness, - then claim the next step by incrementing _max_rollout_version. + Per batch: + 0. await sampler.admit(...) — blocks until the batch may dispatch and + returns the target_step stamp (the staleness gate lives in the + sampler, not here). Per prompt: 1. Acquire _buffer_capacity slot (backpressure) @@ -450,9 +388,7 @@ async def _rollout_pump(self) -> None: 5. Decrement _inflight_rollouts """ sem = asyncio.Semaphore(self._cfg.max_inflight_prompts) - over_sampling = self._cfg.over_sampling - max_staleness = self._cfg.max_weight_staleness_versions - force_in_order = self._cfg.force_in_order + self._rollout_exhausted.clear() print("rollout_pump: starting", flush=True) async def _dispatch_one_prompt( @@ -496,17 +432,15 @@ def _release_permits_if_task_not_started( async with asyncio.TaskGroup() as rollout_tasks: while max_epochs is None or self._current_epoch < max_epochs: for prompt_batch in self._dataloader: - # over_sampling=False: batch-level gate on max_rollout_version. - if not over_sampling: - while ( - self._max_rollout_version - >= self._trainer_version + max_staleness - ): - await asyncio.sleep(0.005) - self._max_rollout_version += 1 - - # target_step = batch dispatch index when force_in_order is on. - target_step = self._max_rollout_version if force_in_order else None + # The sampler owns admission: it blocks until this batch may + # dispatch (per its own staleness algorithm) and returns the + # target_step to stamp (None when the policy doesn't stamp). + # Keeping the gate here means the rollout pump follows + # whichever sampler is selected without a second copy of the + # gating logic to keep in sync. + target_step = await self._sampler.admit( + trainer_version_fn=lambda: self._trainer_version + ) for prompt_idx in range(prompt_batch.size): prompt: DatumSpec = { # type: ignore @@ -538,30 +472,26 @@ def _release_permits_if_task_not_started( self._current_epoch += 1 + self._rollout_exhausted.set() print(f"rollout_pump: completed {self._current_epoch} epoch(s)", flush=True) async def _train_pump(self) -> None: """Per-prompt-group streaming train loop. Per step: - - Lazy ``begin_train_step`` on first ready group. - - Per ready group: optional logprob refresh - (``get_logprobs_from_meta`` / ``get_reference_policy_logprobs_from_meta``) → - ``_advantage_stage`` → ``train_microbatches_from_meta``. - - End-of-step: ``finish_train_step`` → - single ``clear_samples`` → ``_sync_weights``. - - **Concurrency contract:** - ``self._trainer`` is a driver-side ``TQPolicy`` object, not a Ray - actor. Trainer calls run via - ``asyncio.to_thread`` and are awaited sequentially, so the worker's - ``_train_step_state`` accumulators (``local_valid_seqs +=``, - ``mb_losses.append``, ``all_mb_metrics.append``), which are not - concurrency-safe, see exactly one call at a time. ``to_thread`` - keeps the fan-out + internal ``ray.get`` off the event loop so the - rollout pump makes progress during trainer calls; exceptions - surface at the corresponding ``await``. + 1. sampler.evict drops stale groups from the buffer and clears their TQ rows. + 2. sampler.select returns K prompt groups (or None) and drops them from the + buffer; DP rows survive so the trainer can read them. Already trainable — + buffer wrote training-shaped rows at rollout time. + 3. _advantage_stage(train_meta). + 4. trainer.train_microbatches_from_meta + finish_train_step. + 5. dp_client.clear_samples on consumed sample_ids; release _buffer_capacity + per dropped group, then sync. """ + target_groups = ( + self._cfg.target_groups_per_step or self._cfg.min_groups_per_batch + ) + # TODO: fix the prev_logprobs_required and reference_logprobs_required logic prev_logprobs_required = self._cfg.advantage_policy_logprobs_field is not None reference_logprobs_required = ( @@ -569,176 +499,187 @@ async def _train_pump(self) -> None: ) while self._train_steps < self._cfg.max_train_steps: - # __init__ coerces None → min_groups_per_batch (int); - # the assert narrows the Optional[int] type for pyrefly. - assert self._cfg.target_groups_per_step is not None - assert self._cfg.train_global_batch_size is not None - target_groups: int = self._cfg.target_groups_per_step - # Mini-batch aggregation: K groups per begin/finish cycle, where - # K * group_size == train_global_batch_size. Each - # cycle is one opt.step. Mirrors sync's gb_idx loop in the worker. - groups_per_minibatch = ( - self._cfg.train_global_batch_size // self._cfg.group_size - ) - num_minibatches = target_groups // groups_per_minibatch - rollout_exhausted = False - - for mb_idx in range(num_minibatches): - groups_dispatched = 0 - step_open = False - step_min_weight_version: int | None = None - - # No SC-side error handling: a mid-cycle worker failure - # propagates out of run(). The worker restores its own hooks - # on failure (see megatron_policy_worker); abort_train_step - # + a retry policy return with fault-tolerance support. - while groups_dispatched < groups_per_minibatch: - await asyncio.sleep(0) - await self._claim_available_meta() - evicted_meta = await self._evict_stale_claimed() - if evicted_meta is not None: - evicted_groups = count_groups( - evicted_meta, - group_size=self._cfg.group_size, + version_during_step = self._trainer_version + groups_dispatched = 0 + min_sample_version = None + step_open = False + + with self._timer.time("total_step_time"): + while groups_dispatched < target_groups: + # Wait for a selectable batch + with self._timer.time("exposed_generation"): + await asyncio.sleep(0) + + # Evict stale groups + evicted = await self._sampler.evict( + current_train_weight=self._trainer_version, ) - for _ in range(evicted_groups): - self._buffer_capacity.release() - - group_indices = None - if self._claimed_meta is not None and self._claimed_meta.size > 0: - group_indices = self._sampler.select_one_group( - self._claimed_meta, - trainer_version=self._trainer_version, - group_size=self._cfg.group_size, + if evicted: + print( + f" evicted {evicted} stale prompt group(s)", + flush=True, + ) + for _ in range(evicted): + self._buffer_capacity.release() + + # Select a batch + max_prompt_groups = target_groups - groups_dispatched + min_prompt_groups = min( + self._cfg.min_groups_per_batch, + max_prompt_groups, + ) + train_meta, num_groups = await self._sampler.select( + current_train_weight=self._trainer_version, + min_prompt_groups=min_prompt_groups, + max_prompt_groups=max_prompt_groups, ) - if group_indices is None: - await asyncio.sleep(0.005) - continue - - group_meta = self._claimed_meta.subset(group_indices) - self._claimed_meta = self._claimed_meta.drop(group_indices) - - if prev_logprobs_required or reference_logprobs_required: - with self._timed("policy_and_reference_logprobs"): - if prev_logprobs_required: - await asyncio.to_thread( - self._trainer.get_logprobs_from_meta, - group_meta, - ) - if reference_logprobs_required: - await asyncio.to_thread( - self._trainer.get_reference_policy_logprobs_from_meta, - group_meta, + # If no batch is selectable, sleep and retry + if train_meta is None: + if self._rollout_exhausted.is_set(): + buffered_groups = len(self._buffer) + if groups_dispatched == 0 and buffered_groups == 0: + print( + "train_pump: rollout exhausted and " + "buffer drained", + flush=True, + ) + return + raise RuntimeError( + "rollout exhausted before a complete training " + f"step was assembled: dispatched " + f"{groups_dispatched}/{target_groups} prompt " + f"groups with {buffered_groups} group(s) " + f"remaining in the buffer" ) + await asyncio.sleep(0.005) + continue + + # Release buffer capacity + for _ in range(num_groups): + self._buffer_capacity.release() - # Advantage stage — inline in the train pump, not a - # standalone pump task. - with self._timed("advantage_stage"): - group_meta = await self._advantage_stage(group_meta) + # Compute prev_logprobs / ref_logprobs + with self._timer.time("logprob_inference_prep"): + await asyncio.to_thread(self._trainer.prepare_for_lp_inference) + with self._timer.time("policy_and_reference_logprobs"): + if prev_logprobs_required: + await asyncio.to_thread( + self._trainer.get_logprobs_from_meta, train_meta + ) + if reference_logprobs_required: + await asyncio.to_thread( + self._trainer.get_reference_policy_logprobs_from_meta, + train_meta, + ) - if not step_open: - # gbs/mbs default to worker-side cfg when None; SC - # owns only loss_fn (stable across the whole run). + # Compute advantages + with self._timer.time("advantage_calculation"): + train_meta = await self._advantage_stage(train_meta) + + # Train + with self._timer.time("training_prep"): + await asyncio.to_thread(self._trainer.prepare_for_training) + with self._timer.time("policy_training"): + if not step_open: + await asyncio.to_thread( + self._trainer.begin_train_step, + self._loss_fn, + ) + step_open = True await asyncio.to_thread( - self._trainer.begin_train_step, - self._loss_fn, + self._trainer.train_microbatches_from_meta, + train_meta, ) - step_open = True - await asyncio.to_thread( - self._trainer.train_microbatches_from_meta, - group_meta, - ) - groups_dispatched += 1 - self._buffer_capacity.release() - self._step_consumed_sample_ids.extend(group_meta.sample_ids) - group_min_v = min_weight_version(group_meta) - if group_min_v is not None: - step_min_weight_version = ( - group_min_v - if step_min_weight_version is None - else min(step_min_weight_version, group_min_v) + if train_meta.sequence_lengths: + self._step_log_dict["sequence_lengths"].extend( + int(s) for s in train_meta.sequence_lengths ) - if not step_open: - # No groups consumed this mini-batch — either rollouts - # exhausted before any group arrived, or the outer - # loop broke without dispatching. Skip finish/cleanup. - print( - f"train_pump: rollout exhausted at mb {mb_idx} " - f"(no groups for this opt.step)", - flush=True, - ) - break - - # finish_train_step returns step metrics. trainer_version - # is driver-owned (workers don't emit it) and bumps after - # this call succeeds. The new value propagates to rollouts - # via _sync_weights at end of the outer step. Capture - # train_results for the logger emit below. - with self._timed("policy_training"): - train_results = await asyncio.to_thread( - self._trainer.finish_train_step + # Refresh min_sample_version + curr_min_sample_version = min( + t["weight_version"] + for t in train_meta.tags # type: ignore ) + if min_sample_version is not None: + min_sample_version = min( + min_sample_version, curr_min_sample_version + ) + else: + min_sample_version = curr_min_sample_version - # finish_train_step succeeded → opt.step ran on the worker, - # model weights are advanced. Bump the version immediately so - # SC's counter reflects worker state even if clear_samples - # below raises. - prev_trainer_version = self._trainer_version - self._trainer_version += 1 - - consumed_ids = list(self._step_consumed_sample_ids) - self._step_consumed_sample_ids = [] - with self._timed("clear_samples"): + # Remove consumed sample_ids from the buffer await self._call_dp( "clear_samples", - sample_ids=consumed_ids, + sample_ids=list(train_meta.sample_ids), partition_id=self._cfg.partition_id, ) - lag = ( - prev_trainer_version - step_min_weight_version - if step_min_weight_version is not None - else 0 - ) - print( - f"train step {self._train_steps + 1}/" - f"{self._cfg.max_train_steps} " - f"mb {mb_idx + 1}/{num_minibatches} " - f"trainer_v={self._trainer_version} lag={lag} " - f"batch_size={len(consumed_ids)}", - flush=True, - ) - # Log metrics - self._logger.log_metrics( - train_results, step=self._trainer_version, prefix="train" - ) - self._logger.log_metrics( - self._timer.get_timing_metrics(), - step=self._trainer_version, - prefix="timing/train", - step_finished=True, + groups_dispatched += num_groups + + with self._timer.time("policy_training"): + result = await asyncio.to_thread(self._trainer.finish_train_step) + + step_metrics = aggregate_step_metrics(result) + step_metrics.update( + reduce_advantage_pump_metrics(**self._step_log_dict) ) - self._timer.reset() + self._step_log_dict = {k: [] for k in self._step_log_dict} - if rollout_exhausted: - # Inner loop terminated early due to rollout exhaustion; - # the cycle still completed (step_open was True), so we - # finished and logged above. Exit the mini-batch loop now - # — no more groups are coming. - break + self._trainer_version += 1 + self._train_steps += 1 + with self._timer.time("weight_sync"): + await self._sync_weights() + + timing_metrics: dict[str, float] = self._timer.get_timing_metrics( + reduction_op="sum" + ) # type: ignore + + total_time = timing_metrics.get("total_step_time", 0.0) + total_num_gpus = int(ray.cluster_resources().get("GPU", 0)) + if ( + total_time > 0 + and total_num_gpus > 0 + and "global_valid_toks" in step_metrics + ): + timing_metrics["valid_tokens_per_sec_per_gpu"] = ( + step_metrics["global_valid_toks"] / total_time / total_num_gpus + ) - if rollout_exhausted: - # No more rollouts; stop the outer training loop entirely - # rather than running empty mini-batches. - break + print("\n⏱️ Timing:") + print(f" • Total step time: {total_time:.2f}s") + for k, v in sorted( + timing_metrics.items(), key=lambda item: item[1], reverse=True + ): + if k == "total_step_time": + continue + percent = (v / total_time * 100) if total_time > 0 else 0.0 + print(f" • {k}: {v:.2f}s ({percent:.1f}%)") + + # TODO: checkpointing (save_period/top-k metric_name, + # policy.save_checkpoint, dataloader state, TQReplayBuffer state). + # TODO: per-step train_data jsonl dump, vllm metrics logger, + # histogram log, rollout_metrics, seq_logprob_error_metrics, + # pretty-print "Training Results" block, print_performance_metrics. + print(f"step_metrics={step_metrics}", flush=True) + self._logger.log_metrics( + step_metrics, step=self._train_steps, prefix="train" + ) + self._logger.log_metrics( + timing_metrics, step=self._train_steps, prefix="timing/train" + ) + self._timer.reset() - # One sync per outer step covers all opt.steps in this iteration. - with self._timed("sync_weights"): - await self._sync_weights() - self._train_steps += 1 + # min sample version refers to the version each consumed sample was + # generated with; lag = training version - oldest sample version. + lag = version_during_step - min_sample_version # type: ignore + print( + f"train step {self._train_steps}/{self._cfg.max_train_steps} " + f"trainer_v={self._trainer_version} " + f"lag={lag} ", + flush=True, + ) async def _sync_weights(self) -> None: """Drain in-flight rollouts then synchronize weights. @@ -794,11 +735,11 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> KVBatchMeta: select_fields=self._advantage_input_fields(), ) - prompt_ids = _tensor_field(data, PROMPT_IDS_FIELD) - rewards = _squeeze_trailing_unit_dim(_tensor_field(data, REWARD_FIELD)).float() - token_mask = _tensor_field(data, TOKEN_MASK_FIELD).float() - sample_mask = _squeeze_trailing_unit_dim( - _tensor_field(data, SAMPLE_MASK_FIELD) + prompt_ids = tensor_field(data, PROMPT_IDS_FIELD) + rewards = squeeze_trailing_unit_dim(tensor_field(data, REWARD_FIELD)).float() + token_mask = tensor_field(data, TOKEN_MASK_FIELD).float() + sample_mask = squeeze_trailing_unit_dim( + tensor_field(data, SAMPLE_MASK_FIELD) ).float() mask = token_mask * sample_mask.unsqueeze(-1) @@ -806,18 +747,18 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> KVBatchMeta: "total_reward": rewards, } for field_name in self._cfg.advantage_repeated_batch_fields: - repeated_batch[field_name] = _squeeze_trailing_unit_dim( - _tensor_field(data, field_name) + repeated_batch[field_name] = squeeze_trailing_unit_dim( + tensor_field(data, field_name) ) kwargs: dict[str, torch.Tensor] = {} if self._cfg.advantage_policy_logprobs_field is not None: - kwargs["logprobs_policy"] = _tensor_field( + kwargs["logprobs_policy"] = tensor_field( data, self._cfg.advantage_policy_logprobs_field, ) if self._cfg.advantage_reference_logprobs_field is not None: - kwargs["logprobs_reference"] = _tensor_field( + kwargs["logprobs_reference"] = tensor_field( data, self._cfg.advantage_reference_logprobs_field, ) @@ -829,12 +770,17 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> KVBatchMeta: repeated_batch=repeated_batch, **kwargs, ) + response_advantages = torch.masked_select(advantages, mask.bool()) + self._step_log_dict["rewards"].append(rewards.detach().cpu()) + self._step_log_dict["masked_advantages"].append( + response_advantages.detach().cpu() + ) await self._call_dp( "put_samples", sample_ids=meta.sample_ids, partition_id=meta.partition_id, - fields=_fields_for_put( + fields=fields_for_put( meta, {ADVANTAGE_OUTPUT_FIELD: advantages}, ), @@ -843,36 +789,6 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> KVBatchMeta: # ── utility helpers ──────────────────────────────────────────────────── - async def _claim_available_meta(self) -> None: - """Claim currently-ready rows and append them to the local scheduler cache. - - TODO: replace this with a non-consuming metadata listing API. - ``claim_meta`` advances TQ's per-task cursor, so SC must keep a - local cache of claimed-but-not-yet-trained samples for now. - """ - batch_size = self._cfg.max_claim_groups * self._cfg.group_size - meta = await self._call_dp( - "claim_meta", - partition_id=self._cfg.partition_id, - task_name=self._cfg.consumer_task_name, - required_fields=self._claim_required_fields(), - batch_size=batch_size, - blocking=False, - timeout_s=0.0, - ) - if meta.size == 0: - return - if self._claimed_meta is None or self._claimed_meta.size == 0: - self._claimed_meta = meta - else: - self._claimed_meta = self._claimed_meta.concat(meta) - - def _claim_required_fields(self) -> list[str]: - fields = list(self._cfg.claim_required_fields) - if self._cfg.advantage_enabled: - fields.extend(self._advantage_input_fields()) - return list(dict.fromkeys(fields)) - def _advantage_input_fields(self) -> list[str]: fields = [ PROMPT_IDS_FIELD, @@ -886,70 +802,3 @@ def _advantage_input_fields(self) -> list[str]: if self._cfg.advantage_reference_logprobs_field is not None: fields.append(self._cfg.advantage_reference_logprobs_field) return list(dict.fromkeys(fields)) - - async def _evict_stale_claimed(self) -> KVBatchMeta | None: - if self._claimed_meta is None or self._claimed_meta.size == 0: - return None - indices = self._sampler.evictable_indices( - self._claimed_meta, - trainer_version=self._trainer_version, - group_size=self._cfg.group_size, - ) - if not indices: - return None - evicted_meta = self._claimed_meta.subset(indices) - print( - f" evicting {evicted_meta.size} stale samples from " - f"{count_groups(evicted_meta, group_size=self._cfg.group_size)} " - f"prompt group(s)", - flush=True, - ) - await self._call_dp( - "clear_samples", - sample_ids=evicted_meta.sample_ids, - partition_id=evicted_meta.partition_id, - ) - self._claimed_meta = self._claimed_meta.drop(indices) - return evicted_meta - - -def _tensor_field(data: TensorDict, field_name: str) -> torch.Tensor: - value = data[field_name] - if not isinstance(value, torch.Tensor): - raise TypeError( - f"advantage stage expected tensor field {field_name!r}; got {type(value)}" - ) - if value.is_nested: - return torch.nested.to_padded_tensor(value, padding=0) - return value - - -def _squeeze_trailing_unit_dim(value: torch.Tensor) -> torch.Tensor: - if value.dim() >= 2 and value.shape[-1] == 1: - return value.squeeze(-1) - return value - - -def _fields_for_put(meta: KVBatchMeta, fields: dict[str, torch.Tensor]) -> TensorDict: - packed: dict[str, torch.Tensor] = {} - if meta.sequence_lengths is None: - for field_name, value in fields.items(): - packed[field_name] = value.detach().contiguous() - # pyrefly: ignore[bad-argument-type] - return TensorDict(packed, batch_size=[meta.size]) - - lengths = torch.tensor(meta.sequence_lengths, dtype=torch.long) - for field_name, value in fields.items(): - if value.dim() >= 2 and value.shape[1] == int(lengths.max().item()): - rows = [ - value[i, : int(lengths[i].item())].detach().contiguous() - for i in range(meta.size) - ] - packed[field_name] = torch.nested.as_nested_tensor( - rows, - layout=torch.jagged, - ) - else: - packed[field_name] = value.detach().contiguous() - # pyrefly: ignore[bad-argument-type] - return TensorDict(packed, batch_size=[meta.size]) diff --git a/nemo_rl/algorithms/single_controller_utils/utils.py b/nemo_rl/algorithms/single_controller_utils/utils.py new file mode 100644 index 0000000000..4844898170 --- /dev/null +++ b/nemo_rl/algorithms/single_controller_utils/utils.py @@ -0,0 +1,189 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Helpers used by SingleControllerActor.""" + +from __future__ import annotations + +from typing import Any + +import numpy as np +import torch +from tensordict import TensorDict + +from nemo_rl.data_plane import KVBatchMeta + +# Reduction rules for all_mb_metrics. Mirror grpo.py / grpo_sync.py. +_MB_METRIC_MIN: frozenset[str] = frozenset( + {"probs_ratio_min", "probs_ratio_clamped_min"} +) +_MB_METRIC_MAX: frozenset[str] = frozenset( + {"probs_ratio_max", "probs_ratio_clamped_max"} +) +_MB_METRIC_MEAN: frozenset[str] = frozenset( + { + "lr", + "wd", + "reward", + "global_valid_seqs", + "global_valid_toks", + "mean_prompt_length", + } +) + + +def aggregate_step_metrics(train_result: dict[str, Any]) -> dict[str, Any]: + """Reduce per-microbatch metric lists into step-level scalars. + + Args: + train_result: Output of TQPolicy.finish_train_step. + + Returns: + Flat dict of step-level scalars ready for logging. + """ + metrics: dict[str, Any] = {} + loss = train_result.get("loss") + if isinstance(loss, torch.Tensor): + metrics["loss"] = loss.detach().mean().item() + elif loss is not None: + metrics["loss"] = float(loss) + grad_norm = train_result.get("grad_norm") + if isinstance(grad_norm, torch.Tensor): + metrics["grad_norm"] = grad_norm.detach().mean().item() + elif grad_norm is not None: + metrics["grad_norm"] = float(grad_norm) + if "total_flops" in train_result: + metrics["total_flops"] = float(train_result["total_flops"]) + if "num_ranks" in train_result: + metrics["num_ranks"] = int(train_result["num_ranks"]) + + # moe/mtp share the same reduction rules as all_mb_metrics in grpo.py. + mb: dict[str, list[Any]] = {} + if "moe_metrics" in train_result: + mb.update({f"moe/{k}": v for k, v in train_result["moe_metrics"].items()}) + if "mtp_metrics" in train_result: + mb.update({f"mtp/{k}": v for k, v in train_result["mtp_metrics"].items()}) + mb.update(train_result.get("all_mb_metrics", {})) + + for k, v in mb.items(): + if k in _MB_METRIC_MIN: + valid = [x for x in v if not np.isinf(x)] + metrics[k] = float(np.min(valid)) if valid else -1.0 + elif k in _MB_METRIC_MAX: + valid = [x for x in v if not np.isinf(x)] + metrics[k] = float(np.max(valid)) if valid else -1.0 + elif k in _MB_METRIC_MEAN: + metrics[k] = float(np.mean(v)) + else: + metrics[k] = float(np.sum(v)) + return metrics + + +def reduce_advantage_pump_metrics( + rewards: list[torch.Tensor], + masked_advantages: list[torch.Tensor], + sequence_lengths: list[int], +) -> dict[str, float]: + """Reduce per-step accumulators from _advantage_stage into step scalars. + + Args: + rewards: One tensor per advantage_stage call; each row a sample reward. + masked_advantages: Token-masked advantages, one tensor per call. + sequence_lengths: All input_lengths trained on this step. + + Returns: + Dict with reward, advantages/{mean,max,min}, total_num_tokens. + """ + out: dict[str, float] = {} + if rewards: + out["reward"] = float(torch.cat([r.flatten() for r in rewards]).mean()) + if masked_advantages: + cat = torch.cat([a.flatten() for a in masked_advantages]) + if cat.numel() > 0: + out["advantages/mean"] = float(cat.mean()) + out["advantages/max"] = float(cat.max()) + out["advantages/min"] = float(cat.min()) + else: + out["advantages/mean"] = 0.0 + out["advantages/max"] = 0.0 + out["advantages/min"] = 0.0 + if sequence_lengths: + out["total_num_tokens"] = float(sum(sequence_lengths)) + return out + + +def tensor_field(data: TensorDict, field_name: str) -> torch.Tensor: + """Read a tensor column from a TensorDict, depadding if nested. + + Args: + data: TensorDict returned by the data plane. + field_name: Column name to fetch. + + Returns: + Dense tensor (nested columns are padded with zeros). + """ + value = data[field_name] + if not isinstance(value, torch.Tensor): + raise TypeError(f"expected tensor field {field_name!r}; got {type(value)}") + if value.is_nested: + return torch.nested.to_padded_tensor(value, padding=0) + return value + + +def squeeze_trailing_unit_dim(value: torch.Tensor) -> torch.Tensor: + """Drop a trailing dim of size 1 if present. + + Args: + value: Input tensor. + + Returns: + Tensor without the trailing unit dim. + """ + if value.dim() >= 2 and value.shape[-1] == 1: + return value.squeeze(-1) + return value + + +def fields_for_put(meta: KVBatchMeta, fields: dict[str, torch.Tensor]) -> TensorDict: + """Pack tensors for DataPlane put, re-nesting jagged rows when needed. + + Args: + meta: Batch meta whose sequence_lengths drive the nesting. + fields: Field name to dense tensor. + + Returns: + TensorDict shaped for dp_client.put_samples. + """ + packed: dict[str, torch.Tensor] = {} + if meta.sequence_lengths is None: + for field_name, value in fields.items(): + packed[field_name] = value.detach().contiguous() + # pyrefly: ignore[bad-argument-type] + return TensorDict(packed, batch_size=[meta.size]) + + lengths = torch.tensor(meta.sequence_lengths, dtype=torch.long) + for field_name, value in fields.items(): + if value.dim() >= 2 and value.shape[1] == int(lengths.max().item()): + rows = [ + value[i, : int(lengths[i].item())].detach().contiguous() + for i in range(meta.size) + ] + packed[field_name] = torch.nested.as_nested_tensor( + rows, + layout=torch.jagged, + ) + else: + packed[field_name] = value.detach().contiguous() + # pyrefly: ignore[bad-argument-type] + return TensorDict(packed, batch_size=[meta.size]) diff --git a/nemo_rl/algorithms/staleness_sampler.py b/nemo_rl/algorithms/staleness_sampler.py deleted file mode 100644 index a9fbc18a9d..0000000000 --- a/nemo_rl/algorithms/staleness_sampler.py +++ /dev/null @@ -1,220 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Prompt-group batch selection strategies for SingleController metadata.""" - -from __future__ import annotations - -from dataclasses import dataclass -from typing import Optional - -from nemo_rl.data_plane import KVBatchMeta - - -@dataclass(frozen=True) -class PromptGroup: - """Indices and scheduling metadata for one prompt group.""" - - group_id: str - indices: list[int] - weight_version: int | None - committed: bool - expected_num_samples: int - - @property - def is_complete(self) -> bool: - return len(self.indices) == self.expected_num_samples - - -class StalenessSampler: - """Select complete prompt groups inside a version staleness window.""" - - def __init__(self, max_staleness_versions: int): - self.max_staleness_versions = max_staleness_versions - - def select_indices( - self, - meta: KVBatchMeta, - *, - trainer_version: int, - min_groups: int, - group_size: int, - ) -> Optional[list[int]]: - eligible: list[tuple[int, int, PromptGroup]] = [] - for group in _iter_groups(meta, group_size): - if not group.committed or not group.is_complete: - continue - if group.weight_version is None or group.weight_version > trainer_version: - continue - lag = trainer_version - group.weight_version - if lag > self.max_staleness_versions: - continue - eligible.append((lag, group.indices[0], group)) - - if len(eligible) < min_groups: - return None - - eligible.sort(key=lambda item: (item[0], item[1])) - groups = [item[2] for item in eligible[:min_groups]] - return _flatten_group_indices(groups) - - def select_one_group( - self, - meta: KVBatchMeta, - *, - trainer_version: int, - group_size: int, - ) -> Optional[list[int]]: - eligible: list[tuple[int, int, PromptGroup]] = [] - for group in _iter_groups(meta, group_size): - if not group.committed or not group.is_complete: - continue - if group.weight_version is None or group.weight_version > trainer_version: - continue - lag = trainer_version - group.weight_version - if lag > self.max_staleness_versions: - continue - eligible.append((lag, group.indices[0], group)) - - if not eligible: - return None - - eligible.sort(key=lambda item: (item[0], item[1])) - return _flatten_group_indices([eligible[0][2]]) - - def evictable_indices( - self, - meta: KVBatchMeta, - *, - trainer_version: int, - group_size: int, - ) -> list[int]: - groups = [] - for group in _iter_groups(meta, group_size): - if group.weight_version is None or not group.is_complete: - continue - lag = trainer_version - group.weight_version - if lag > self.max_staleness_versions: - groups.append(group) - return _flatten_group_indices(groups) - - -def count_groups( - meta: KVBatchMeta, - *, - group_size: int, -) -> int: - """Count complete prompt groups represented by ``meta``.""" - return sum(1 for group in _iter_groups(meta, group_size) if group.is_complete) - - -def incomplete_group_indices( - meta: KVBatchMeta, - *, - group_size: int, -) -> list[int]: - """Flat indices of groups that are uncommitted or short of expected samples. - - Such groups are invisible to both selection (``select_*`` skip them) and - eviction (``evictable_indices`` skips them), so once no more samples can - arrive they are permanently stuck. Callers drop them after rollout - shutdown so exhaustion checks can fire. - """ - return _flatten_group_indices( - [ - group - for group in _iter_groups(meta, group_size) - if not group.committed or not group.is_complete - ] - ) - - -def min_weight_version(meta: KVBatchMeta) -> int | None: - """Smallest ``weight_version`` across per-sample tags, or None if absent.""" - versions = [ - v for v in (_weight_version(tag) for tag in meta.tags or []) if v is not None - ] - return min(versions) if versions else None - - -def _iter_groups( - meta: KVBatchMeta, - group_size: int, -) -> list[PromptGroup]: - tags = meta.tags or [{} for _ in meta.sample_ids] - grouped: dict[str, list[int]] = {} - first_tag: dict[str, dict] = {} - - for idx, sample_id in enumerate(meta.sample_ids): - tag = tags[idx] if idx < len(tags) else {} - group_id = str(tag.get("group_id") or _group_id_from_sample_id(sample_id)) - grouped.setdefault(group_id, []).append(idx) - first_tag.setdefault(group_id, tag) - - groups: list[PromptGroup] = [] - for group_id, indices in grouped.items(): - tag = first_tag[group_id] - expected = _as_int( - tag.get( - "expected_num_samples", - tag.get( - "expected_num_keys", - tag.get("generations_per_prompt", group_size), - ), - ) - ) - groups.append( - PromptGroup( - group_id=group_id, - indices=indices, - weight_version=_weight_version(tag), - committed=_as_bool(tag.get("committed", True)), - expected_num_samples=expected or group_size, - ) - ) - groups.sort(key=lambda group: group.indices[0]) - return groups - - -def _flatten_group_indices(groups: list[PromptGroup]) -> list[int]: - return [idx for group in groups for idx in group.indices] - - -def _group_id_from_sample_id(sample_id: str) -> str: - prefix, sep, suffix = sample_id.rpartition("_g") - if sep and suffix.isdigit(): - return prefix - return sample_id - - -def _weight_version(tag: dict) -> int | None: - value = tag.get("weight_version", tag.get("version")) - return _as_int(value) - - -def _as_int(value) -> int | None: - if value is None: - return None - try: - return int(value) - except (TypeError, ValueError): - return None - - -def _as_bool(value) -> bool: - if isinstance(value, bool): - return value - if isinstance(value, str): - return value.lower() in {"1", "true", "yes"} - return bool(value) diff --git a/pyrefly.toml b/pyrefly.toml index 9a4237d347..fa5696e661 100644 --- a/pyrefly.toml +++ b/pyrefly.toml @@ -48,6 +48,7 @@ project-includes = [ "nemo_rl/algorithms/async_utils/__init__.py", "nemo_rl/algorithms/async_utils/interfaces.py", "nemo_rl/algorithms/async_utils/replay_buffer.py", + "nemo_rl/algorithms/async_utils/staleness_sampler.py", "nemo_rl/algorithms/async_utils/trajectory_collector.py", "nemo_rl/algorithms/logits_sampling_utils.py", "nemo_rl/algorithms/loss/__init__.py", @@ -56,7 +57,7 @@ project-includes = [ "nemo_rl/algorithms/opd.py", "nemo_rl/algorithms/reward_functions.py", "nemo_rl/algorithms/single_controller.py", - "nemo_rl/algorithms/staleness_sampler.py", + "nemo_rl/algorithms/single_controller_utils/utils.py", "nemo_rl/algorithms/utils.py", "nemo_rl/algorithms/x_token/__init__.py", "nemo_rl/algorithms/x_token/utils.py", diff --git a/tests/unit/single_controller/_dp_fakes.py b/tests/unit/single_controller/_dp_fakes.py new file mode 100644 index 0000000000..1f26ea6d77 --- /dev/null +++ b/tests/unit/single_controller/_dp_fakes.py @@ -0,0 +1,158 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared TQ data-plane fakes for single_controller tests.""" + +from __future__ import annotations + +from typing import Any + +import ray +import torch +from tensordict import TensorDict + +from nemo_rl.data_plane.adapters.noop import NoOpDataPlaneClient + +_PARTITION_ID = "rollout_data" +# TQReplayBuffer.commit writes these training-row fields per prompt; +# _advantage_stage writes ``advantages`` back on top. +_BULK_FIELDS = [ + "input_ids", + "input_lengths", + "generation_logprobs", + "token_mask", + "sample_mask", + "prompt_ids_for_adv", + "total_reward", +] +_ADV_FIELD = "advantages" + + +@ray.remote(num_cpus=0) # pragma: no cover +class _TQActor: + """Ray-wrapped NoOpDataPlaneClient for cross-process TQ inspection.""" + + def __init__( + self, + partition_id: str, + fields: list[str], + num_samples: int, + consumer_tasks: list[str], + ) -> None: + self._client = NoOpDataPlaneClient() + self._client.register_partition( + partition_id=partition_id, + fields=list(fields), + num_samples=int(num_samples), + consumer_tasks=list(consumer_tasks), + ) + + def put_samples( + self, + sample_ids: list[str], + partition_id: str, + fields: TensorDict | None = None, + tags: list[dict[str, Any]] | None = None, + ) -> Any: + return self._client.put_samples( + sample_ids=sample_ids, + partition_id=partition_id, + fields=fields, + tags=tags, + ) + + def get_samples( + self, + sample_ids: list[str], + partition_id: str, + select_fields: list[str], + ) -> TensorDict: + return self._client.get_samples( + sample_ids=sample_ids, + partition_id=partition_id, + select_fields=list(select_fields), + ) + + def clear_samples(self, sample_ids: list[str], partition_id: str) -> Any: + return self._client.clear_samples( + sample_ids=sample_ids, partition_id=partition_id + ) + + def claim_meta(self, **kwargs: Any) -> Any: + return self._client.claim_meta(**kwargs) + + def get_tags( + self, partition_id: str, sample_ids: list[str] + ) -> list[dict[str, Any]]: + rec = self._client._partitions[partition_id] + return [dict(rec.tags.get(sid, {})) for sid in sample_ids] + + def peek_count(self, partition_id: str) -> int: + return len(self._client._partitions[partition_id].rows) + + +class _SyncDPAdapter: + """Sync DataPlaneClient over a Ray actor handle. Pads nested tensors before transport.""" + + def __init__(self, handle: Any) -> None: + self._handle = handle + + def put_samples( + self, + sample_ids: list[str], + partition_id: str, + fields: TensorDict | None = None, + tags: list[dict[str, Any]] | None = None, + ) -> Any: + if fields is not None: + fields = self._padded(fields) + return ray.get( + self._handle.put_samples.remote( + sample_ids=sample_ids, + partition_id=partition_id, + fields=fields, + tags=tags, + ) + ) + + def get_samples( + self, + sample_ids: list[str], + partition_id: str, + select_fields: list[str], + ) -> TensorDict: + return ray.get( + self._handle.get_samples.remote( + sample_ids=sample_ids, + partition_id=partition_id, + select_fields=select_fields, + ) + ) + + def clear_samples(self, sample_ids: list[str], partition_id: str) -> Any: + return ray.get( + self._handle.clear_samples.remote( + sample_ids=sample_ids, partition_id=partition_id + ) + ) + + @staticmethod + def _padded(td: TensorDict) -> TensorDict: + out: dict[str, torch.Tensor] = {} + for k in td.keys(): + v = td.get(k) + if isinstance(v, torch.Tensor) and v.is_nested: + v = torch.nested.to_padded_tensor(v, padding=0) + out[k] = v + return TensorDict(out, batch_size=td.batch_size) diff --git a/tests/unit/single_controller/test_rollout_pump.py b/tests/unit/single_controller/test_rollout_pump.py index b2e82685f2..4a2ba61cce 100644 --- a/tests/unit/single_controller/test_rollout_pump.py +++ b/tests/unit/single_controller/test_rollout_pump.py @@ -24,14 +24,18 @@ import pytest import ray import torch -from tensordict import TensorDict from nemo_rl.algorithms.async_utils.replay_buffer import TQReplayBuffer +from nemo_rl.algorithms.async_utils.staleness_sampler import ( + InOrderSampler, + WeightFifoSampler, + WindowedSampler, + WindowedSamplerConfig, +) from nemo_rl.algorithms.single_controller import ( SingleControllerActor, SingleControllerConfig, ) -from nemo_rl.data_plane.adapters.noop import NoOpDataPlaneClient from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.experience.rollout_manager import RolloutManager @@ -46,123 +50,25 @@ rollout_cluster, # noqa: F401 rollout_tokenizer, # noqa: F401 ) - -_PARTITION_ID = "rollout_data" -# TQReplayBuffer.commit tensorizes each PromptGroupRecord and writes -# ``generations_per_prompt`` training rows directly to TQ. -_BULK_FIELDS = [ - "input_ids", - "input_lengths", - "generation_logprobs", - "token_mask", - "sample_mask", - "prompt_ids_for_adv", - "total_reward", -] - - -@ray.remote(num_cpus=0) -class _TQActor: - """Ray-wrapped NoOpDataPlaneClient for cross-process TQ inspection.""" - - def __init__( - self, - partition_id: str, - fields: list[str], - num_samples: int, - consumer_tasks: list[str], - ) -> None: - self._client = NoOpDataPlaneClient() - self._client.register_partition( - partition_id=partition_id, - fields=list(fields), - num_samples=int(num_samples), - consumer_tasks=list(consumer_tasks), - ) - - def put_samples( - self, - sample_ids: list[str], - partition_id: str, - fields: TensorDict | None = None, - tags: list[dict[str, Any]] | None = None, - ) -> Any: - return self._client.put_samples( - sample_ids=sample_ids, - partition_id=partition_id, - fields=fields, - tags=tags, - ) - - def claim_meta(self, **kwargs: Any) -> Any: - return self._client.claim_meta(**kwargs) - - def get_samples( - self, - sample_ids: list[str], - partition_id: str, - select_fields: list[str], - ) -> TensorDict: - return self._client.get_samples( - sample_ids=sample_ids, - partition_id=partition_id, - select_fields=list(select_fields), - ) - - def get_tags( - self, partition_id: str, sample_ids: list[str] - ) -> list[dict[str, Any]]: - rec = self._client._partitions[partition_id] - return [dict(rec.tags.get(sid, {})) for sid in sample_ids] - - def peek_count(self, partition_id: str) -> int: - return len(self._client._partitions[partition_id].rows) - - -class _SyncDPAdapter: - """Sync DataPlaneClient over a Ray actor handle. Pads nested tensors before transport.""" - - def __init__(self, handle: Any) -> None: - self._handle = handle - - def put_samples( - self, - sample_ids: list[str], - partition_id: str, - fields: TensorDict | None = None, - tags: list[dict[str, Any]] | None = None, - ) -> Any: - if fields is not None: - fields = self._padded(fields) - return ray.get( - self._handle.put_samples.remote( - sample_ids=sample_ids, - partition_id=partition_id, - fields=fields, - tags=tags, - ) - ) - - @staticmethod - def _padded(td: TensorDict) -> TensorDict: - out: dict[str, torch.Tensor] = {} - for k in td.keys(): - v = td.get(k) - if isinstance(v, torch.Tensor) and v.is_nested: - v = torch.nested.to_padded_tensor(v, padding=0) - out[k] = v - return TensorDict(out, batch_size=td.batch_size) +from tests.unit.single_controller._dp_fakes import ( + _BULK_FIELDS, + _PARTITION_ID, + _SyncDPAdapter, + _TQActor, +) @pytest.mark.parametrize( - ("force_in_order", "expected_target_steps"), + ("make_sampler", "expected_target_steps"), [ - (False, [None, None]), - (True, [0, 1]), + # weight_fifo gates but does not stamp target_step. + (lambda buf: WeightFifoSampler(buf, max_staleness_versions=1), [None, None]), + # in_order stamps the dispatch index as target_step. + (lambda buf: InOrderSampler(buf, max_lookahead_versions=1), [0, 1]), ], ) def test_rollout_pump_stamps_target_steps( - force_in_order: bool, + make_sampler, expected_target_steps: list[int | None], ) -> None: class _RecordingBuffer: @@ -187,29 +93,30 @@ async def generate_and_push( ctrl = object.__new__(controller_cls) ctrl._cfg = SimpleNamespace( max_inflight_prompts=2, - over_sampling=False, - max_weight_staleness_versions=1, - force_in_order=force_in_order, diagnostics=False, max_num_epochs=1, ) ctrl._rollout_manager = _RecordingRolloutManager(buffer) + # The sampler owns admission + target_step stamping (the dispatch counter + # lives on the sampler, not the actor). + ctrl._sampler = make_sampler(buffer) prompt_batch = BatchedDataDict( {"message_log": [[{"role": "user", "content": "prompt"}]]} ) ctrl._dataloader = [prompt_batch, prompt_batch] ctrl._rollout_permitted = asyncio.Event() ctrl._rollout_permitted.set() + ctrl._rollout_exhausted = asyncio.Event() ctrl._buffer_capacity = asyncio.Semaphore(2) ctrl._inflight_rollouts = 0 ctrl._dispatched_rollouts = set() - ctrl._max_rollout_version = -1 ctrl._trainer_version = 0 ctrl._current_epoch = 0 asyncio.run(ctrl._rollout_pump()) assert buffer.target_step_list == expected_target_steps + assert ctrl._rollout_exhausted.is_set() def test_rollout_pump_failure_cancels_sibling_and_releases_capacity() -> None: @@ -244,13 +151,12 @@ async def _main() -> None: ctrl = object.__new__(controller_cls) ctrl._cfg = SimpleNamespace( max_inflight_prompts=2, - over_sampling=True, - max_weight_staleness_versions=1, - force_in_order=False, diagnostics=False, max_num_epochs=1, ) ctrl._rollout_manager = manager + # Over-sampled windowed policy: admit never gates (buffer unused here). + ctrl._sampler = WindowedSampler(None, max_staleness_versions=1) ctrl._dataloader = [ BatchedDataDict( { @@ -263,10 +169,10 @@ async def _main() -> None: ] ctrl._rollout_permitted = asyncio.Event() ctrl._rollout_permitted.set() + ctrl._rollout_exhausted = asyncio.Event() ctrl._buffer_capacity = asyncio.Semaphore(2) ctrl._inflight_rollouts = 0 ctrl._dispatched_rollouts = set() - ctrl._max_rollout_version = -1 ctrl._trainer_version = 0 ctrl._current_epoch = 0 @@ -278,6 +184,7 @@ async def _main() -> None: assert ctrl._inflight_rollouts == 0 assert ctrl._buffer_capacity._value == 2 assert ctrl._dispatched_rollouts == set() + assert not ctrl._rollout_exhausted.is_set() asyncio.run(_main()) @@ -323,22 +230,21 @@ async def _main() -> None: ctrl = object.__new__(controller_cls) ctrl._cfg = SimpleNamespace( max_inflight_prompts=1, - over_sampling=True, - max_weight_staleness_versions=1, - force_in_order=False, diagnostics=False, max_num_epochs=1, ) ctrl._rollout_manager = _NeverCalledRolloutManager() + # Over-sampled windowed policy: admit never gates (buffer unused here). + ctrl._sampler = WindowedSampler(None, max_staleness_versions=1) ctrl._dataloader = [ BatchedDataDict({"message_log": [[{"role": "user", "content": "prompt"}]]}) ] ctrl._rollout_permitted = asyncio.Event() ctrl._rollout_permitted.set() + ctrl._rollout_exhausted = asyncio.Event() ctrl._buffer_capacity = real_semaphore(1) ctrl._inflight_rollouts = 0 ctrl._dispatched_rollouts = set() - ctrl._max_rollout_version = -1 ctrl._trainer_version = 0 ctrl._current_epoch = 0 @@ -349,6 +255,7 @@ async def _main() -> None: assert created_semaphores[0]._value == 1 assert ctrl._inflight_rollouts == 0 assert ctrl._dispatched_rollouts == set() + assert ctrl._rollout_exhausted.is_set() asyncio.run(_main()) @@ -359,14 +266,14 @@ def test_rollout_pump_writes_expected_tq_data( single_multi_step_calculator_input_sample, # noqa: F811 tmp_path, ): - """SC._rollout_pump writes max_rollout_prompts * num_generations rows to TQ with the expected fields and tags.""" + """SC._rollout_pump writes num_prompts * num_generations rows to TQ with the expected fields and tags.""" vllm_generation, tokenizer, task_to_env, _, _ = multi_step_setup_vllm_async input_sample = single_multi_step_calculator_input_sample num_generations = 2 - max_rollout_prompts = 2 + num_prompts = 2 # TQReplayBuffer.commit writes ``num_generations`` training rows per prompt. - expected_samples = max_rollout_prompts * num_generations + expected_samples = num_prompts * num_generations max_seq_len = 1024 max_rollout_turns = input_sample["extra_env_info"]["max_steps"] + 1 @@ -379,15 +286,13 @@ def test_rollout_pump_writes_expected_tq_data( dp_adapter = _SyncDPAdapter(tq_actor) cfg = SingleControllerConfig.model_construct( - batch_selection_strategy="staleness_window", - max_weight_staleness_versions=1, + sampler=WindowedSamplerConfig(max_staleness_versions=1), min_groups_per_batch=1, group_size=num_generations, - max_inflight_prompts=max_rollout_prompts, - max_buffered_rollouts=max_rollout_prompts, + max_inflight_prompts=num_prompts, + max_buffered_rollouts=num_prompts, max_train_steps=1, max_num_epochs=None, - over_sampling=True, partition_id=_PARTITION_ID, logger={ "log_dir": str(tmp_path / "logs"), @@ -400,7 +305,7 @@ def test_rollout_pump_writes_expected_tq_data( ) # Wrap each value in a single-element list so size==1 and v[0] returns the original field. batched_sample = BatchedDataDict({k: [v] for k, v in input_sample.items()}) - dataloader = [batched_sample] * max_rollout_prompts + dataloader = [batched_sample] * num_prompts tq_buffer = TQReplayBuffer( dp_adapter, @@ -419,15 +324,15 @@ def test_rollout_pump_writes_expected_tq_data( ) ctrl = SingleControllerActor.remote( cfg=cfg, - prompts=[], dp_client_handle=dp_adapter, gen_handle=vllm_generation, trainer_handle=object(), weight_synchronizer=object(), loss_fn=None, - advantage_estimator=None, rollout_manager=rollout_manager, + advantage_estimator=None, dataloader=dataloader, + tq_buffer=tq_buffer, ) vllm_generation.prepare_for_generation() @@ -468,7 +373,7 @@ def test_rollout_pump_writes_expected_tq_data( head, _, tail = sid.rpartition("_g") assert head and tail.isdigit(), f"unexpected sample_id: {sid}" group_ids.add(head) - assert len(group_ids) == max_rollout_prompts + assert len(group_ids) == num_prompts bulk = ray.get( tq_actor.get_samples.remote( diff --git a/tests/unit/single_controller/test_sampler_interface.py b/tests/unit/single_controller/test_sampler_interface.py new file mode 100644 index 0000000000..8e0ef3e20c --- /dev/null +++ b/tests/unit/single_controller/test_sampler_interface.py @@ -0,0 +1,308 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Contract tests for the split PromptGroupSampler policies + factory. + +CPU-only; mirrors the FakeBuffer surface used by test_staleness_sampler.py. +Covers what the split buys over the monolithic sampler: per-policy admission, +InOrderSampler's target_step-keyed evict (evict/select agree by construction), +and FQN loading of an out-of-repo sampler. +""" + +from __future__ import annotations + +import asyncio + +import pytest + +from nemo_rl.algorithms.async_utils.staleness_sampler import ( + InOrderSampler, + InOrderSamplerConfig, + PromptGroupSampler, + WeightFifoSampler, + WindowedSampler, + WindowedSamplerConfig, + create_sampler, +) +from nemo_rl.data_plane import KVBatchMeta + + +class FakeBuffer: + """Minimal TQReplayBuffer surface the samplers read/mutate.""" + + def __init__(self, partition_id: str = "rollout_data") -> None: + self._partition_id = partition_id + self.meta_list: list[KVBatchMeta | None] = [] + self.start_weight_list: list[int] = [] + self.end_weight_list: list[int] = [] + self.target_step_list: list[int | None] = [] + self.ready_list: list[bool] = [] + self.remove_calls: list[tuple[list[int], bool]] = [] + + def add( + self, + group_id: str, + weight: int, + *, + ready: bool = True, + target_step: int | None = None, + ) -> None: + meta = KVBatchMeta( + partition_id=self._partition_id, + task_name=None, + sample_ids=[f"{group_id}_g0"], + tags=[{"weight_version": weight, "group_id": group_id}], + ) + self.meta_list.append(meta if ready else None) + self.start_weight_list.append(weight) + self.end_weight_list.append(weight) + self.target_step_list.append(target_step) + self.ready_list.append(ready) + + async def remove(self, idxs: list[int], remove_in_dp: bool) -> int: + self.remove_calls.append((list(idxs), remove_in_dp)) + for i in sorted(idxs, reverse=True): + del self.meta_list[i] + del self.start_weight_list[i] + del self.end_weight_list[i] + del self.target_step_list[i] + del self.ready_list[i] + return len(idxs) + + +def _run(coro): + return asyncio.run(coro) + + +class TestBuiltinsImplementInterface: + @pytest.mark.parametrize( + "sampler", + [ + WindowedSampler(FakeBuffer(), max_staleness_versions=1), + WeightFifoSampler(FakeBuffer(), max_staleness_versions=1), + InOrderSampler(FakeBuffer(), max_lookahead_versions=1), + ], + ) + def test_isinstance_protocol(self, sampler): + assert isinstance(sampler, PromptGroupSampler) + + +class TestAdmission: + def test_windowed_never_gates_and_never_stamps(self): + s = WindowedSampler(FakeBuffer(), max_staleness_versions=2) + # trainer stuck at 0, but over-sampled admission returns immediately. + assert _run(s.admit(trainer_version_fn=lambda: 0)) is None + assert _run(s.admit(trainer_version_fn=lambda: 0)) is None + + def test_in_order_stamps_monotonic_dispatch_index(self): + s = InOrderSampler(FakeBuffer(), max_lookahead_versions=5) + assert _run(s.admit(trainer_version_fn=lambda: 10)) == 0 + assert _run(s.admit(trainer_version_fn=lambda: 10)) == 1 + + def test_weight_fifo_gates_on_lookahead_and_does_not_stamp(self): + # dispatch_index starts at -1; window 0 => admits exactly one batch + # ahead of the trainer, then blocks. Assert the second admit would block + # by giving it a trainer_version that keeps the gate closed. + s = WeightFifoSampler(FakeBuffer(), max_staleness_versions=0) + assert _run(s.admit(trainer_version_fn=lambda: 0)) is None # -1 -> 0 + # Now dispatch_index=0, trainer=0, window=0 -> 0 >= 0 blocks forever. + with pytest.raises(asyncio.TimeoutError): + _run(asyncio.wait_for(s.admit(trainer_version_fn=lambda: 0), timeout=0.05)) + + +class TestInOrderEvictMatchesSelect: + """The bug the split fixes: monolithic evict keyed on weight could drop a + slot whose target_step was still upcoming. InOrderSampler keys evict on + target_step, so it never drops a slot select would later match.""" + + def test_future_target_not_evicted_even_if_weight_out_of_window(self): + buf = FakeBuffer() + # weight far below the window, but target_step is still upcoming. + buf.add("g", weight=0, ready=True, target_step=2) + s = InOrderSampler(buf, max_lookahead_versions=1) + removed = _run(s.evict(current_train_weight=2)) + assert removed == 0 # target_step 2 == current, not past -> kept + assert len(buf.target_step_list) == 1 + + def test_past_target_ready_slot_is_evicted(self): + buf = FakeBuffer() + buf.add("g", weight=0, ready=True, target_step=1) + s = InOrderSampler(buf, max_lookahead_versions=1) + removed = _run(s.evict(current_train_weight=3)) # target 1 < 3 -> stale + assert removed == 1 + + def test_unready_slot_is_never_evicted(self): + buf = FakeBuffer() + buf.add("g", weight=0, ready=False, target_step=1) + s = InOrderSampler(buf, max_lookahead_versions=1) + # past target, but unready -> skipped to avoid the commit race. + assert _run(s.evict(current_train_weight=5)) == 0 + + +class TestFactory: + def test_windowed_config_builds_windowed(self): + s = create_sampler( + FakeBuffer(), WindowedSamplerConfig(max_staleness_versions=3) + ) + assert isinstance(s, WindowedSampler) + assert s.max_staleness_versions == 3 + + def test_in_order_config_builds_in_order(self): + s = create_sampler(FakeBuffer(), InOrderSamplerConfig(max_lookahead_versions=2)) + assert isinstance(s, InOrderSampler) + assert s.max_lookahead_versions == 2 + + def test_weight_fifo_config_builds_weight_fifo(self): + from nemo_rl.algorithms.async_utils.staleness_sampler import ( + WeightFifoSamplerConfig, + ) + + s = create_sampler( + FakeBuffer(), WeightFifoSamplerConfig(max_staleness_versions=4) + ) + assert isinstance(s, WeightFifoSampler) + assert s.max_staleness_versions == 4 + + +class TestCustomFqnSampler: + def test_custom_target_loads_out_of_repo_sampler(self): + # A user sampler defined anywhere importable; here, this test module. + from nemo_rl.algorithms.async_utils.staleness_sampler import ( + CustomSamplerConfig, + ) + + s = create_sampler( + FakeBuffer(), + CustomSamplerConfig( + target=f"{__name__}:EchoSampler", max_lookahead_versions=1 + ), + ) + assert isinstance(s, EchoSampler) + assert isinstance(s, PromptGroupSampler) + assert s.max_lookahead_versions == 1 + + +class TestWindowedSelect: + def test_selects_ready_groups_in_window(self): + buf = FakeBuffer() + buf.add("a", weight=3) + buf.add("b", weight=5) # current + buf.add("c", weight=1) # below window (5-2=3) + s = WindowedSampler(buf, max_staleness_versions=2) + meta, n = _run( + s.select(current_train_weight=5, min_prompt_groups=1, max_prompt_groups=8) + ) + assert n == 2 # a(3) and b(5); c(1) excluded + assert len(buf.start_weight_list) == 1 # only c remains + + def test_below_min_returns_none(self): + buf = FakeBuffer() + buf.add("a", weight=5) + s = WindowedSampler(buf, max_staleness_versions=2) + assert _run( + s.select(current_train_weight=5, min_prompt_groups=2, max_prompt_groups=8) + ) == (None, 0) + + def test_unready_excluded(self): + buf = FakeBuffer() + buf.add("a", weight=5, ready=False) + s = WindowedSampler(buf, max_staleness_versions=2) + assert _run( + s.select(current_train_weight=5, min_prompt_groups=1, max_prompt_groups=8) + ) == (None, 0) + + def test_freshest_first_orders_by_lag(self): + buf = FakeBuffer() + buf.add("old", weight=1) + buf.add("new", weight=5) + s = WindowedSampler(buf, max_staleness_versions=10, sample_freshest_first=True) + meta, n = _run( + s.select(current_train_weight=5, min_prompt_groups=1, max_prompt_groups=1) + ) + # freshest (weight 5) picked first -> "old" (weight 1) remains. + assert n == 1 + assert buf.start_weight_list == [1] + + +class TestWeightFifoSelect: + def test_drains_oldest_in_window_weight_first(self): + buf = FakeBuffer() + buf.add("old1", weight=3) + buf.add("new", weight=5) + buf.add("old2", weight=3) + s = WeightFifoSampler(buf, max_staleness_versions=5) + meta, n = _run( + s.select(current_train_weight=5, min_prompt_groups=1, max_prompt_groups=8) + ) + assert n == 2 # both weight-3 groups; weight-5 waits its turn + assert buf.start_weight_list == [5] + + def test_waits_for_partial_oldest_batch(self): + buf = FakeBuffer() + buf.add("old", weight=3) + s = WeightFifoSampler(buf, max_staleness_versions=5) + # oldest weight has only 1 group but min is 2 -> wait (None), don't skip + # ahead to a newer weight. + assert _run( + s.select(current_train_weight=5, min_prompt_groups=2, max_prompt_groups=8) + ) == (None, 0) + + def test_empty_window_returns_none(self): + buf = FakeBuffer() + buf.add("future", weight=9) + s = WeightFifoSampler(buf, max_staleness_versions=2) + assert _run( + s.select(current_train_weight=5, min_prompt_groups=1, max_prompt_groups=8) + ) == (None, 0) + + +class TestInOrderSelect: + def test_matches_target_step_ignoring_weight_window(self): + buf = FakeBuffer() + # weight far outside any window, but target_step == trainer version. + buf.add("g", weight=100, target_step=5) + s = InOrderSampler(buf, max_lookahead_versions=1) + meta, n = _run( + s.select(current_train_weight=5, min_prompt_groups=1, max_prompt_groups=8) + ) + assert n == 1 + + def test_non_matching_target_not_selected(self): + buf = FakeBuffer() + buf.add("g", weight=5, target_step=6) + s = InOrderSampler(buf, max_lookahead_versions=1) + assert _run( + s.select(current_train_weight=5, min_prompt_groups=1, max_prompt_groups=8) + ) == (None, 0) + + +class TestDefaultEvictSkipsUnready: + def test_windowed_evict_drops_ready_below_window(self): + buf = FakeBuffer() + buf.add("stale", weight=0, ready=True) + buf.add("fresh", weight=5, ready=True) + s = WindowedSampler(buf, max_staleness_versions=1) + removed = _run(s.evict(current_train_weight=5)) # min_valid = 4 + assert removed == 1 + assert buf.start_weight_list == [5] + + def test_windowed_evict_skips_unready_stale(self): + buf = FakeBuffer() + buf.add("stale_unready", weight=0, ready=False) + s = WindowedSampler(buf, max_staleness_versions=1) + assert _run(s.evict(current_train_weight=5)) == 0 + + +class EchoSampler(InOrderSampler): + """Stand-in for a user-defined sampler loaded by FQN.""" diff --git a/tests/unit/single_controller/test_sc_utils_helpers.py b/tests/unit/single_controller/test_sc_utils_helpers.py new file mode 100644 index 0000000000..4d462b81bb --- /dev/null +++ b/tests/unit/single_controller/test_sc_utils_helpers.py @@ -0,0 +1,186 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for single_controller_utils.utils pure helpers.""" + +from __future__ import annotations + +import math + +import pytest +import torch +from tensordict import TensorDict + +from nemo_rl.algorithms.single_controller_utils.utils import ( + aggregate_step_metrics, + fields_for_put, + reduce_advantage_pump_metrics, + squeeze_trailing_unit_dim, + tensor_field, +) +from nemo_rl.data_plane import KVBatchMeta + + +def _meta(size: int, sequence_lengths: list[int] | None = None) -> KVBatchMeta: + return KVBatchMeta( + partition_id="rollout_data", + task_name="train", + sample_ids=[f"s{i}" for i in range(size)], + sequence_lengths=sequence_lengths, + ) + + +class TestSqueezeTrailingUnitDim: + def test_squeezes_trailing_unit_dim(self) -> None: + out = squeeze_trailing_unit_dim(torch.zeros(4, 1)) + assert out.shape == (4,) + + def test_leaves_1d_untouched(self) -> None: + out = squeeze_trailing_unit_dim(torch.zeros(4)) + assert out.shape == (4,) + + def test_leaves_non_unit_trailing_dim(self) -> None: + out = squeeze_trailing_unit_dim(torch.zeros(4, 3)) + assert out.shape == (4, 3) + + +class TestTensorField: + def test_returns_dense_tensor(self) -> None: + td = TensorDict({"x": torch.arange(6).reshape(2, 3)}, batch_size=[2]) + out = tensor_field(td, "x") + assert torch.equal(out, torch.arange(6).reshape(2, 3)) + + def test_pads_nested_tensor(self) -> None: + nested = torch.nested.as_nested_tensor( + [torch.tensor([1, 2, 3]), torch.tensor([4, 5])], + layout=torch.jagged, + ) + td = TensorDict({"x": nested}, batch_size=[2]) + out = tensor_field(td, "x") + assert not out.is_nested + assert out.shape == (2, 3) + assert out[1].tolist() == [4, 5, 0] + + def test_non_tensor_raises_type_error(self) -> None: + td = TensorDict({"x": torch.zeros(2)}, batch_size=[2], non_blocking=False) + td.set_non_tensor("meta", ["a", "b"]) + with pytest.raises(TypeError): + tensor_field(td, "meta") + + +class TestAggregateStepMetrics: + def test_scalar_loss_and_grad_norm_tensors(self) -> None: + result = { + "loss": torch.tensor([1.0, 3.0]), + "grad_norm": torch.tensor(2.0), + } + out = aggregate_step_metrics(result) + assert out["loss"] == pytest.approx(2.0) + assert out["grad_norm"] == pytest.approx(2.0) + + def test_float_loss_and_optional_scalars(self) -> None: + out = aggregate_step_metrics({"loss": 0.5, "total_flops": 10, "num_ranks": 4}) + assert out["loss"] == pytest.approx(0.5) + assert out["total_flops"] == pytest.approx(10.0) + assert out["num_ranks"] == 4 + assert "grad_norm" not in out + + def test_mb_metric_reduction_rules(self) -> None: + result = { + "all_mb_metrics": { + "probs_ratio_min": [0.4, 0.2, 0.9], + "probs_ratio_max": [0.4, 0.2, 0.9], + "lr": [0.1, 0.3], + "some_sum_metric": [1.0, 2.0, 3.0], + } + } + out = aggregate_step_metrics(result) + assert out["probs_ratio_min"] == pytest.approx(0.2) + assert out["probs_ratio_max"] == pytest.approx(0.9) + assert out["lr"] == pytest.approx(0.2) + assert out["some_sum_metric"] == pytest.approx(6.0) + + def test_min_max_all_inf_falls_back_to_neg_one(self) -> None: + result = { + "all_mb_metrics": { + "probs_ratio_min": [math.inf, math.inf], + "probs_ratio_max": [math.inf], + } + } + out = aggregate_step_metrics(result) + assert out["probs_ratio_min"] == -1.0 + assert out["probs_ratio_max"] == -1.0 + + def test_moe_and_mtp_metrics_are_prefixed(self) -> None: + result = { + "moe_metrics": {"load_balance": [1.0, 3.0]}, + "mtp_metrics": {"acc": [2.0, 2.0]}, + } + out = aggregate_step_metrics(result) + assert out["moe/load_balance"] == pytest.approx(4.0) + assert out["mtp/acc"] == pytest.approx(4.0) + + +class TestReduceAdvantagePumpMetrics: + def test_reward_and_advantages_and_tokens(self) -> None: + out = reduce_advantage_pump_metrics( + rewards=[torch.tensor([1.0, 3.0])], + masked_advantages=[torch.tensor([-1.0, 0.0, 2.0])], + sequence_lengths=[4, 6], + ) + assert out["reward"] == pytest.approx(2.0) + assert out["advantages/mean"] == pytest.approx(1.0 / 3.0) + assert out["advantages/max"] == pytest.approx(2.0) + assert out["advantages/min"] == pytest.approx(-1.0) + assert out["total_num_tokens"] == pytest.approx(10.0) + + def test_empty_advantages_tensor_yields_zeros(self) -> None: + out = reduce_advantage_pump_metrics( + rewards=[], + masked_advantages=[torch.empty(0)], + sequence_lengths=[], + ) + assert out["advantages/mean"] == 0.0 + assert out["advantages/max"] == 0.0 + assert out["advantages/min"] == 0.0 + assert "reward" not in out + assert "total_num_tokens" not in out + + def test_all_empty_inputs_returns_empty_dict(self) -> None: + assert reduce_advantage_pump_metrics([], [], []) == {} + + +class TestFieldsForPut: + def test_no_sequence_lengths_packs_contiguous(self) -> None: + meta = _meta(2, sequence_lengths=None) + out = fields_for_put(meta, {"advantages": torch.zeros(2, 3)}) + assert out.batch_size == torch.Size([2]) + assert not out["advantages"].is_nested + assert out["advantages"].shape == (2, 3) + + def test_renests_padded_rows_by_sequence_length(self) -> None: + meta = _meta(2, sequence_lengths=[3, 2]) + value = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 0.0]]) + out = fields_for_put(meta, {"advantages": value}) + assert out["advantages"].is_nested + rows = out["advantages"].unbind() + assert rows[0].tolist() == [1.0, 2.0, 3.0] + assert rows[1].tolist() == [4.0, 5.0] + + def test_non_matching_width_stays_contiguous(self) -> None: + meta = _meta(2, sequence_lengths=[3, 2]) + value = torch.zeros(2, 1) + out = fields_for_put(meta, {"scalar": value}) + assert not out["scalar"].is_nested + assert out["scalar"].shape == (2, 1) diff --git a/tests/unit/single_controller/test_single_controller.py b/tests/unit/single_controller/test_single_controller.py new file mode 100644 index 0000000000..e9ce01b224 --- /dev/null +++ b/tests/unit/single_controller/test_single_controller.py @@ -0,0 +1,166 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for SingleController initialization and pump lifecycle.""" + +import asyncio +from types import SimpleNamespace + +import pytest + +import nemo_rl.algorithms.single_controller as single_controller +from nemo_rl.algorithms.single_controller import ( + SingleControllerActor, + SingleControllerConfig, +) +from nemo_rl.data_plane import KVBatchMeta +from nemo_rl.utils.timer import Timer + + +def test_rejects_multiple_optimizer_steps_per_rl_step(monkeypatch) -> None: + monkeypatch.setattr(single_controller, "Logger", lambda _: object()) + cfg = SingleControllerConfig.model_construct( + min_groups_per_batch=1, + target_groups_per_step=2, + group_size=4, + train_global_batch_size=4, + advantage_enabled=False, + logger={}, + ) + controller_cls = SingleControllerActor.__ray_metadata__.modified_class + + with pytest.raises( + ValueError, + match=( + r"train_global_batch_size \(4\) must equal " + r"target_groups_per_step \(2\) \* group_size \(4\) = 8; " + r"multiple optimizer steps per RL step are not supported" + ), + ): + controller_cls( + cfg=cfg, + dp_client_handle=None, + gen_handle=None, + trainer_handle=None, + weight_synchronizer=None, + loss_fn=None, + rollout_manager=SimpleNamespace(_tq_buffer=None), + ) + + +class _EmptySampler: + async def evict(self, *, current_train_weight: int) -> int: + del current_train_weight + return 0 + + async def select(self, **kwargs): + del kwargs + return None, 0 + + +class _OneThenEmptySampler(_EmptySampler): + def __init__(self, meta: KVBatchMeta) -> None: + self._meta: KVBatchMeta | None = meta + + async def select(self, **kwargs): + del kwargs + if self._meta is None: + return None, 0 + meta = self._meta + self._meta = None + return meta, 1 + + +class _EmptyBuffer: + def __len__(self) -> int: + return 0 + + +class _NoOpTrainer: + def prepare_for_lp_inference(self) -> None: + pass + + def prepare_for_training(self) -> None: + pass + + def begin_train_step(self, loss_fn) -> None: + del loss_fn + + def train_microbatches_from_meta(self, meta: KVBatchMeta) -> None: + del meta + + +class _NoOpDataPlane: + def clear_samples(self, **kwargs) -> None: + del kwargs + + +def _train_pump_controller(*, sampler) -> object: + controller_cls = SingleControllerActor.__ray_metadata__.modified_class + ctrl = object.__new__(controller_cls) + ctrl._cfg = SimpleNamespace( + target_groups_per_step=2, + min_groups_per_batch=1, + max_train_steps=1, + advantage_policy_logprobs_field=None, + advantage_reference_logprobs_field=None, + advantage_enabled=False, + partition_id="rollout_data", + ) + ctrl._sampler = sampler + ctrl._buffer = _EmptyBuffer() + ctrl._buffer_capacity = asyncio.Semaphore(2) + ctrl._rollout_exhausted = asyncio.Event() + ctrl._rollout_exhausted.set() + ctrl._trainer = _NoOpTrainer() + ctrl._loss_fn = None + ctrl._dp_client = _NoOpDataPlane() + ctrl._timer = Timer() + ctrl._trainer_version = 0 + ctrl._train_steps = 0 + ctrl._step_log_dict = { + "rewards": [], + "masked_advantages": [], + "sequence_lengths": [], + } + return ctrl + + +def test_train_pump_stops_after_rollout_exhaustion_and_buffer_drain() -> None: + ctrl = _train_pump_controller(sampler=_EmptySampler()) + + asyncio.run(asyncio.wait_for(ctrl._train_pump(), timeout=1.0)) + + assert ctrl._train_steps == 0 + + +def test_train_pump_fails_if_rollout_exhausts_during_partial_step() -> None: + meta = KVBatchMeta( + partition_id="rollout_data", + task_name="train", + sample_ids=["sample-0"], + fields=[], + sequence_lengths=[1], + tags=[{"weight_version": 0}], + ) + ctrl = _train_pump_controller(sampler=_OneThenEmptySampler(meta)) + + with pytest.raises( + RuntimeError, + match=( + r"rollout exhausted before a complete training step was assembled: " + r"dispatched 1/2 prompt groups" + ), + ): + asyncio.run(asyncio.wait_for(ctrl._train_pump(), timeout=1.0)) diff --git a/tests/unit/single_controller/test_train_pump.py b/tests/unit/single_controller/test_train_pump.py index 1763e0efca..b6134df948 100644 --- a/tests/unit/single_controller/test_train_pump.py +++ b/tests/unit/single_controller/test_train_pump.py @@ -12,56 +12,162 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Dry-run test: SC._train_pump drives the split-API cycle per mini-batch. - -TODO: once the SingleController setup entrypoint lands (real TQPolicy, -WeightSynchronizer, and advantage_estimator wiring), replace the fakes -below with the real components and turn this into an integration test. -Today the trainer / weight_synchronizer / advantage_estimator are all -stubs — the test exercises the pump's control flow (claim → adv-stage → -split-API cycle → clear → sync) but not the collaborators' semantics. -""" +"""End-to-end tests for SC._train_pump.""" from __future__ import annotations +import gc +import math +from types import SimpleNamespace +from typing import Any + import pytest import ray import torch from tensordict import TensorDict +from nemo_rl.algorithms.async_utils.replay_buffer import TQReplayBuffer +from nemo_rl.algorithms.async_utils.staleness_sampler import WindowedSamplerConfig from nemo_rl.algorithms.single_controller import ( SingleControllerActor, SingleControllerConfig, ) -from nemo_rl.data_plane.adapters.noop import NoOpDataPlaneClient - -_PARTITION_ID = "rollout_data" -_ROLLOUT_FIELDS = [ +from nemo_rl.algorithms.utils import get_tokenizer +from nemo_rl.data_plane import KVBatchMeta +from nemo_rl.distributed.virtual_cluster import RayVirtualCluster +from nemo_rl.models.generation import configure_generation_config +from nemo_rl.models.policy.tq_policy import TQPolicy +from tests.unit.models.policy.test_megatron_worker import create_megatron_test_config +from tests.unit.single_controller._dp_fakes import _PARTITION_ID +from tests.unit.test_utils import SimpleLossFn + +# Union of DP_TRAIN_FIELDS (TQPolicy) + rollout extras (total_reward, +# prompt_ids_for_adv) — the partition schema must cover every field any +# producer/consumer touches. +_REGISTERED_FIELDS = [ "input_ids", "input_lengths", + "generation_logprobs", + "prev_logprobs", + "reference_policy_logprobs", + "advantages", "token_mask", "sample_mask", "total_reward", "prompt_ids_for_adv", ] -_ADV_FIELD = "advantages" -@ray.remote(num_cpus=0) -class _DPActor(NoOpDataPlaneClient): - """Ray-wrapped NoOpDataPlaneClient for cross-process DP inspection.""" +def _simple_tq_cfg() -> dict: + """Simple in-process TQ backend cfg (mirrors tests/unit/data_plane/conftest.py).""" + return { + "enabled": True, + "impl": "transfer_queue", + "backend": "simple", + "storage_capacity": 1024, + "num_storage_units": 1, + "claim_meta_poll_interval_s": 0.5, + "global_segment_size": 8589934592, # 8 GiB + "local_buffer_size": 1073741824, # 1 GiB + } + + +def _populate_group( + dp_client, + *, + group_uuid: str, + group_size: int, + seq_len: int, + weight_version: int, +) -> KVBatchMeta: + """Write one complete prompt group to DP and return its meta.""" + sample_ids = [f"{group_uuid}_g{i}" for i in range(group_size)] + fields = TensorDict( + { + "input_ids": torch.randint(0, 1000, (group_size, seq_len)).long(), + "input_lengths": torch.tensor([seq_len] * group_size).long(), + "token_mask": torch.ones(group_size, seq_len, dtype=torch.long), + "sample_mask": torch.ones(group_size, dtype=torch.long), + "generation_logprobs": torch.zeros( + group_size, seq_len, dtype=torch.float32 + ), + # prev_logprobs / reference_policy_logprobs are read by + # TQPolicy workers as part of DP_TRAIN_FIELDS; pre-seed with + # zeros since the loss (SimpleLossFn) ignores them. + "prev_logprobs": torch.zeros(group_size, seq_len, dtype=torch.float32), + "reference_policy_logprobs": torch.zeros( + group_size, seq_len, dtype=torch.float32 + ), + "total_reward": torch.arange(group_size, dtype=torch.float32) * 0.5 + 1.0, + "prompt_ids_for_adv": torch.zeros(group_size, seq_len, dtype=torch.long), + }, + batch_size=(group_size,), + ) + tags = [{"weight_version": int(weight_version)} for _ in range(group_size)] + dp_client.put_samples( + sample_ids=sample_ids, + partition_id=_PARTITION_ID, + fields=fields, + tags=tags, + ) + return KVBatchMeta( + partition_id=_PARTITION_ID, + task_name="train", + sample_ids=sample_ids, + fields=list(fields.keys()), + sequence_lengths=[seq_len] * group_size, + tags=[dict(t) for t in tags], + ) + + +def _prepopulate_buffer( + buffer: TQReplayBuffer, meta: KVBatchMeta, *, weight_version: int +) -> None: + """Insert a ready slot into TQReplayBuffer, mirroring what commit() sets.""" + buffer.meta_list.append(meta) + buffer.start_weight_list.append(int(weight_version)) + buffer.end_weight_list.append(int(weight_version)) + buffer.target_step_list.append(None) + buffer.ready_list.append(True) + # Group id follows pack_payload's "{group_uuid}_g{i}" convention. + group_id = meta.sample_ids[0].rpartition("_g")[0] + buffer._group_ids.append(group_id) + + +@pytest.fixture(scope="function") +def train_cluster(): + """Single-GPU virtual cluster for the trainer worker group.""" + cluster = RayVirtualCluster( + name="test-sc-train-cluster", + bundle_ct_per_node_list=[1], + use_gpus=True, + num_gpus_per_node=1, + max_colocated_worker_groups=1, + ) + try: + yield cluster + finally: + cluster.shutdown() + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + - def claim_meta(self, *args, **kwargs): - meta = super().claim_meta(*args, **kwargs) - rec = self._partitions[meta.partition_id] - meta.tags = [dict(rec.tags.get(sid, {})) for sid in meta.sample_ids] - return meta +class _FakeAdvEstimator: + """Passthrough advantage: returns rewards unchanged.""" - def peek_count(self, partition_id: str) -> int: - return len(self._partitions[partition_id].rows) + def compute_advantage( + self, + prompt_ids: torch.Tensor, + rewards: torch.Tensor, + mask: torch.Tensor, + repeated_batch: dict, + **kwargs, + ) -> torch.Tensor: + return rewards.unsqueeze(-1).expand(mask.shape).detach().clone() -@ray.remote(num_cpus=0) +@ray.remote(num_cpus=0) # pragma: no cover class _CallLog: """Ordered append-only log the fakes below write to.""" @@ -76,29 +182,43 @@ def get(self) -> list[tuple[str, dict]]: return list(self._entries) -class _FakeTrainer: - """Sync driver-side TQPolicy stand-in.""" +class _RecordingLogger: + """Forward metrics logged inside SC to a Ray actor visible to the test.""" - def __init__(self, log_handle) -> None: + def __init__(self, log_handle: Any) -> None: self._log = log_handle - def begin_train_step(self, loss_fn) -> None: - ray.get(self._log.record.remote("begin_train_step", {})) - - def train_microbatches_from_meta(self, meta) -> None: + def log_metrics( + self, + metrics: dict[str, Any], + step: int, + prefix: str | None = "", + ) -> None: ray.get( self._log.record.remote( - "train_microbatches_from_meta", {"meta_size": int(meta.size)} + "metrics", + { + "metrics": dict(metrics), + "step": int(step), + "prefix": prefix, + }, ) ) - def finish_train_step(self) -> dict: - ray.get(self._log.record.remote("finish_train_step", {})) - return {"loss": 0.0} + +@ray.remote(num_cpus=1, num_gpus=0) +class _RecordingSingleControllerActor( + SingleControllerActor.__ray_metadata__.modified_class +): + """SingleControllerActor variant with an observable logger.""" + + def __init__(self, *, metric_log_handle: Any, **kwargs: Any) -> None: + super().__init__(**kwargs) + self._logger = _RecordingLogger(metric_log_handle) class _FakeWeightSync: - """Async WeightSynchronizer stand-in.""" + """Records the version SC syncs at end-of-step, via a Ray actor log.""" def __init__(self, log_handle) -> None: self._log = log_handle @@ -107,194 +227,137 @@ async def sync_weights(self, version: int) -> None: ray.get(self._log.record.remote("sync_weights", {"version": int(version)})) -class _FakeAdvEstimator: - """Records shapes SC hands the estimator; returns a per-sample scalar.""" - - def __init__(self, log_handle) -> None: - self._log = log_handle +@pytest.mark.mcore +@pytest.mark.hf_gated +def test_train_pump_drives_mcore_training_step( + train_cluster, + tiny_llama_model_path, + tmp_path, +): + """SC._train_pump runs one outer step against a real Megatron TQPolicy.""" + train_steps = 2 + train_gbs = 4 + num_generations = 2 + num_prompts = 2 + + policy_cfg = create_megatron_test_config(tiny_llama_model_path, tp=1, pp=1) + policy_cfg["train_global_batch_size"] = train_gbs + policy_cfg["train_micro_batch_size"] = train_gbs + tokenizer = get_tokenizer(policy_cfg["tokenizer"]) + policy_cfg["generation"] = configure_generation_config( + policy_cfg["generation"], tokenizer + ) - def compute_advantage( - self, - prompt_ids: torch.Tensor, - rewards: torch.Tensor, - mask: torch.Tensor, - repeated_batch: dict, - **kwargs, - ) -> torch.Tensor: - ray.get( - self._log.record.remote( - "compute_advantage", - { - "rewards_shape": tuple(rewards.shape), - "mask_shape": tuple(mask.shape), - "prompt_ids_shape": tuple(prompt_ids.shape), - "rb_keys": sorted(repeated_batch.keys()), - "kwargs_keys": sorted(kwargs.keys()), - }, - ) - ) - return rewards.detach().clone() + dp_cfg = _simple_tq_cfg() + # TQPolicy's ctor bootstraps the TQ controller (bootstrap=True) and + # fans out setup_data_plane to workers (bootstrap=False on workers). + trainer = TQPolicy( + cluster=train_cluster, + config=policy_cfg, + tokenizer=tokenizer, + dp_cfg=dp_cfg, + tq_partition_id=_PARTITION_ID, + ) + try: + # Reuse TQPolicy's driver-side dp_client for SC — same controller, + # no double-bootstrap. + dp_client = trainer.dp_client -def _populate_group( - dp, - group_uuid: str, - group_size: int, - seq_len: int, - weight_version: int, -) -> list[str]: - """Write one complete prompt group to DP.""" - sample_ids = [f"{group_uuid}_g{i}" for i in range(group_size)] - fields = TensorDict( - { - "input_ids": torch.arange(group_size * seq_len) - .reshape(group_size, seq_len) - .long(), - "input_lengths": torch.tensor([seq_len] * group_size).long(), - "token_mask": torch.ones(group_size, seq_len, dtype=torch.long), - "sample_mask": torch.ones(group_size, 1, dtype=torch.long), - "total_reward": ( - torch.arange(group_size, dtype=torch.float32) * 0.5 + 1.0 - ).unsqueeze(-1), - "prompt_ids_for_adv": torch.zeros(group_size, seq_len, dtype=torch.long), - }, - batch_size=(group_size,), - ) - tags = [ - { - "weight_version": int(weight_version), - "expected_num_samples": int(group_size), - "committed": True, - "input_lengths": int(seq_len), - } - for _ in range(group_size) - ] - ray.get( - dp.put_samples.remote( - sample_ids=sample_ids, + # Register the partition once with the full field union. All + # pre-populated samples across every step live here at once. + dp_client.register_partition( partition_id=_PARTITION_ID, - fields=fields, - tags=tags, + fields=_REGISTERED_FIELDS, + num_samples=train_steps * train_gbs, + consumer_tasks=["prev_lp", "ref_lp", "train"], ) - ) - return sample_ids - - -@pytest.fixture(scope="module") -def ray_cluster(): - if not ray.is_initialized(): - ray.init(ignore_reinit_error=True, num_cpus=4, log_to_driver=False) - yield - ray.shutdown() - - -def test_train_pump_dry_run(ray_cluster, tmp_path): - """SC._train_pump runs 2 mini-batches × 1 group each, then a single _sync_weights.""" - del ray_cluster - group_size = 2 - seq_len = 4 - num_groups = 2 # target_groups_per_step - # train_gbs == group_size → groups_per_minibatch=1, num_minibatches=2. - train_gbs = group_size - - dp = _DPActor.remote() - log = _CallLog.remote() - ray.get( - dp.register_partition.remote( + + # Real TQReplayBuffer, pre-populated with one ready batch (num_prompts + # groups of num_generations samples) per step, stamped with the + # step's weight_version so the sampler can advance across steps. + tq_buffer = TQReplayBuffer( + dp_client, partition_id=_PARTITION_ID, - fields=[*_ROLLOUT_FIELDS, _ADV_FIELD], - num_samples=group_size * num_groups * 4, - consumer_tasks=["train"], + pad_value_dict={"input_ids": int(tokenizer.pad_token_id or 0)}, ) - ) - for g in range(num_groups): - _populate_group( - dp, - group_uuid=f"prompt{g}", - group_size=group_size, - seq_len=seq_len, - weight_version=0, + for step in range(train_steps): + for g in range(num_prompts): + meta = _populate_group( + dp_client, + group_uuid=f"step{step}_prompt{g}", + group_size=num_generations, + seq_len=16, + weight_version=step, + ) + _prepopulate_buffer(tq_buffer, meta, weight_version=step) + + log = _CallLog.remote() + weight_sync = _FakeWeightSync(log) + adv_est = _FakeAdvEstimator() + # Rollout manager stub — SC.__init__ only touches ._tq_buffer. + rollout_manager = SimpleNamespace(_tq_buffer=None) + + cfg = SingleControllerConfig.model_construct( + sampler=WindowedSamplerConfig(max_staleness_versions=1), + min_groups_per_batch=num_prompts, + target_groups_per_step=num_prompts, + group_size=num_generations, + max_inflight_prompts=num_prompts, + max_buffered_rollouts=num_prompts, + max_train_steps=train_steps, + max_num_epochs=None, + train_global_batch_size=train_gbs, + partition_id=_PARTITION_ID, + advantage_enabled=True, + advantage_repeated_batch_fields=[], + advantage_policy_logprobs_field=None, + advantage_reference_logprobs_field=None, + diagnostics=False, + logger={ + "log_dir": str(tmp_path / "logs"), + "wandb_enabled": False, + "swanlab_enabled": False, + "tensorboard_enabled": False, + "mlflow_enabled": False, + "monitor_gpus": False, + }, ) - trainer = _FakeTrainer(log) - weight_sync = _FakeWeightSync(log) - adv_est = _FakeAdvEstimator(log) - - cfg = SingleControllerConfig.model_construct( - max_weight_staleness_versions=1, - min_groups_per_batch=num_groups, - target_groups_per_step=num_groups, - group_size=group_size, - batch_selection_strategy="staleness_window", - max_inflight_prompts=8, - max_buffered_rollouts=8, - max_train_steps=1, - max_rollout_prompts=num_groups, - train_global_batch_size=train_gbs, - partition_id=_PARTITION_ID, - consumer_task_name="train", - claim_required_fields=["input_ids"], - max_claim_groups=8, - advantage_enabled=True, - advantage_repeated_batch_fields=[], - advantage_policy_logprobs_field=None, - advantage_reference_logprobs_field=None, - diagnostics=False, - logger={ - "log_dir": str(tmp_path / "logs"), - "wandb_enabled": False, - "swanlab_enabled": False, - "tensorboard_enabled": False, - "mlflow_enabled": False, - "monitor_gpus": False, - }, - ) - - ctrl = SingleControllerActor.remote( - cfg=cfg, - prompts=[], # rollout pump is not started - dp_client_handle=dp, - gen_handle=None, - trainer_handle=trainer, - weight_synchronizer=weight_sync, - loss_fn=object(), - advantage_estimator=adv_est, - ) + ctrl = _RecordingSingleControllerActor.remote( + metric_log_handle=log, + cfg=cfg, + dp_client_handle=dp_client, + gen_handle=None, + trainer_handle=trainer, + weight_synchronizer=weight_sync, + loss_fn=SimpleLossFn(), + rollout_manager=rollout_manager, + advantage_estimator=adv_est, + dataloader=None, + tq_buffer=tq_buffer, + ) - # _train_pump exits when max_train_steps is reached — no cancel needed. - ray.get(ctrl._train_pump.remote()) - - state = ray.get(ctrl.ping.remote()) - assert state["train_steps"] == 1 - assert state["trainer_version"] == 2 # one bump per mini-batch - - entries = ray.get(log.get.remote()) - kinds = [k for k, _ in entries] - assert kinds == [ - "compute_advantage", - "begin_train_step", - "train_microbatches_from_meta", - "finish_train_step", - "compute_advantage", - "begin_train_step", - "train_microbatches_from_meta", - "finish_train_step", - "sync_weights", - ] - - # sync_weights fires with the post-bump trainer_version. - sync_payload = next(p for k, p in entries if k == "sync_weights") - assert sync_payload["version"] == 2 - - for k, p in entries: - if k == "train_microbatches_from_meta": - assert p["meta_size"] == group_size - elif k == "compute_advantage": - assert p["rewards_shape"] == (group_size,) - assert p["mask_shape"] == (group_size, seq_len) - assert p["prompt_ids_shape"] == (group_size, seq_len) - assert p["rb_keys"] == ["total_reward"] - assert p["kwargs_keys"] == [] - - # clear_samples was called after each finish_train_step. - assert ray.get(dp.peek_count.remote(_PARTITION_ID)) == 0 + # train_steps outer steps, each: sampler.select → advantage stage → begin/microbatches/finish → sync. + ray.get(ctrl._train_pump.remote()) + + state = ray.get(ctrl.ping.remote()) + assert state["train_steps"] == train_steps + assert state["trainer_version"] == train_steps + + entries = ray.get(log.get.remote()) + sync_versions = [p["version"] for k, p in entries if k == "sync_weights"] + assert sync_versions == list(range(1, train_steps + 1)) + + train_metrics = [ + p["metrics"] + for kind, p in entries + if kind == "metrics" and p["prefix"] == "train" + ] + assert len(train_metrics) == train_steps + for metrics in train_metrics: + assert math.isfinite(metrics["reward"]) + assert math.isfinite(metrics["advantages/mean"]) + + finally: + trainer.shutdown()