From 30aad15874be92b38485aaf6d8e8ec7a6907caa2 Mon Sep 17 00:00:00 2001 From: Yuki Huang Date: Wed, 15 Jul 2026 08:59:46 -0700 Subject: [PATCH 01/17] =?UTF-8?q?feat(sc):=20train=20path=20=E2=80=94=20St?= =?UTF-8?q?alenessSampler=20+=20=5Ftrain=5Fpump=20rewrite?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Yuki Huang --- .../async_utils/staleness_sampler.py | 160 ++++++ nemo_rl/algorithms/single_controller.py | 387 ++++++--------- nemo_rl/algorithms/staleness_sampler.py | 220 --------- pyrefly.toml | 1 + .../test_staleness_sampler.py | 466 ++++++++++++++++++ 5 files changed, 789 insertions(+), 445 deletions(-) create mode 100644 nemo_rl/algorithms/async_utils/staleness_sampler.py delete mode 100644 nemo_rl/algorithms/staleness_sampler.py create mode 100644 tests/unit/single_controller/test_staleness_sampler.py 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..144dabdd3d --- /dev/null +++ b/nemo_rl/algorithms/async_utils/staleness_sampler.py @@ -0,0 +1,160 @@ +# 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 selection over a TQReplayBuffer.""" + +from nemo_rl.algorithms.async_utils.replay_buffer import TQReplayBuffer +from nemo_rl.data_plane import KVBatchMeta + + +class StalenessSampler: + """Pick complete prompt groups from a TQReplayBuffer. + + Args: + buffer: Shared TQReplayBuffer holding the candidate slots. + max_staleness_versions: Max weight-version gap a sample may have from the trainer. + sample_freshest_first: Prefer smallest lag when picking from the in-window set. + require_order: Take only from the oldest in-window weight_version and wait for its batch to fill. + force_in_order: Match each slot's target_step against current_train_weight, ignoring the window; mirrors legacy async_grpo target_weight semantics. + """ + + def __init__( + self, + buffer: TQReplayBuffer, + max_staleness_versions: int, + sample_freshest_first: bool = False, + require_order: bool = False, + force_in_order: bool = False, + ) -> None: + if max_staleness_versions < 0: + raise ValueError( + f"max_staleness_versions must be non-negative, got " + f"{max_staleness_versions}" + ) + if require_order and sample_freshest_first: + raise ValueError( + "require_order and sample_freshest_first are mutually exclusive" + ) + self._buffer = buffer + self.max_staleness_versions = max_staleness_versions + self.sample_freshest_first = sample_freshest_first + self.require_order = require_order + self.force_in_order = force_in_order + + async def select( + self, + *, + current_train_weight: int, + min_prompt_groups: int, + max_prompt_groups: int, + ) -> tuple[KVBatchMeta | None, int]: + """Concat up to max_prompt_groups eligible groups and drop them from the buffer. + + Eligibility = ready and weight in + [current_train_weight - max_staleness_versions, current_train_weight]. + DataPlane rows survive the local drop; caller clears them at step boundary. + + Args: + current_train_weight: Current trainer weight version. + min_prompt_groups: Minimum groups required; returns (None, 0) below this. + max_prompt_groups: Cap on groups returned when the threshold is met. + + Returns: + meta: Concatenated KVBatchMeta, or None if not enough groups. + num_groups: Number of prompt groups in meta; 0 when meta is 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})" + ) + + if self.force_in_order: + # target_step exact match; staleness window ignored. + valid_idxs = [ + i + for i, target in enumerate(self._buffer.target_step_list) + if target == current_train_weight and self._buffer.ready_list[i] + ] + else: + min_valid_version = max( + 0, current_train_weight - self.max_staleness_versions + ) + if self.require_order: + 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] + ] + else: + 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 len(valid_idxs) < min_prompt_groups: + return None, 0 + + if self.sample_freshest_first: + valid_idxs.sort( + key=lambda i: ( + current_train_weight - self._buffer.start_weight_list[i], + i, + ) + ) + + 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), + ) + + async def evict(self, *, current_train_weight: int) -> int: + """Drop groups whose weight falls below the staleness window. + + Future entries (weight > current_train_weight) are left alone. + + Args: + current_train_weight: Current trainer weight version; groups with + weight < current_train_weight - max_staleness_versions are dropped. + + Returns: + Number of group entries removed from the buffer. + """ + min_valid_version = max(0, current_train_weight - self.max_staleness_versions) + stale_idxs = [ + i + for i, weight in enumerate(self._buffer.start_weight_list) + if weight < min_valid_version + ] + if not stale_idxs: + return 0 + return await self._buffer.remove(stale_idxs, remove_in_dp=True) diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 16a0461f31..5ae67f11b4 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -544,201 +544,193 @@ 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_pump(train_meta). + 4. trainer.train_microbatch_from_meta + finish_train_step. + 5. dp_client.clear_samples on consumed sample_ids; release _buffer_capacity + per dropped group, then sync. """ + adv_cfg = self._advantage_cfg + grpo_cfg = self._master_config.grpo + # 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 = ( - self._cfg.advantage_reference_logprobs_field is not None - ) + prev_logprobs_required = adv_cfg.policy_logprobs_field is not None + reference_logprobs_required = adv_cfg.reference_logprobs_field is not None + + while self._train_steps < grpo_cfg["max_num_steps"]: + groups_dispatched = 0 + min_sample_version = None + step_open = False + + with self._timer.time("total_step_time"): + while groups_dispatched < grpo_cfg["num_prompts_per_step"]: + # 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, + ) + if evicted: + print( + f" evicted {evicted} stale prompt group(s)", + flush=True, + ) + for _ in range(evicted): + self._buffer_capacity.release() - 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, + # Select a batch + max_prompt_groups = ( + grpo_cfg["num_prompts_per_step"] - groups_dispatched + ) + min_prompt_groups = min( + self._async_cfg.min_prompt_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, ) - for _ in range(evicted_groups): + + # If no batch is selectable, sleep and retry + if train_meta is None: + await asyncio.sleep(0.005) + continue + + # Release buffer capacity + for _ in range(num_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, - ) + # 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 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, - ) - - # 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) - - 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, + ) + + if train_meta.sequence_lengths: + self._step_log_dict["sequence_lengths"].extend( + int(s) for s in train_meta.sequence_lengths ) - step_open = True - await asyncio.to_thread( - self._trainer.train_microbatches_from_meta, - group_meta, + # Refresh min_sample_version + curr_min_sample_version = min( + t["weight_version"] + for t in train_meta.tags # type: ignore ) - 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 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 + + # Remove consumed sample_ids from the buffer + await self._call_dp( + "clear_samples", + sample_ids=list(train_meta.sample_ids), + partition_id=self._partition_id, + ) + + groups_dispatched += num_groups 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)", + "train_pump: rollout exhausted before any group ready", 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 - ) - - # 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 + with self._timer.time("policy_training"): + result = await asyncio.to_thread(self._trainer.finish_train_step) - consumed_ids = list(self._step_consumed_sample_ids) - self._step_consumed_sample_ids = [] - with self._timed("clear_samples"): - await self._call_dp( - "clear_samples", - sample_ids=consumed_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, + step_metrics = aggregate_step_metrics(result) + step_metrics.update( + reduce_advantage_pump_metrics(**self._step_log_dict) ) + self._step_log_dict = {k: [] for k in self._step_log_dict} - # 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, + 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) + cluster_cfg = self._master_config.cluster + total_num_gpus = cluster_cfg["num_nodes"] * cluster_cfg["gpus_per_node"] + if total_time > 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 ) - self._timer.reset() - 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 - - 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 = current trainer version - oldest sample version. + lag = self._trainer_version - min_sample_version # type: ignore + print( + f"train step {self._train_steps}/{grpo_cfg['max_num_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. @@ -843,36 +835,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, @@ -887,31 +849,6 @@ def _advantage_input_fields(self) -> list[str]: 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] 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..a70ea16100 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", diff --git a/tests/unit/single_controller/test_staleness_sampler.py b/tests/unit/single_controller/test_staleness_sampler.py new file mode 100644 index 0000000000..ce15f34acf --- /dev/null +++ b/tests/unit/single_controller/test_staleness_sampler.py @@ -0,0 +1,466 @@ +# 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 StalenessSampler (pure filter over TQReplayBuffer state).""" + +from __future__ import annotations + +import asyncio + +import pytest + +from nemo_rl.algorithms.async_utils.staleness_sampler import StalenessSampler +from nemo_rl.data_plane import KVBatchMeta + + +class FakeBuffer: + """Minimal TQReplayBuffer surface used by StalenessSampler tests.""" + + 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.ready_list: list[bool] = [] + self.remove_calls: list[tuple[list[int], bool]] = [] + + def add( + self, + group_id: str, + weight: int, + group_size: int = 1, + ready: bool = True, + end_weight: int | None = None, + ) -> KVBatchMeta: + sample_ids = [f"{group_id}_g{i}" for i in range(group_size)] + meta = KVBatchMeta( + partition_id=self._partition_id, + task_name=None, + sample_ids=sample_ids, + tags=[{"weight_version": weight, "group_id": group_id}] * group_size, + ) + self.meta_list.append(meta if ready else None) + self.start_weight_list.append(weight) + self.end_weight_list.append(weight if end_weight is None else end_weight) + self.ready_list.append(ready) + return meta + + 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.ready_list[i] + return len(idxs) + + +def _run(coro): + return asyncio.run(coro) + + +class TestStalenessSamplerSelect: + def test_select_returns_none_when_insufficient(self): + buf = FakeBuffer() + buf.add("g0", weight=5) + sampler = StalenessSampler(buf, max_staleness_versions=2) + + result = _run( + sampler.select( + current_train_weight=5, min_prompt_groups=2, max_prompt_groups=2 + ) + ) + assert result == (None, 0) + + def test_select_returns_none_on_empty_buffer(self): + buf = FakeBuffer() + sampler = StalenessSampler(buf, max_staleness_versions=2) + + result = _run( + sampler.select( + current_train_weight=5, min_prompt_groups=1, max_prompt_groups=1 + ) + ) + assert result == (None, 0) + + def test_select_filters_by_staleness_window(self): + buf = FakeBuffer() + # Weights 3, 4, 5, 2, 6 against trainer=5, max_staleness=2: + # lags = 2, 1, 0, 3 (stale), -1 (future) + for i, w in enumerate([3, 4, 5, 2, 6]): + buf.add(f"g{i}", weight=w) + sampler = StalenessSampler( + buf, max_staleness_versions=2, sample_freshest_first=True + ) + + selected, num_groups = _run( + sampler.select( + current_train_weight=5, min_prompt_groups=2, max_prompt_groups=2 + ) + ) + + assert selected is not None + # Freshest first → g2 (lag 0), g1 (lag 1) + assert selected.sample_ids == ["g2_g0", "g1_g0"] + assert num_groups == 2 + + def test_select_freshest_first_orders_by_lag(self): + buf = FakeBuffer() + for w in [3, 4, 5]: + buf.add(f"v{w}", weight=w) + sampler = StalenessSampler( + buf, max_staleness_versions=5, sample_freshest_first=True + ) + + selected, num_groups = _run( + sampler.select( + current_train_weight=6, min_prompt_groups=2, max_prompt_groups=2 + ) + ) + assert selected is not None + assert selected.sample_ids == ["v5_g0", "v4_g0"] + assert num_groups == 2 + + def test_select_fifo_orders_by_insertion(self): + buf = FakeBuffer() + for w in [3, 4, 5]: + buf.add(f"v{w}", weight=w) + sampler = StalenessSampler( + buf, max_staleness_versions=5, sample_freshest_first=False + ) + + selected, num_groups = _run( + sampler.select( + current_train_weight=6, min_prompt_groups=2, max_prompt_groups=2 + ) + ) + assert selected is not None + assert selected.sample_ids == ["v3_g0", "v4_g0"] + assert num_groups == 2 + + def test_select_skips_future_weight(self): + buf = FakeBuffer() + buf.add("now", weight=5) + buf.add("future", weight=7) + sampler = StalenessSampler(buf, max_staleness_versions=10) + + selected, num_groups = _run( + sampler.select( + current_train_weight=5, min_prompt_groups=1, max_prompt_groups=1 + ) + ) + + assert selected is not None + assert selected.sample_ids == ["now_g0"] + assert num_groups == 1 + + def test_select_concats_groups(self): + buf = FakeBuffer() + buf.add("g0", weight=5, group_size=2) + buf.add("g1", weight=5, group_size=2) + sampler = StalenessSampler(buf, max_staleness_versions=0) + + selected, num_groups = _run( + sampler.select( + current_train_weight=5, min_prompt_groups=2, max_prompt_groups=2 + ) + ) + + assert selected is not None + assert selected.sample_ids == [ + "g0_g0", + "g0_g1", + "g1_g0", + "g1_g1", + ] + # Two groups concatenated, each of size 2 → 4 sample_ids total. + assert num_groups == 2 + + def test_select_strict_on_policy_requires_exact_version(self): + buf = FakeBuffer() + for i, w in enumerate([4, 5, 5, 6]): + buf.add(f"g{i}", weight=w) + sampler = StalenessSampler(buf, max_staleness_versions=0) + + # 3 eligible (need weight=5), only have 2 + result = _run( + sampler.select( + current_train_weight=5, min_prompt_groups=3, max_prompt_groups=3 + ) + ) + assert result == (None, 0) + + # Buffer still intact: select with min=3 returned None without dropping anything. + selected, num_groups = _run( + sampler.select( + current_train_weight=5, min_prompt_groups=2, max_prompt_groups=2 + ) + ) + assert selected is not None + assert selected.sample_ids == ["g1_g0", "g2_g0"] + assert num_groups == 2 + + def test_select_drops_returned_entries_from_buffer(self): + buf = FakeBuffer() + for i, w in enumerate([5, 5, 5]): + buf.add(f"g{i}", weight=w) + sampler = StalenessSampler(buf, max_staleness_versions=0) + + first_meta, first_num_groups = _run( + sampler.select( + current_train_weight=5, min_prompt_groups=1, max_prompt_groups=1 + ) + ) + assert first_meta is not None + assert first_meta.sample_ids == ["g0_g0"] + assert first_num_groups == 1 + assert buf.start_weight_list == [5, 5] + # remove_in_dp=False; DP rows kept for trainer. + assert buf.remove_calls[-1][1] is False + + second_meta, second_num_groups = _run( + sampler.select( + current_train_weight=5, min_prompt_groups=1, max_prompt_groups=1 + ) + ) + assert second_meta is not None + assert second_meta.sample_ids == ["g1_g0"] + assert second_num_groups == 1 + + def test_select_rejects_zero_min_prompt_groups(self): + buf = FakeBuffer() + sampler = StalenessSampler(buf, max_staleness_versions=0) + with pytest.raises(ValueError): + _run( + sampler.select( + current_train_weight=0, min_prompt_groups=0, max_prompt_groups=0 + ) + ) + + def test_select_rejects_max_less_than_min(self): + buf = FakeBuffer() + for i in range(3): + buf.add(f"g{i}", weight=5) + sampler = StalenessSampler(buf, max_staleness_versions=0) + + with pytest.raises(ValueError): + _run( + sampler.select( + current_train_weight=5, min_prompt_groups=2, max_prompt_groups=1 + ) + ) + + def test_select_caps_at_max_prompt_groups(self): + buf = FakeBuffer() + for i in range(5): + buf.add(f"g{i}", weight=5) + sampler = StalenessSampler(buf, max_staleness_versions=0) + + selected, num_groups = _run( + sampler.select( + current_train_weight=5, min_prompt_groups=2, max_prompt_groups=3 + ) + ) + + assert selected is not None + # FIFO order; capped at max=3 even though 5 are eligible. + assert selected.sample_ids == ["g0_g0", "g1_g0", "g2_g0"] + assert num_groups == 3 + # The remaining two stay in the buffer. + assert buf.start_weight_list == [5, 5] + + def test_select_takes_all_available_when_between_min_and_max(self): + buf = FakeBuffer() + for i in range(3): + buf.add(f"g{i}", weight=5) + sampler = StalenessSampler(buf, max_staleness_versions=0) + + selected, num_groups = _run( + sampler.select( + current_train_weight=5, min_prompt_groups=2, max_prompt_groups=8 + ) + ) + + assert selected is not None + assert selected.sample_ids == ["g0_g0", "g1_g0", "g2_g0"] + assert num_groups == 3 + assert buf.start_weight_list == [] + + +class TestStalenessSamplerEvict: + def test_evict_removes_stale_groups(self): + buf = FakeBuffer() + # trainer=5, max_staleness=1 → lag >1 means stale (weights 0, 1, 2 stale; 4, 5 fresh) + for i, w in enumerate([0, 1, 4, 5, 2]): + buf.add(f"g{i}", weight=w) + sampler = StalenessSampler(buf, max_staleness_versions=1) + + dropped = _run(sampler.evict(current_train_weight=5)) + + assert dropped == 3 + assert buf.start_weight_list == [4, 5] + # Survivors' sample_ids + assert [m.sample_ids[0] for m in buf.meta_list] == ["g2_g0", "g3_g0"] + + def test_evict_returns_zero_when_nothing_stale(self): + buf = FakeBuffer() + for w in [4, 5]: + buf.add(f"v{w}", weight=w) + sampler = StalenessSampler(buf, max_staleness_versions=1) + + assert _run(sampler.evict(current_train_weight=5)) == 0 + assert buf.remove_calls == [] + + def test_evict_keeps_future_groups(self): + buf = FakeBuffer() + buf.add("future", weight=7) + sampler = StalenessSampler(buf, max_staleness_versions=0) + + assert _run(sampler.evict(current_train_weight=5)) == 0 + assert buf.start_weight_list == [7] + + def test_evict_drops_whole_group(self): + buf = FakeBuffer() + buf.add("stale", weight=1, group_size=4) + buf.add("fresh", weight=5, group_size=4) + sampler = StalenessSampler(buf, max_staleness_versions=1) + + dropped = _run(sampler.evict(current_train_weight=5)) + + assert dropped == 1 + assert buf.remove_calls == [([0], True)] + assert buf.start_weight_list == [5] + assert [m.sample_ids[0] for m in buf.meta_list] == ["fresh_g0"] + + +class TestStalenessSamplerInit: + def test_rejects_negative_max_staleness(self): + buf = FakeBuffer() + with pytest.raises(ValueError): + StalenessSampler(buf, max_staleness_versions=-1) + + def test_rejects_require_order_with_freshest_first(self): + buf = FakeBuffer() + with pytest.raises(ValueError): + StalenessSampler( + buf, + max_staleness_versions=0, + sample_freshest_first=True, + require_order=True, + ) + + +class TestStalenessSamplerReady: + def test_default_mode_skips_unready_slots(self): + buf = FakeBuffer() + buf.add("g0", weight=5, ready=False) + buf.add("g1", weight=5, ready=True) + sampler = StalenessSampler(buf, max_staleness_versions=0) + + selected, num_groups = _run( + sampler.select( + current_train_weight=5, min_prompt_groups=1, max_prompt_groups=1 + ) + ) + + assert selected is not None + assert selected.sample_ids == ["g1_g0"] + assert num_groups == 1 + + def test_default_mode_waits_when_too_few_ready(self): + buf = FakeBuffer() + buf.add("g0", weight=5, ready=False) + buf.add("g1", weight=5, ready=True) + sampler = StalenessSampler(buf, max_staleness_versions=0) + + result = _run( + sampler.select( + current_train_weight=5, min_prompt_groups=2, max_prompt_groups=2 + ) + ) + assert result == (None, 0) + + +class TestStalenessSamplerRequireOrder: + def test_consumes_oldest_batch_first(self): + buf = FakeBuffer() + # Two complete batches: v=4 then v=5; require_order must take v=4 first. + for i, w in enumerate((4, 4, 5, 5)): + buf.add(f"v{w}_{i}", weight=w) + sampler = StalenessSampler(buf, max_staleness_versions=1, require_order=True) + + selected, num_groups = _run( + sampler.select( + current_train_weight=5, min_prompt_groups=2, max_prompt_groups=2 + ) + ) + + assert selected is not None + # Insertion-order FIFO inside the oldest batch. + assert selected.sample_ids == ["v4_0_g0", "v4_1_g0"] + assert num_groups == 2 + assert buf.start_weight_list == [5, 5] + + def test_waits_when_oldest_batch_partially_ready(self): + buf = FakeBuffer() + # Oldest batch v=4 has 1 ready + 1 unready; v=5 batch is fully ready. + # require_order must NOT skip ahead to v=5. + buf.add("v4_a", weight=4, ready=True) + buf.add("v4_b", weight=4, ready=False) + buf.add("v5_a", weight=5, ready=True) + buf.add("v5_b", weight=5, ready=True) + sampler = StalenessSampler(buf, max_staleness_versions=1, require_order=True) + + result = _run( + sampler.select( + current_train_weight=5, min_prompt_groups=2, max_prompt_groups=2 + ) + ) + assert result == (None, 0) + # Buffer untouched: nothing removed. + assert buf.start_weight_list == [4, 4, 5, 5] + assert buf.ready_list == [True, False, True, True] + + def test_returns_none_when_oldest_batch_not_filled(self): + buf = FakeBuffer() + buf.add("v4_a", weight=4, ready=True) + # Only 1 ready in oldest batch; need 2. + sampler = StalenessSampler(buf, max_staleness_versions=1, require_order=True) + + result = _run( + sampler.select( + current_train_weight=5, min_prompt_groups=2, max_prompt_groups=2 + ) + ) + assert result == (None, 0) + + def test_ignores_future_versions_when_picking_target(self): + buf = FakeBuffer() + # Trainer at 5, staleness 1: window is [4, 5]; v=7 (future) must not + # become the oldest target. + buf.add("v7", weight=7, ready=True) + buf.add("v5_a", weight=5, ready=True) + buf.add("v5_b", weight=5, ready=True) + sampler = StalenessSampler(buf, max_staleness_versions=1, require_order=True) + + selected, num_groups = _run( + sampler.select( + current_train_weight=5, min_prompt_groups=2, max_prompt_groups=2 + ) + ) + + assert selected is not None + assert selected.sample_ids == ["v5_a_g0", "v5_b_g0"] + assert num_groups == 2 + assert buf.start_weight_list == [7] From 056c12e84a6fedbf9790b0f4ae7fed830c9d488b Mon Sep 17 00:00:00 2001 From: Yuki Huang Date: Fri, 17 Jul 2026 09:45:59 -0700 Subject: [PATCH 02/17] refactor(sc): extract SC helpers to single_controller_utils.utils, rewire StalenessSampler init Signed-off-by: Yuki Huang --- nemo_rl/algorithms/single_controller.py | 90 +++------ .../single_controller_utils/utils.py | 189 ++++++++++++++++++ 2 files changed, 221 insertions(+), 58 deletions(-) create mode 100644 nemo_rl/algorithms/single_controller_utils/utils.py diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 5ae67f11b4..a4280385d9 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -50,13 +50,15 @@ import ray import torch from pydantic import BaseModel, Field -from tensordict import TensorDict +from nemo_rl.algorithms.async_utils.staleness_sampler import StalenessSampler 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 @@ -245,6 +247,7 @@ def __init__( advantage_estimator: Any | None = None, rollout_manager: Any = None, dataloader: Any = None, + tq_buffer: Any = None, ) -> None: self._cfg = cfg self._prompts = prompts @@ -256,6 +259,7 @@ def __init__( self._advantage_estimator = advantage_estimator self._rollout_manager = rollout_manager self._dataloader = dataloader + self._buffer = tq_buffer # Built here, not on the driver: Logger backends (wandb/tb/...) hold # _thread.lock that Ray can't cloudpickle into the actor. @@ -328,7 +332,13 @@ def __init__( # 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) + + self._sampler = StalenessSampler( + self._buffer, + max_staleness_versions=cfg.max_weight_staleness_versions, + require_order=not cfg.over_sampling, + force_in_order=cfg.force_in_order, + ) # ── asyncio state ────────────────────────────────────────────────── # Gate: cleared during _sync_weights, set when generation may proceed @@ -354,6 +364,12 @@ 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 @@ -786,11 +802,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) @@ -798,18 +814,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, ) @@ -826,7 +842,7 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> KVBatchMeta: "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}, ), @@ -848,45 +864,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)) - - -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..e202d99c4d --- /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_pump into step scalars. + + Args: + rewards: One tensor per advantage_pump 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]) From 8ea11d27e07707bd4e28d08e9241c98dcd36695e Mon Sep 17 00:00:00 2001 From: Yuki Huang Date: Fri, 17 Jul 2026 09:54:42 -0700 Subject: [PATCH 03/17] docs(sc): fix stale references in module/pump/warn docstrings Signed-off-by: Yuki Huang --- nemo_rl/algorithms/single_controller.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index a4280385d9..99c1320328 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 @@ -182,7 +182,7 @@ def warn_if_staleness_window_below_minibatches( 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 — + sampler skips them and ``sampler.evict`` 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 @@ -564,8 +564,8 @@ async def _train_pump(self) -> None: 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_pump(train_meta). - 4. trainer.train_microbatch_from_meta + finish_train_step. + 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. """ From 67ad762a107a25657dc24edeb170f44834011efb Mon Sep 17 00:00:00 2001 From: Yuki Huang Date: Fri, 17 Jul 2026 10:04:53 -0700 Subject: [PATCH 04/17] test(sc): cover StalenessSampler force_in_order path Signed-off-by: Yuki Huang --- .../test_staleness_sampler.py | 63 +++++++++++++++++++ 1 file changed, 63 insertions(+) diff --git a/tests/unit/single_controller/test_staleness_sampler.py b/tests/unit/single_controller/test_staleness_sampler.py index ce15f34acf..5e53061638 100644 --- a/tests/unit/single_controller/test_staleness_sampler.py +++ b/tests/unit/single_controller/test_staleness_sampler.py @@ -32,6 +32,7 @@ def __init__(self, partition_id: str = "rollout_data") -> None: 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]] = [] @@ -42,6 +43,7 @@ def add( group_size: int = 1, ready: bool = True, end_weight: int | None = None, + target_step: int | None = None, ) -> KVBatchMeta: sample_ids = [f"{group_id}_g{i}" for i in range(group_size)] meta = KVBatchMeta( @@ -53,6 +55,7 @@ def add( self.meta_list.append(meta if ready else None) self.start_weight_list.append(weight) self.end_weight_list.append(weight if end_weight is None else end_weight) + self.target_step_list.append(target_step) self.ready_list.append(ready) return meta @@ -62,6 +65,7 @@ async def remove(self, idxs: list[int], remove_in_dp: bool) -> int: 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) @@ -464,3 +468,62 @@ def test_ignores_future_versions_when_picking_target(self): assert selected.sample_ids == ["v5_a_g0", "v5_b_g0"] assert num_groups == 2 assert buf.start_weight_list == [7] + + +class TestStalenessSamplerForceInOrder: + def test_selects_only_groups_whose_target_step_matches_current(self): + buf = FakeBuffer() + # Two matching-target groups and one future-target group. + buf.add("g_match_a", weight=5, target_step=5) + buf.add("g_match_b", weight=5, target_step=5) + buf.add("g_future", weight=5, target_step=7) + sampler = StalenessSampler(buf, max_staleness_versions=0, force_in_order=True) + + selected, num_groups = _run( + sampler.select( + current_train_weight=5, min_prompt_groups=1, max_prompt_groups=4 + ) + ) + + assert selected is not None + assert selected.sample_ids == ["g_match_a_g0", "g_match_b_g0"] + assert num_groups == 2 + # The future-target group survives; only the two matches were removed. + assert buf.target_step_list == [7] + assert buf.start_weight_list == [5] + + def test_ignores_stale_target_step(self): + buf = FakeBuffer() + # Trainer at 5; a slot with target_step=4 must NOT satisfy force_in_order. + buf.add("g_stale", weight=4, target_step=4) + sampler = StalenessSampler( + buf, + # Wide staleness window: force_in_order must ignore it entirely. + max_staleness_versions=10, + force_in_order=True, + ) + + result = _run( + sampler.select( + current_train_weight=5, min_prompt_groups=1, max_prompt_groups=4 + ) + ) + + assert result == (None, 0) + # Buffer untouched. + assert buf.target_step_list == [4] + + def test_ignores_unready_group_with_matching_target(self): + buf = FakeBuffer() + buf.add("g_unready", weight=5, target_step=5, ready=False) + sampler = StalenessSampler(buf, max_staleness_versions=0, force_in_order=True) + + result = _run( + sampler.select( + current_train_weight=5, min_prompt_groups=1, max_prompt_groups=4 + ) + ) + + assert result == (None, 0) + # Buffer untouched. + assert buf.ready_list == [False] From e990a7c714556be3dfbc59aa82c27bb4743a9b53 Mon Sep 17 00:00:00 2001 From: Yuki Huang Date: Fri, 17 Jul 2026 10:07:14 -0700 Subject: [PATCH 05/17] =?UTF-8?q?refactor(sc):=20rename=20StalenessSampler?= =?UTF-8?q?.require=5Forder=20=E2=86=92=20strict=5Fweight=5Ffifo?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Yuki Huang --- .../async_utils/staleness_sampler.py | 12 ++++----- nemo_rl/algorithms/single_controller.py | 2 +- .../test_staleness_sampler.py | 26 ++++++++++++------- 3 files changed, 24 insertions(+), 16 deletions(-) diff --git a/nemo_rl/algorithms/async_utils/staleness_sampler.py b/nemo_rl/algorithms/async_utils/staleness_sampler.py index 144dabdd3d..c5a2f443a0 100644 --- a/nemo_rl/algorithms/async_utils/staleness_sampler.py +++ b/nemo_rl/algorithms/async_utils/staleness_sampler.py @@ -25,7 +25,7 @@ class StalenessSampler: buffer: Shared TQReplayBuffer holding the candidate slots. max_staleness_versions: Max weight-version gap a sample may have from the trainer. sample_freshest_first: Prefer smallest lag when picking from the in-window set. - require_order: Take only from the oldest in-window weight_version and wait for its batch to fill. + strict_weight_fifo: Consume in strict weight_version FIFO — take only from the oldest in-window weight_version and wait for its batch to fill. force_in_order: Match each slot's target_step against current_train_weight, ignoring the window; mirrors legacy async_grpo target_weight semantics. """ @@ -34,7 +34,7 @@ def __init__( buffer: TQReplayBuffer, max_staleness_versions: int, sample_freshest_first: bool = False, - require_order: bool = False, + strict_weight_fifo: bool = False, force_in_order: bool = False, ) -> None: if max_staleness_versions < 0: @@ -42,14 +42,14 @@ def __init__( f"max_staleness_versions must be non-negative, got " f"{max_staleness_versions}" ) - if require_order and sample_freshest_first: + if strict_weight_fifo and sample_freshest_first: raise ValueError( - "require_order and sample_freshest_first are mutually exclusive" + "strict_weight_fifo and sample_freshest_first are mutually exclusive" ) self._buffer = buffer self.max_staleness_versions = max_staleness_versions self.sample_freshest_first = sample_freshest_first - self.require_order = require_order + self.strict_weight_fifo = strict_weight_fifo self.force_in_order = force_in_order async def select( @@ -93,7 +93,7 @@ async def select( min_valid_version = max( 0, current_train_weight - self.max_staleness_versions ) - if self.require_order: + if self.strict_weight_fifo: in_window = [ weight for weight in self._buffer.start_weight_list diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 99c1320328..d209cbe5e5 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -336,7 +336,7 @@ def __init__( self._sampler = StalenessSampler( self._buffer, max_staleness_versions=cfg.max_weight_staleness_versions, - require_order=not cfg.over_sampling, + strict_weight_fifo=not cfg.over_sampling, force_in_order=cfg.force_in_order, ) diff --git a/tests/unit/single_controller/test_staleness_sampler.py b/tests/unit/single_controller/test_staleness_sampler.py index 5e53061638..df6a27fae9 100644 --- a/tests/unit/single_controller/test_staleness_sampler.py +++ b/tests/unit/single_controller/test_staleness_sampler.py @@ -354,14 +354,14 @@ def test_rejects_negative_max_staleness(self): with pytest.raises(ValueError): StalenessSampler(buf, max_staleness_versions=-1) - def test_rejects_require_order_with_freshest_first(self): + def test_rejects_strict_weight_fifo_with_freshest_first(self): buf = FakeBuffer() with pytest.raises(ValueError): StalenessSampler( buf, max_staleness_versions=0, sample_freshest_first=True, - require_order=True, + strict_weight_fifo=True, ) @@ -396,13 +396,15 @@ def test_default_mode_waits_when_too_few_ready(self): assert result == (None, 0) -class TestStalenessSamplerRequireOrder: +class TestStalenessSamplerStrictWeightFifo: def test_consumes_oldest_batch_first(self): buf = FakeBuffer() - # Two complete batches: v=4 then v=5; require_order must take v=4 first. + # Two complete batches: v=4 then v=5; strict_weight_fifo must take v=4 first. for i, w in enumerate((4, 4, 5, 5)): buf.add(f"v{w}_{i}", weight=w) - sampler = StalenessSampler(buf, max_staleness_versions=1, require_order=True) + sampler = StalenessSampler( + buf, max_staleness_versions=1, strict_weight_fifo=True + ) selected, num_groups = _run( sampler.select( @@ -419,12 +421,14 @@ def test_consumes_oldest_batch_first(self): def test_waits_when_oldest_batch_partially_ready(self): buf = FakeBuffer() # Oldest batch v=4 has 1 ready + 1 unready; v=5 batch is fully ready. - # require_order must NOT skip ahead to v=5. + # strict_weight_fifo must NOT skip ahead to v=5. buf.add("v4_a", weight=4, ready=True) buf.add("v4_b", weight=4, ready=False) buf.add("v5_a", weight=5, ready=True) buf.add("v5_b", weight=5, ready=True) - sampler = StalenessSampler(buf, max_staleness_versions=1, require_order=True) + sampler = StalenessSampler( + buf, max_staleness_versions=1, strict_weight_fifo=True + ) result = _run( sampler.select( @@ -440,7 +444,9 @@ def test_returns_none_when_oldest_batch_not_filled(self): buf = FakeBuffer() buf.add("v4_a", weight=4, ready=True) # Only 1 ready in oldest batch; need 2. - sampler = StalenessSampler(buf, max_staleness_versions=1, require_order=True) + sampler = StalenessSampler( + buf, max_staleness_versions=1, strict_weight_fifo=True + ) result = _run( sampler.select( @@ -456,7 +462,9 @@ def test_ignores_future_versions_when_picking_target(self): buf.add("v7", weight=7, ready=True) buf.add("v5_a", weight=5, ready=True) buf.add("v5_b", weight=5, ready=True) - sampler = StalenessSampler(buf, max_staleness_versions=1, require_order=True) + sampler = StalenessSampler( + buf, max_staleness_versions=1, strict_weight_fifo=True + ) selected, num_groups = _run( sampler.select( From 401196cb8561b8ea412b8198d7bfadc1e6866148 Mon Sep 17 00:00:00 2001 From: Yuki Huang Date: Fri, 17 Jul 2026 10:08:31 -0700 Subject: [PATCH 06/17] test(sc): drop redundant StalenessSampler.select cases subsumed by neighbours Signed-off-by: Yuki Huang --- .../test_staleness_sampler.py | 28 ------------------- 1 file changed, 28 deletions(-) diff --git a/tests/unit/single_controller/test_staleness_sampler.py b/tests/unit/single_controller/test_staleness_sampler.py index df6a27fae9..a0ef59e103 100644 --- a/tests/unit/single_controller/test_staleness_sampler.py +++ b/tests/unit/single_controller/test_staleness_sampler.py @@ -87,17 +87,6 @@ def test_select_returns_none_when_insufficient(self): ) assert result == (None, 0) - def test_select_returns_none_on_empty_buffer(self): - buf = FakeBuffer() - sampler = StalenessSampler(buf, max_staleness_versions=2) - - result = _run( - sampler.select( - current_train_weight=5, min_prompt_groups=1, max_prompt_groups=1 - ) - ) - assert result == (None, 0) - def test_select_filters_by_staleness_window(self): buf = FakeBuffer() # Weights 3, 4, 5, 2, 6 against trainer=5, max_staleness=2: @@ -119,23 +108,6 @@ def test_select_filters_by_staleness_window(self): assert selected.sample_ids == ["g2_g0", "g1_g0"] assert num_groups == 2 - def test_select_freshest_first_orders_by_lag(self): - buf = FakeBuffer() - for w in [3, 4, 5]: - buf.add(f"v{w}", weight=w) - sampler = StalenessSampler( - buf, max_staleness_versions=5, sample_freshest_first=True - ) - - selected, num_groups = _run( - sampler.select( - current_train_weight=6, min_prompt_groups=2, max_prompt_groups=2 - ) - ) - assert selected is not None - assert selected.sample_ids == ["v5_g0", "v4_g0"] - assert num_groups == 2 - def test_select_fifo_orders_by_insertion(self): buf = FakeBuffer() for w in [3, 4, 5]: From 7126ef7c9df717c513c917779e7111818324766a Mon Sep 17 00:00:00 2001 From: Yuki Huang Date: Fri, 17 Jul 2026 22:17:16 -0700 Subject: [PATCH 07/17] pyrefly Signed-off-by: Yuki Huang --- pyrefly.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyrefly.toml b/pyrefly.toml index a70ea16100..fa5696e661 100644 --- a/pyrefly.toml +++ b/pyrefly.toml @@ -57,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", From 25a8c5f03ce13d71774a2defef476ec7412daf2e Mon Sep 17 00:00:00 2001 From: Yuki Huang Date: Fri, 17 Jul 2026 22:29:26 -0700 Subject: [PATCH 08/17] refactor(sc): drop unused claim-state config fields and init attrs Signed-off-by: Yuki Huang --- nemo_rl/algorithms/single_controller.py | 14 -------------- tests/unit/single_controller/test_train_pump.py | 3 --- 2 files changed, 17 deletions(-) diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index d209cbe5e5..be569861d9 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -43,7 +43,6 @@ import asyncio import time -from contextlib import nullcontext from functools import partial from typing import Any, Literal, Optional @@ -144,9 +143,6 @@ class SingleControllerConfig(BaseModel, extra="allow"): # 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 @@ -373,8 +369,6 @@ def __init__( # 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=" @@ -385,14 +379,6 @@ def __init__( 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]: diff --git a/tests/unit/single_controller/test_train_pump.py b/tests/unit/single_controller/test_train_pump.py index 1763e0efca..5df6a51c8e 100644 --- a/tests/unit/single_controller/test_train_pump.py +++ b/tests/unit/single_controller/test_train_pump.py @@ -232,9 +232,6 @@ def test_train_pump_dry_run(ray_cluster, tmp_path): 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, From 4effa0f61fbccecafcf2d9e8d63be3d83d2116c5 Mon Sep 17 00:00:00 2001 From: Yuki Huang Date: Sat, 18 Jul 2026 00:29:02 -0700 Subject: [PATCH 09/17] test(sc): add real-Megatron e2e train_pump test, fix stale _train_pump attrs Signed-off-by: Yuki Huang --- nemo_rl/algorithms/single_controller.py | 39 +- tests/unit/single_controller/_dp_fakes.py | 158 ++++++ .../single_controller/test_rollout_pump.py | 115 +---- .../unit/single_controller/test_train_pump.py | 464 +++++++++--------- 4 files changed, 429 insertions(+), 347 deletions(-) create mode 100644 tests/unit/single_controller/_dp_fakes.py diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index be569861d9..e99b7872ef 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -253,9 +253,12 @@ def __init__( 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. @@ -555,20 +558,23 @@ async def _train_pump(self) -> None: 5. dp_client.clear_samples on consumed sample_ids; release _buffer_capacity per dropped group, then sync. """ - adv_cfg = self._advantage_cfg - grpo_cfg = self._master_config.grpo + 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 = adv_cfg.policy_logprobs_field is not None - reference_logprobs_required = adv_cfg.reference_logprobs_field is not None + prev_logprobs_required = self._cfg.advantage_policy_logprobs_field is not None + reference_logprobs_required = ( + self._cfg.advantage_reference_logprobs_field is not None + ) - while self._train_steps < grpo_cfg["max_num_steps"]: + while self._train_steps < self._cfg.max_train_steps: groups_dispatched = 0 min_sample_version = None step_open = False with self._timer.time("total_step_time"): - while groups_dispatched < grpo_cfg["num_prompts_per_step"]: + while groups_dispatched < target_groups: # Wait for a selectable batch with self._timer.time("exposed_generation"): await asyncio.sleep(0) @@ -586,11 +592,9 @@ async def _train_pump(self) -> None: self._buffer_capacity.release() # Select a batch - max_prompt_groups = ( - grpo_cfg["num_prompts_per_step"] - groups_dispatched - ) + max_prompt_groups = target_groups - groups_dispatched min_prompt_groups = min( - self._async_cfg.min_prompt_groups_per_batch, + self._cfg.min_groups_per_batch, max_prompt_groups, ) train_meta, num_groups = await self._sampler.select( @@ -662,7 +666,7 @@ async def _train_pump(self) -> None: await self._call_dp( "clear_samples", sample_ids=list(train_meta.sample_ids), - partition_id=self._partition_id, + partition_id=self._cfg.partition_id, ) groups_dispatched += num_groups @@ -693,9 +697,12 @@ async def _train_pump(self) -> None: ) # type: ignore total_time = timing_metrics.get("total_step_time", 0.0) - cluster_cfg = self._master_config.cluster - total_num_gpus = cluster_cfg["num_nodes"] * cluster_cfg["gpus_per_node"] - if total_time > 0 and "global_valid_toks" in step_metrics: + 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 ) @@ -728,7 +735,7 @@ async def _train_pump(self) -> None: # generated with; lag = current trainer version - oldest sample version. lag = self._trainer_version - min_sample_version # type: ignore print( - f"train step {self._train_steps}/{grpo_cfg['max_num_steps']} " + f"train step {self._train_steps}/{self._cfg.max_train_steps} " f"trainer_v={self._trainer_version} " f"lag={lag} ", flush=True, diff --git a/tests/unit/single_controller/_dp_fakes.py b/tests/unit/single_controller/_dp_fakes.py new file mode 100644 index 0000000000..3cf117bfaf --- /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.add 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..ca21a12735 100644 --- a/tests/unit/single_controller/test_rollout_pump.py +++ b/tests/unit/single_controller/test_rollout_pump.py @@ -24,14 +24,12 @@ import pytest import ray import torch -from tensordict import TensorDict from nemo_rl.algorithms.async_utils.replay_buffer import TQReplayBuffer 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,112 +44,12 @@ 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( @@ -428,6 +326,7 @@ def test_rollout_pump_writes_expected_tq_data( advantage_estimator=None, rollout_manager=rollout_manager, dataloader=dataloader, + tq_buffer=tq_buffer, ) vllm_generation.prepare_for_generation() diff --git a/tests/unit/single_controller/test_train_pump.py b/tests/unit/single_controller/test_train_pump.py index 5df6a51c8e..f15cea9849 100644 --- a/tests/unit/single_controller/test_train_pump.py +++ b/tests/unit/single_controller/test_train_pump.py @@ -12,56 +12,159 @@ # 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 +from types import SimpleNamespace + import pytest import ray import torch from tensordict import TensorDict +from nemo_rl.algorithms.async_utils.replay_buffer import TQReplayBuffer 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 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 - def peek_count(self, partition_id: str) -> int: - return len(self._partitions[partition_id].rows) +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], + ) -@ray.remote(num_cpus=0) +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() + + +class _FakeAdvEstimator: + """Passthrough advantage: returns rewards unchanged.""" + + def compute_advantage( + self, + prompt_ids: torch.Tensor, + rewards: torch.Tensor, + mask: torch.Tensor, + repeated_batch: dict, + **kwargs, + ) -> torch.Tensor: + return rewards.detach().clone() + + +@ray.remote(num_cpus=0) # pragma: no cover class _CallLog: """Ordered append-only log the fakes below write to.""" @@ -76,29 +179,8 @@ def get(self) -> list[tuple[str, dict]]: return list(self._entries) -class _FakeTrainer: - """Sync driver-side TQPolicy stand-in.""" - - def __init__(self, log_handle) -> 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: - ray.get( - self._log.record.remote( - "train_microbatches_from_meta", {"meta_size": int(meta.size)} - ) - ) - - def finish_train_step(self) -> dict: - ray.get(self._log.record.remote("finish_train_step", {})) - return {"loss": 0.0} - - 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,191 +189,127 @@ 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.""" +@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.""" + num_generations = 2 + num_prompts = 2 + seq_len = 16 + train_gbs = num_prompts * num_generations + + 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"] = num_generations + tokenizer = get_tokenizer(policy_cfg["tokenizer"]) + policy_cfg["generation"] = configure_generation_config( + policy_cfg["generation"], tokenizer + ) - def __init__(self, log_handle) -> None: - self._log = log_handle + 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, + ) - 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() + try: + # Reuse TQPolicy's driver-side dp_client for SC — same controller, + # no double-bootstrap. + dp_client = trainer.dp_client + # Register the partition once with the full field union. + dp_client.register_partition( + partition_id=_PARTITION_ID, + fields=_REGISTERED_FIELDS, + num_samples=train_gbs * 4, + consumer_tasks=["prev_lp", "ref_lp", "train"], + ) -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, + # Real TQReplayBuffer, pre-populated with 2 ready groups so + # StalenessSampler.select fires immediately. + tq_buffer = TQReplayBuffer( + dp_client, partition_id=_PARTITION_ID, - fields=fields, - tags=tags, + pad_value_dict={"input_ids": int(tokenizer.pad_token_id or 0)}, ) - ) - 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( + for g in range(num_prompts): + meta = _populate_group( + dp_client, + group_uuid=f"prompt{g}", + group_size=num_generations, + seq_len=seq_len, + weight_version=0, + ) + _prepopulate_buffer(tq_buffer, meta, weight_version=0) + + 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( + max_weight_staleness_versions=1, + min_groups_per_batch=num_prompts, + target_groups_per_step=num_prompts, + group_size=num_generations, + batch_selection_strategy="staleness_window", + max_inflight_prompts=num_prompts, + max_buffered_rollouts=num_prompts, + max_train_steps=1, + max_num_epochs=None, + max_rollout_prompts=num_prompts, + over_sampling=True, + train_global_batch_size=train_gbs, partition_id=_PARTITION_ID, - fields=[*_ROLLOUT_FIELDS, _ADV_FIELD], - num_samples=group_size * num_groups * 4, - consumer_tasks=["train"], + 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, + }, ) - ) - 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, + + ctrl = SingleControllerActor.remote( + cfg=cfg, + prompts=[], + dp_client_handle=dp_client, + gen_handle=None, + trainer_handle=trainer, + weight_synchronizer=weight_sync, + loss_fn=SimpleLossFn(), + advantage_estimator=adv_est, + rollout_manager=rollout_manager, + dataloader=None, + tq_buffer=tq_buffer, ) - 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, - 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, - }, - ) + # One outer step: sampler.select → advantage stage → begin/microbatches/finish → sync. + ray.get(ctrl._train_pump.remote()) - 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, - ) + state = ray.get(ctrl.ping.remote()) + assert state["train_steps"] == 1 + assert state["trainer_version"] == 1 + + entries = ray.get(log.get.remote()) + sync_versions = [p["version"] for k, p in entries if k == "sync_weights"] + assert sync_versions == [1] - # _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 + finally: + trainer.shutdown() From 71fe4139fed76846392310cbea8aff0263150f8b Mon Sep 17 00:00:00 2001 From: Yuki Huang Date: Sat, 18 Jul 2026 01:35:34 -0700 Subject: [PATCH 10/17] refactor(sc): drop unused max_rollout_prompts config field and prompts init arg Signed-off-by: Yuki Huang --- nemo_rl/algorithms/single_controller.py | 8 ++------ tests/unit/single_controller/test_rollout_pump.py | 1 - tests/unit/single_controller/test_train_pump.py | 2 -- 3 files changed, 2 insertions(+), 9 deletions(-) diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index e99b7872ef..1beb4a4341 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -126,11 +126,9 @@ class SingleControllerConfig(BaseModel, extra="allow"): # 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 @@ -234,7 +232,6 @@ class SingleControllerActor: def __init__( self, cfg: SingleControllerConfig, - prompts: list[str], dp_client_handle: Any, gen_handle: Any, trainer_handle: Any, @@ -246,7 +243,6 @@ def __init__( 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 diff --git a/tests/unit/single_controller/test_rollout_pump.py b/tests/unit/single_controller/test_rollout_pump.py index ca21a12735..008ddc45eb 100644 --- a/tests/unit/single_controller/test_rollout_pump.py +++ b/tests/unit/single_controller/test_rollout_pump.py @@ -317,7 +317,6 @@ 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(), diff --git a/tests/unit/single_controller/test_train_pump.py b/tests/unit/single_controller/test_train_pump.py index f15cea9849..2b5d7d72a7 100644 --- a/tests/unit/single_controller/test_train_pump.py +++ b/tests/unit/single_controller/test_train_pump.py @@ -267,7 +267,6 @@ def test_train_pump_drives_mcore_training_step( max_buffered_rollouts=num_prompts, max_train_steps=1, max_num_epochs=None, - max_rollout_prompts=num_prompts, over_sampling=True, train_global_batch_size=train_gbs, partition_id=_PARTITION_ID, @@ -288,7 +287,6 @@ def test_train_pump_drives_mcore_training_step( ctrl = SingleControllerActor.remote( cfg=cfg, - prompts=[], dp_client_handle=dp_client, gen_handle=None, trainer_handle=trainer, From 4fe817d0a2447be50d67df191e1d0044c33dee35 Mon Sep 17 00:00:00 2001 From: Yuki Huang Date: Sat, 18 Jul 2026 02:05:18 -0700 Subject: [PATCH 11/17] test(sc): multi-step train_pump, cover force_in_order weight-outside-window Signed-off-by: Yuki Huang --- .../test_staleness_sampler.py | 17 +++++++ .../unit/single_controller/test_train_pump.py | 45 ++++++++++--------- 2 files changed, 41 insertions(+), 21 deletions(-) diff --git a/tests/unit/single_controller/test_staleness_sampler.py b/tests/unit/single_controller/test_staleness_sampler.py index a0ef59e103..69f580caf7 100644 --- a/tests/unit/single_controller/test_staleness_sampler.py +++ b/tests/unit/single_controller/test_staleness_sampler.py @@ -493,6 +493,23 @@ def test_ignores_stale_target_step(self): # Buffer untouched. assert buf.target_step_list == [4] + def test_selects_when_target_matches_even_if_weight_outside_window(self): + buf = FakeBuffer() + # start_weight=100 is well outside any realistic staleness window, + # but target_step matches current — force_in_order must ignore the window. + buf.add("g_match_stale_weight", weight=100, target_step=5) + sampler = StalenessSampler(buf, max_staleness_versions=0, force_in_order=True) + + selected, num_groups = _run( + sampler.select( + current_train_weight=5, min_prompt_groups=1, max_prompt_groups=4 + ) + ) + + assert selected is not None + assert selected.sample_ids == ["g_match_stale_weight_g0"] + assert num_groups == 1 + def test_ignores_unready_group_with_matching_target(self): buf = FakeBuffer() buf.add("g_unready", weight=5, target_step=5, ready=False) diff --git a/tests/unit/single_controller/test_train_pump.py b/tests/unit/single_controller/test_train_pump.py index 2b5d7d72a7..5c50c16871 100644 --- a/tests/unit/single_controller/test_train_pump.py +++ b/tests/unit/single_controller/test_train_pump.py @@ -197,14 +197,14 @@ def test_train_pump_drives_mcore_training_step( 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 - seq_len = 16 - train_gbs = num_prompts * num_generations 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"] = num_generations + policy_cfg["train_micro_batch_size"] = train_gbs tokenizer = get_tokenizer(policy_cfg["tokenizer"]) policy_cfg["generation"] = configure_generation_config( policy_cfg["generation"], tokenizer @@ -226,30 +226,33 @@ def test_train_pump_drives_mcore_training_step( # no double-bootstrap. dp_client = trainer.dp_client - # Register the partition once with the full field union. + # 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=_REGISTERED_FIELDS, - num_samples=train_gbs * 4, + num_samples=train_steps * train_gbs, consumer_tasks=["prev_lp", "ref_lp", "train"], ) - # Real TQReplayBuffer, pre-populated with 2 ready groups so - # StalenessSampler.select fires immediately. + # 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, pad_value_dict={"input_ids": int(tokenizer.pad_token_id or 0)}, ) - for g in range(num_prompts): - meta = _populate_group( - dp_client, - group_uuid=f"prompt{g}", - group_size=num_generations, - seq_len=seq_len, - weight_version=0, - ) - _prepopulate_buffer(tq_buffer, meta, 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) @@ -265,7 +268,7 @@ def test_train_pump_drives_mcore_training_step( batch_selection_strategy="staleness_window", max_inflight_prompts=num_prompts, max_buffered_rollouts=num_prompts, - max_train_steps=1, + max_train_steps=train_steps, max_num_epochs=None, over_sampling=True, train_global_batch_size=train_gbs, @@ -298,16 +301,16 @@ def test_train_pump_drives_mcore_training_step( tq_buffer=tq_buffer, ) - # One outer step: sampler.select → advantage stage → begin/microbatches/finish → sync. + # 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"] == 1 - assert state["trainer_version"] == 1 + 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 == [1] + assert sync_versions == list(range(1, train_steps + 1)) finally: trainer.shutdown() From f721dc07679c05d134b13e94dc57bdd728551f5d Mon Sep 17 00:00:00 2001 From: ruit Date: Tue, 21 Jul 2026 22:39:46 -0700 Subject: [PATCH 12/17] test(sc): fix stale replay buffer method reference Signed-off-by: ruit --- tests/unit/single_controller/_dp_fakes.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/single_controller/_dp_fakes.py b/tests/unit/single_controller/_dp_fakes.py index 3cf117bfaf..1f26ea6d77 100644 --- a/tests/unit/single_controller/_dp_fakes.py +++ b/tests/unit/single_controller/_dp_fakes.py @@ -25,7 +25,7 @@ from nemo_rl.data_plane.adapters.noop import NoOpDataPlaneClient _PARTITION_ID = "rollout_data" -# TQReplayBuffer.add writes these training-row fields per prompt; +# TQReplayBuffer.commit writes these training-row fields per prompt; # _advantage_stage writes ``advantages`` back on top. _BULK_FIELDS = [ "input_ids", From 3af580e5b09fb7e28e269507cd2d8129ab99107c Mon Sep 17 00:00:00 2001 From: ruit Date: Wed, 22 Jul 2026 06:31:20 -0700 Subject: [PATCH 13/17] fix(sc): address train path review feedback Signed-off-by: ruit --- .../algorithms/async_utils/staleness_sampler.py | 3 +++ nemo_rl/algorithms/single_controller.py | 7 ++++--- .../algorithms/single_controller_utils/utils.py | 4 ++-- .../unit/single_controller/test_rollout_pump.py | 16 ++++++++-------- tests/unit/single_controller/test_train_pump.py | 2 +- 5 files changed, 18 insertions(+), 14 deletions(-) diff --git a/nemo_rl/algorithms/async_utils/staleness_sampler.py b/nemo_rl/algorithms/async_utils/staleness_sampler.py index c5a2f443a0..86251dbfba 100644 --- a/nemo_rl/algorithms/async_utils/staleness_sampler.py +++ b/nemo_rl/algorithms/async_utils/staleness_sampler.py @@ -69,6 +69,9 @@ async def select( current_train_weight: Current trainer weight version. min_prompt_groups: Minimum groups required; returns (None, 0) below this. max_prompt_groups: Cap on groups returned when the threshold is met. + Selection is greedy without waiting: it returns all currently + eligible groups up to this cap (never deliberately fewer, and it + does not wait for the cap to fill). Returns: meta: Concatenated KVBatchMeta, or None if not enough groups. diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 1beb4a4341..c10f44e7f8 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -237,8 +237,8 @@ def __init__( 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: @@ -565,6 +565,7 @@ async def _train_pump(self) -> None: ) while self._train_steps < self._cfg.max_train_steps: + version_during_step = self._trainer_version groups_dispatched = 0 min_sample_version = None step_open = False @@ -728,8 +729,8 @@ async def _train_pump(self) -> None: self._timer.reset() # min sample version refers to the version each consumed sample was - # generated with; lag = current trainer version - oldest sample version. - lag = self._trainer_version - min_sample_version # type: ignore + # 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} " diff --git a/nemo_rl/algorithms/single_controller_utils/utils.py b/nemo_rl/algorithms/single_controller_utils/utils.py index e202d99c4d..4844898170 100644 --- a/nemo_rl/algorithms/single_controller_utils/utils.py +++ b/nemo_rl/algorithms/single_controller_utils/utils.py @@ -95,10 +95,10 @@ def reduce_advantage_pump_metrics( masked_advantages: list[torch.Tensor], sequence_lengths: list[int], ) -> dict[str, float]: - """Reduce per-step accumulators from _advantage_pump into step scalars. + """Reduce per-step accumulators from _advantage_stage into step scalars. Args: - rewards: One tensor per advantage_pump call; each row a sample reward. + 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. diff --git a/tests/unit/single_controller/test_rollout_pump.py b/tests/unit/single_controller/test_rollout_pump.py index 008ddc45eb..372072403d 100644 --- a/tests/unit/single_controller/test_rollout_pump.py +++ b/tests/unit/single_controller/test_rollout_pump.py @@ -257,14 +257,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 @@ -281,8 +281,8 @@ def test_rollout_pump_writes_expected_tq_data( max_weight_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, @@ -298,7 +298,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, @@ -322,8 +322,8 @@ def test_rollout_pump_writes_expected_tq_data( 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, ) @@ -366,7 +366,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_train_pump.py b/tests/unit/single_controller/test_train_pump.py index 5c50c16871..150fd9fde7 100644 --- a/tests/unit/single_controller/test_train_pump.py +++ b/tests/unit/single_controller/test_train_pump.py @@ -295,8 +295,8 @@ def test_train_pump_drives_mcore_training_step( trainer_handle=trainer, weight_synchronizer=weight_sync, loss_fn=SimpleLossFn(), - advantage_estimator=adv_est, rollout_manager=rollout_manager, + advantage_estimator=adv_est, dataloader=None, tq_buffer=tq_buffer, ) From 52e60e8d0ce30108091442553b28f50ee65e33bb Mon Sep 17 00:00:00 2001 From: ruit Date: Wed, 22 Jul 2026 07:13:56 -0700 Subject: [PATCH 14/17] fix(sc): restore advantage metrics and coverage Signed-off-by: ruit --- nemo_rl/algorithms/single_controller.py | 5 + .../test_sc_utils_helpers.py | 186 ++++++++++++++++++ .../unit/single_controller/test_train_pump.py | 52 ++++- 3 files changed, 241 insertions(+), 2 deletions(-) create mode 100644 tests/unit/single_controller/test_sc_utils_helpers.py diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index c10f44e7f8..2047e70c9f 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -827,6 +827,11 @@ 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", 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_train_pump.py b/tests/unit/single_controller/test_train_pump.py index 150fd9fde7..bc2b64cf23 100644 --- a/tests/unit/single_controller/test_train_pump.py +++ b/tests/unit/single_controller/test_train_pump.py @@ -17,7 +17,9 @@ from __future__ import annotations import gc +import math from types import SimpleNamespace +from typing import Any import pytest import ray @@ -161,7 +163,7 @@ def compute_advantage( repeated_batch: dict, **kwargs, ) -> torch.Tensor: - return rewards.detach().clone() + return rewards.unsqueeze(-1).expand(mask.shape).detach().clone() @ray.remote(num_cpus=0) # pragma: no cover @@ -179,6 +181,41 @@ def get(self) -> list[tuple[str, dict]]: return list(self._entries) +class _RecordingLogger: + """Forward metrics logged inside SC to a Ray actor visible to the test.""" + + def __init__(self, log_handle: Any) -> None: + self._log = log_handle + + def log_metrics( + self, + metrics: dict[str, Any], + step: int, + prefix: str | None = "", + ) -> None: + ray.get( + self._log.record.remote( + "metrics", + { + "metrics": dict(metrics), + "step": int(step), + "prefix": prefix, + }, + ) + ) + + +@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: """Records the version SC syncs at end-of-step, via a Ray actor log.""" @@ -288,7 +325,8 @@ def test_train_pump_drives_mcore_training_step( }, ) - ctrl = SingleControllerActor.remote( + ctrl = _RecordingSingleControllerActor.remote( + metric_log_handle=log, cfg=cfg, dp_client_handle=dp_client, gen_handle=None, @@ -312,5 +350,15 @@ def test_train_pump_drives_mcore_training_step( 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() From 00ac501222a6eeced0e8d9759892119128f14fc0 Mon Sep 17 00:00:00 2001 From: Terry Kong Date: Wed, 22 Jul 2026 23:38:20 -0700 Subject: [PATCH 15/17] refactor(sc): interface-driven staleness samplers (design demo for #3220) (#3316) Signed-off-by: Terry Kong --- .../async_utils/staleness_sampler.py | 487 ++++++++++++---- nemo_rl/algorithms/single_controller.py | 173 ++---- .../single_controller/test_rollout_pump.py | 39 +- .../test_sampler_interface.py | 304 ++++++++++ .../test_staleness_sampler.py | 526 ------------------ .../unit/single_controller/test_train_pump.py | 5 +- 6 files changed, 773 insertions(+), 761 deletions(-) create mode 100644 tests/unit/single_controller/test_sampler_interface.py delete mode 100644 tests/unit/single_controller/test_staleness_sampler.py diff --git a/nemo_rl/algorithms/async_utils/staleness_sampler.py b/nemo_rl/algorithms/async_utils/staleness_sampler.py index 86251dbfba..c8d699343b 100644 --- a/nemo_rl/algorithms/async_utils/staleness_sampler.py +++ b/nemo_rl/algorithms/async_utils/staleness_sampler.py @@ -12,71 +12,162 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Prompt-group selection over a TQReplayBuffer.""" +"""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 -class StalenessSampler: - """Pick complete prompt groups from a TQReplayBuffer. - Args: - buffer: Shared TQReplayBuffer holding the candidate slots. - max_staleness_versions: Max weight-version gap a sample may have from the trainer. - sample_freshest_first: Prefer smallest lag when picking from the in-window set. - strict_weight_fifo: Consume in strict weight_version FIFO — take only from the oldest in-window weight_version and wait for its batch to fill. - force_in_order: Match each slot's target_step against current_train_weight, ignoring the window; mirrors legacy async_grpo target_weight semantics. +@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. """ - def __init__( + 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, - buffer: TQReplayBuffer, - max_staleness_versions: int, - sample_freshest_first: bool = False, - strict_weight_fifo: bool = False, - force_in_order: bool = False, - ) -> None: - if max_staleness_versions < 0: - raise ValueError( - f"max_staleness_versions must be non-negative, got " - f"{max_staleness_versions}" - ) - if strict_weight_fifo and sample_freshest_first: - raise ValueError( - "strict_weight_fifo and sample_freshest_first are mutually exclusive" - ) + *, + 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 - self.max_staleness_versions = max_staleness_versions - self.sample_freshest_first = sample_freshest_first - self.strict_weight_fifo = strict_weight_fifo - self.force_in_order = force_in_order + # 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[KVBatchMeta | None, int]: - """Concat up to max_prompt_groups eligible groups and drop them from the buffer. + ) -> tuple[Optional[KVBatchMeta], int]: + ... - Eligibility = ready and weight in - [current_train_weight - max_staleness_versions, current_train_weight]. - DataPlane rows survive the local drop; caller clears them at step boundary. - - Args: - current_train_weight: Current trainer weight version. - min_prompt_groups: Minimum groups required; returns (None, 0) below this. - max_prompt_groups: Cap on groups returned when the threshold is met. - Selection is greedy without waiting: it returns all currently - eligible groups up to this cap (never deliberately fewer, and it - does not wait for the cap to fill). + async def evict(self, *, current_train_weight: int) -> int: + """Default: drop *ready* groups below the weight window. - Returns: - meta: Concatenated KVBatchMeta, or None if not enough groups. - num_groups: Number of prompt groups in meta; 0 when meta is None. + 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: @@ -85,42 +176,77 @@ async def select( f"min_prompt_groups ({min_prompt_groups})" ) - if self.force_in_order: - # target_step exact match; staleness window ignored. - valid_idxs = [ - i - for i, target in enumerate(self._buffer.target_step_list) - if target == current_train_weight and self._buffer.ready_list[i] - ] - else: - min_valid_version = max( - 0, current_train_weight - self.max_staleness_versions - ) - if self.strict_weight_fifo: - 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] - ] - else: - 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] - ] + 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: ( @@ -128,36 +254,205 @@ async def select( i, ) ) + return await self._finalize_selection( + valid_idxs, min_prompt_groups, max_prompt_groups + ) - 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) +class _GatedSampler(BaseSampler): + """Base for policies that admit exactly one dispatch batch per trainer step. - return ( - selected_metas[0].concat(*selected_metas[1:]), # type: ignore - len(selected_idxs), - ) + The gate bounds how far generation may run ahead of the trainer + (``gate_window`` versions of lookahead). + """ - async def evict(self, *, current_train_weight: int) -> int: - """Drop groups whose weight falls below the staleness window. + 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 - Future entries (weight > current_train_weight) are left alone. + def _eviction_window(self) -> int: + return self._gate_window - Args: - current_train_weight: Current trainer weight version; groups with - weight < current_train_weight - max_staleness_versions are dropped. + 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) - Returns: - Number of group entries removed from the buffer. - """ + 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) - stale_idxs = [ + 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 < min_valid_version + 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 " + f"{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 2047e70c9f..5246a7ec11 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -50,7 +50,11 @@ import torch from pydantic import BaseModel, Field -from nemo_rl.algorithms.async_utils.staleness_sampler import StalenessSampler +from nemo_rl.algorithms.async_utils.staleness_sampler import ( + InOrderSamplerConfig, + SamplerConfig, + create_sampler, +) from nemo_rl.algorithms.grpo import GRPOLoggerConfig from nemo_rl.algorithms.single_controller_utils.utils import ( aggregate_step_metrics, @@ -100,29 +104,26 @@ 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 @@ -164,57 +165,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 ``sampler.evict`` 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: @@ -266,27 +216,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 @@ -323,17 +256,28 @@ def __init__( f"({cfg.group_size}) so a mini-batch contains " f"whole prompt groups" ) - # 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( - self._buffer, - max_staleness_versions=cfg.max_weight_staleness_versions, - strict_weight_fifo=not cfg.over_sampling, - force_in_order=cfg.force_in_order, + # 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"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." + ) # ── asyncio state ────────────────────────────────────────────────── # Gate: cleared during _sync_weights, set when generation may proceed @@ -343,9 +287,9 @@ def __init__( # 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. @@ -370,8 +314,7 @@ def __init__( self._current_epoch: int = 0 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}", @@ -437,9 +380,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) @@ -451,9 +395,6 @@ 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 print("rollout_pump: starting", flush=True) async def _dispatch_one_prompt( @@ -497,17 +438,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 diff --git a/tests/unit/single_controller/test_rollout_pump.py b/tests/unit/single_controller/test_rollout_pump.py index 372072403d..9b32628b23 100644 --- a/tests/unit/single_controller/test_rollout_pump.py +++ b/tests/unit/single_controller/test_rollout_pump.py @@ -26,6 +26,12 @@ import torch 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, @@ -53,14 +59,16 @@ @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: @@ -85,13 +93,13 @@ 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"}]]} ) @@ -101,7 +109,6 @@ async def generate_and_push( 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 @@ -142,13 +149,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( { @@ -164,7 +170,6 @@ async def _main() -> None: 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 @@ -221,13 +226,12 @@ 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"}]]}) ] @@ -236,7 +240,6 @@ async def _main() -> None: 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 @@ -277,15 +280,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=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"), 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..386798386e --- /dev/null +++ b/tests/unit/single_controller/test_sampler_interface.py @@ -0,0 +1,304 @@ +# 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_staleness_sampler.py b/tests/unit/single_controller/test_staleness_sampler.py deleted file mode 100644 index 69f580caf7..0000000000 --- a/tests/unit/single_controller/test_staleness_sampler.py +++ /dev/null @@ -1,526 +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. - -"""Unit tests for StalenessSampler (pure filter over TQReplayBuffer state).""" - -from __future__ import annotations - -import asyncio - -import pytest - -from nemo_rl.algorithms.async_utils.staleness_sampler import StalenessSampler -from nemo_rl.data_plane import KVBatchMeta - - -class FakeBuffer: - """Minimal TQReplayBuffer surface used by StalenessSampler tests.""" - - 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, - group_size: int = 1, - ready: bool = True, - end_weight: int | None = None, - target_step: int | None = None, - ) -> KVBatchMeta: - sample_ids = [f"{group_id}_g{i}" for i in range(group_size)] - meta = KVBatchMeta( - partition_id=self._partition_id, - task_name=None, - sample_ids=sample_ids, - tags=[{"weight_version": weight, "group_id": group_id}] * group_size, - ) - self.meta_list.append(meta if ready else None) - self.start_weight_list.append(weight) - self.end_weight_list.append(weight if end_weight is None else end_weight) - self.target_step_list.append(target_step) - self.ready_list.append(ready) - return meta - - 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 TestStalenessSamplerSelect: - def test_select_returns_none_when_insufficient(self): - buf = FakeBuffer() - buf.add("g0", weight=5) - sampler = StalenessSampler(buf, max_staleness_versions=2) - - result = _run( - sampler.select( - current_train_weight=5, min_prompt_groups=2, max_prompt_groups=2 - ) - ) - assert result == (None, 0) - - def test_select_filters_by_staleness_window(self): - buf = FakeBuffer() - # Weights 3, 4, 5, 2, 6 against trainer=5, max_staleness=2: - # lags = 2, 1, 0, 3 (stale), -1 (future) - for i, w in enumerate([3, 4, 5, 2, 6]): - buf.add(f"g{i}", weight=w) - sampler = StalenessSampler( - buf, max_staleness_versions=2, sample_freshest_first=True - ) - - selected, num_groups = _run( - sampler.select( - current_train_weight=5, min_prompt_groups=2, max_prompt_groups=2 - ) - ) - - assert selected is not None - # Freshest first → g2 (lag 0), g1 (lag 1) - assert selected.sample_ids == ["g2_g0", "g1_g0"] - assert num_groups == 2 - - def test_select_fifo_orders_by_insertion(self): - buf = FakeBuffer() - for w in [3, 4, 5]: - buf.add(f"v{w}", weight=w) - sampler = StalenessSampler( - buf, max_staleness_versions=5, sample_freshest_first=False - ) - - selected, num_groups = _run( - sampler.select( - current_train_weight=6, min_prompt_groups=2, max_prompt_groups=2 - ) - ) - assert selected is not None - assert selected.sample_ids == ["v3_g0", "v4_g0"] - assert num_groups == 2 - - def test_select_skips_future_weight(self): - buf = FakeBuffer() - buf.add("now", weight=5) - buf.add("future", weight=7) - sampler = StalenessSampler(buf, max_staleness_versions=10) - - selected, num_groups = _run( - sampler.select( - current_train_weight=5, min_prompt_groups=1, max_prompt_groups=1 - ) - ) - - assert selected is not None - assert selected.sample_ids == ["now_g0"] - assert num_groups == 1 - - def test_select_concats_groups(self): - buf = FakeBuffer() - buf.add("g0", weight=5, group_size=2) - buf.add("g1", weight=5, group_size=2) - sampler = StalenessSampler(buf, max_staleness_versions=0) - - selected, num_groups = _run( - sampler.select( - current_train_weight=5, min_prompt_groups=2, max_prompt_groups=2 - ) - ) - - assert selected is not None - assert selected.sample_ids == [ - "g0_g0", - "g0_g1", - "g1_g0", - "g1_g1", - ] - # Two groups concatenated, each of size 2 → 4 sample_ids total. - assert num_groups == 2 - - def test_select_strict_on_policy_requires_exact_version(self): - buf = FakeBuffer() - for i, w in enumerate([4, 5, 5, 6]): - buf.add(f"g{i}", weight=w) - sampler = StalenessSampler(buf, max_staleness_versions=0) - - # 3 eligible (need weight=5), only have 2 - result = _run( - sampler.select( - current_train_weight=5, min_prompt_groups=3, max_prompt_groups=3 - ) - ) - assert result == (None, 0) - - # Buffer still intact: select with min=3 returned None without dropping anything. - selected, num_groups = _run( - sampler.select( - current_train_weight=5, min_prompt_groups=2, max_prompt_groups=2 - ) - ) - assert selected is not None - assert selected.sample_ids == ["g1_g0", "g2_g0"] - assert num_groups == 2 - - def test_select_drops_returned_entries_from_buffer(self): - buf = FakeBuffer() - for i, w in enumerate([5, 5, 5]): - buf.add(f"g{i}", weight=w) - sampler = StalenessSampler(buf, max_staleness_versions=0) - - first_meta, first_num_groups = _run( - sampler.select( - current_train_weight=5, min_prompt_groups=1, max_prompt_groups=1 - ) - ) - assert first_meta is not None - assert first_meta.sample_ids == ["g0_g0"] - assert first_num_groups == 1 - assert buf.start_weight_list == [5, 5] - # remove_in_dp=False; DP rows kept for trainer. - assert buf.remove_calls[-1][1] is False - - second_meta, second_num_groups = _run( - sampler.select( - current_train_weight=5, min_prompt_groups=1, max_prompt_groups=1 - ) - ) - assert second_meta is not None - assert second_meta.sample_ids == ["g1_g0"] - assert second_num_groups == 1 - - def test_select_rejects_zero_min_prompt_groups(self): - buf = FakeBuffer() - sampler = StalenessSampler(buf, max_staleness_versions=0) - with pytest.raises(ValueError): - _run( - sampler.select( - current_train_weight=0, min_prompt_groups=0, max_prompt_groups=0 - ) - ) - - def test_select_rejects_max_less_than_min(self): - buf = FakeBuffer() - for i in range(3): - buf.add(f"g{i}", weight=5) - sampler = StalenessSampler(buf, max_staleness_versions=0) - - with pytest.raises(ValueError): - _run( - sampler.select( - current_train_weight=5, min_prompt_groups=2, max_prompt_groups=1 - ) - ) - - def test_select_caps_at_max_prompt_groups(self): - buf = FakeBuffer() - for i in range(5): - buf.add(f"g{i}", weight=5) - sampler = StalenessSampler(buf, max_staleness_versions=0) - - selected, num_groups = _run( - sampler.select( - current_train_weight=5, min_prompt_groups=2, max_prompt_groups=3 - ) - ) - - assert selected is not None - # FIFO order; capped at max=3 even though 5 are eligible. - assert selected.sample_ids == ["g0_g0", "g1_g0", "g2_g0"] - assert num_groups == 3 - # The remaining two stay in the buffer. - assert buf.start_weight_list == [5, 5] - - def test_select_takes_all_available_when_between_min_and_max(self): - buf = FakeBuffer() - for i in range(3): - buf.add(f"g{i}", weight=5) - sampler = StalenessSampler(buf, max_staleness_versions=0) - - selected, num_groups = _run( - sampler.select( - current_train_weight=5, min_prompt_groups=2, max_prompt_groups=8 - ) - ) - - assert selected is not None - assert selected.sample_ids == ["g0_g0", "g1_g0", "g2_g0"] - assert num_groups == 3 - assert buf.start_weight_list == [] - - -class TestStalenessSamplerEvict: - def test_evict_removes_stale_groups(self): - buf = FakeBuffer() - # trainer=5, max_staleness=1 → lag >1 means stale (weights 0, 1, 2 stale; 4, 5 fresh) - for i, w in enumerate([0, 1, 4, 5, 2]): - buf.add(f"g{i}", weight=w) - sampler = StalenessSampler(buf, max_staleness_versions=1) - - dropped = _run(sampler.evict(current_train_weight=5)) - - assert dropped == 3 - assert buf.start_weight_list == [4, 5] - # Survivors' sample_ids - assert [m.sample_ids[0] for m in buf.meta_list] == ["g2_g0", "g3_g0"] - - def test_evict_returns_zero_when_nothing_stale(self): - buf = FakeBuffer() - for w in [4, 5]: - buf.add(f"v{w}", weight=w) - sampler = StalenessSampler(buf, max_staleness_versions=1) - - assert _run(sampler.evict(current_train_weight=5)) == 0 - assert buf.remove_calls == [] - - def test_evict_keeps_future_groups(self): - buf = FakeBuffer() - buf.add("future", weight=7) - sampler = StalenessSampler(buf, max_staleness_versions=0) - - assert _run(sampler.evict(current_train_weight=5)) == 0 - assert buf.start_weight_list == [7] - - def test_evict_drops_whole_group(self): - buf = FakeBuffer() - buf.add("stale", weight=1, group_size=4) - buf.add("fresh", weight=5, group_size=4) - sampler = StalenessSampler(buf, max_staleness_versions=1) - - dropped = _run(sampler.evict(current_train_weight=5)) - - assert dropped == 1 - assert buf.remove_calls == [([0], True)] - assert buf.start_weight_list == [5] - assert [m.sample_ids[0] for m in buf.meta_list] == ["fresh_g0"] - - -class TestStalenessSamplerInit: - def test_rejects_negative_max_staleness(self): - buf = FakeBuffer() - with pytest.raises(ValueError): - StalenessSampler(buf, max_staleness_versions=-1) - - def test_rejects_strict_weight_fifo_with_freshest_first(self): - buf = FakeBuffer() - with pytest.raises(ValueError): - StalenessSampler( - buf, - max_staleness_versions=0, - sample_freshest_first=True, - strict_weight_fifo=True, - ) - - -class TestStalenessSamplerReady: - def test_default_mode_skips_unready_slots(self): - buf = FakeBuffer() - buf.add("g0", weight=5, ready=False) - buf.add("g1", weight=5, ready=True) - sampler = StalenessSampler(buf, max_staleness_versions=0) - - selected, num_groups = _run( - sampler.select( - current_train_weight=5, min_prompt_groups=1, max_prompt_groups=1 - ) - ) - - assert selected is not None - assert selected.sample_ids == ["g1_g0"] - assert num_groups == 1 - - def test_default_mode_waits_when_too_few_ready(self): - buf = FakeBuffer() - buf.add("g0", weight=5, ready=False) - buf.add("g1", weight=5, ready=True) - sampler = StalenessSampler(buf, max_staleness_versions=0) - - result = _run( - sampler.select( - current_train_weight=5, min_prompt_groups=2, max_prompt_groups=2 - ) - ) - assert result == (None, 0) - - -class TestStalenessSamplerStrictWeightFifo: - def test_consumes_oldest_batch_first(self): - buf = FakeBuffer() - # Two complete batches: v=4 then v=5; strict_weight_fifo must take v=4 first. - for i, w in enumerate((4, 4, 5, 5)): - buf.add(f"v{w}_{i}", weight=w) - sampler = StalenessSampler( - buf, max_staleness_versions=1, strict_weight_fifo=True - ) - - selected, num_groups = _run( - sampler.select( - current_train_weight=5, min_prompt_groups=2, max_prompt_groups=2 - ) - ) - - assert selected is not None - # Insertion-order FIFO inside the oldest batch. - assert selected.sample_ids == ["v4_0_g0", "v4_1_g0"] - assert num_groups == 2 - assert buf.start_weight_list == [5, 5] - - def test_waits_when_oldest_batch_partially_ready(self): - buf = FakeBuffer() - # Oldest batch v=4 has 1 ready + 1 unready; v=5 batch is fully ready. - # strict_weight_fifo must NOT skip ahead to v=5. - buf.add("v4_a", weight=4, ready=True) - buf.add("v4_b", weight=4, ready=False) - buf.add("v5_a", weight=5, ready=True) - buf.add("v5_b", weight=5, ready=True) - sampler = StalenessSampler( - buf, max_staleness_versions=1, strict_weight_fifo=True - ) - - result = _run( - sampler.select( - current_train_weight=5, min_prompt_groups=2, max_prompt_groups=2 - ) - ) - assert result == (None, 0) - # Buffer untouched: nothing removed. - assert buf.start_weight_list == [4, 4, 5, 5] - assert buf.ready_list == [True, False, True, True] - - def test_returns_none_when_oldest_batch_not_filled(self): - buf = FakeBuffer() - buf.add("v4_a", weight=4, ready=True) - # Only 1 ready in oldest batch; need 2. - sampler = StalenessSampler( - buf, max_staleness_versions=1, strict_weight_fifo=True - ) - - result = _run( - sampler.select( - current_train_weight=5, min_prompt_groups=2, max_prompt_groups=2 - ) - ) - assert result == (None, 0) - - def test_ignores_future_versions_when_picking_target(self): - buf = FakeBuffer() - # Trainer at 5, staleness 1: window is [4, 5]; v=7 (future) must not - # become the oldest target. - buf.add("v7", weight=7, ready=True) - buf.add("v5_a", weight=5, ready=True) - buf.add("v5_b", weight=5, ready=True) - sampler = StalenessSampler( - buf, max_staleness_versions=1, strict_weight_fifo=True - ) - - selected, num_groups = _run( - sampler.select( - current_train_weight=5, min_prompt_groups=2, max_prompt_groups=2 - ) - ) - - assert selected is not None - assert selected.sample_ids == ["v5_a_g0", "v5_b_g0"] - assert num_groups == 2 - assert buf.start_weight_list == [7] - - -class TestStalenessSamplerForceInOrder: - def test_selects_only_groups_whose_target_step_matches_current(self): - buf = FakeBuffer() - # Two matching-target groups and one future-target group. - buf.add("g_match_a", weight=5, target_step=5) - buf.add("g_match_b", weight=5, target_step=5) - buf.add("g_future", weight=5, target_step=7) - sampler = StalenessSampler(buf, max_staleness_versions=0, force_in_order=True) - - selected, num_groups = _run( - sampler.select( - current_train_weight=5, min_prompt_groups=1, max_prompt_groups=4 - ) - ) - - assert selected is not None - assert selected.sample_ids == ["g_match_a_g0", "g_match_b_g0"] - assert num_groups == 2 - # The future-target group survives; only the two matches were removed. - assert buf.target_step_list == [7] - assert buf.start_weight_list == [5] - - def test_ignores_stale_target_step(self): - buf = FakeBuffer() - # Trainer at 5; a slot with target_step=4 must NOT satisfy force_in_order. - buf.add("g_stale", weight=4, target_step=4) - sampler = StalenessSampler( - buf, - # Wide staleness window: force_in_order must ignore it entirely. - max_staleness_versions=10, - force_in_order=True, - ) - - result = _run( - sampler.select( - current_train_weight=5, min_prompt_groups=1, max_prompt_groups=4 - ) - ) - - assert result == (None, 0) - # Buffer untouched. - assert buf.target_step_list == [4] - - def test_selects_when_target_matches_even_if_weight_outside_window(self): - buf = FakeBuffer() - # start_weight=100 is well outside any realistic staleness window, - # but target_step matches current — force_in_order must ignore the window. - buf.add("g_match_stale_weight", weight=100, target_step=5) - sampler = StalenessSampler(buf, max_staleness_versions=0, force_in_order=True) - - selected, num_groups = _run( - sampler.select( - current_train_weight=5, min_prompt_groups=1, max_prompt_groups=4 - ) - ) - - assert selected is not None - assert selected.sample_ids == ["g_match_stale_weight_g0"] - assert num_groups == 1 - - def test_ignores_unready_group_with_matching_target(self): - buf = FakeBuffer() - buf.add("g_unready", weight=5, target_step=5, ready=False) - sampler = StalenessSampler(buf, max_staleness_versions=0, force_in_order=True) - - result = _run( - sampler.select( - current_train_weight=5, min_prompt_groups=1, max_prompt_groups=4 - ) - ) - - assert result == (None, 0) - # Buffer untouched. - assert buf.ready_list == [False] diff --git a/tests/unit/single_controller/test_train_pump.py b/tests/unit/single_controller/test_train_pump.py index bc2b64cf23..b6134df948 100644 --- a/tests/unit/single_controller/test_train_pump.py +++ b/tests/unit/single_controller/test_train_pump.py @@ -27,6 +27,7 @@ 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, @@ -298,16 +299,14 @@ def test_train_pump_drives_mcore_training_step( rollout_manager = SimpleNamespace(_tq_buffer=None) cfg = SingleControllerConfig.model_construct( - max_weight_staleness_versions=1, + sampler=WindowedSamplerConfig(max_staleness_versions=1), min_groups_per_batch=num_prompts, target_groups_per_step=num_prompts, group_size=num_generations, - batch_selection_strategy="staleness_window", max_inflight_prompts=num_prompts, max_buffered_rollouts=num_prompts, max_train_steps=train_steps, max_num_epochs=None, - over_sampling=True, train_global_batch_size=train_gbs, partition_id=_PARTITION_ID, advantage_enabled=True, From 8257e1dab17ff38a694ab510c1a9c76ee3a10b78 Mon Sep 17 00:00:00 2001 From: ruit Date: Thu, 23 Jul 2026 01:40:56 -0700 Subject: [PATCH 16/17] fix(sc): enforce complete training steps Signed-off-by: ruit --- nemo_rl/algorithms/single_controller.py | 66 +++---- .../single_controller/test_rollout_pump.py | 6 + .../test_single_controller.py | 166 ++++++++++++++++++ 3 files changed, 207 insertions(+), 31 deletions(-) create mode 100644 tests/unit/single_controller/test_single_controller.py diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 5246a7ec11..da63413e3d 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -131,13 +131,10 @@ class SingleControllerConfig(BaseModel, extra="allow"): # 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 @@ -165,7 +162,6 @@ class SingleControllerConfig(BaseModel, extra="allow"): logger: GRPOLoggerConfig - @ray.remote(num_cpus=1, num_gpus=0) # pragma: no cover class SingleControllerActor: """CPU-only Ray actor that orchestrates the RL training loop. @@ -229,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 @@ -242,19 +237,12 @@ 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: - 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})" - ) - if cfg.train_global_batch_size % cfg.group_size != 0: + if cfg.train_global_batch_size != samples_per_step: 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"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" ) # Staleness policy owns admit (rollout side) + select/evict (train side) # + is_on_policy / required_buffer_capacity (derived facts). ``name`` in @@ -284,6 +272,11 @@ def __init__( 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 @@ -395,6 +388,7 @@ async def _rollout_pump(self) -> None: 5. Decrement _inflight_rollouts """ sem = asyncio.Semaphore(self._cfg.max_inflight_prompts) + self._rollout_exhausted.clear() print("rollout_pump: starting", flush=True) async def _dispatch_one_prompt( @@ -478,6 +472,7 @@ 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: @@ -541,6 +536,22 @@ async def _train_pump(self) -> None: # 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 @@ -607,13 +618,6 @@ async def _train_pump(self) -> None: groups_dispatched += num_groups - if not step_open: - print( - "train_pump: rollout exhausted before any group ready", - flush=True, - ) - break - with self._timer.time("policy_training"): result = await asyncio.to_thread(self._trainer.finish_train_step) diff --git a/tests/unit/single_controller/test_rollout_pump.py b/tests/unit/single_controller/test_rollout_pump.py index 9b32628b23..4a2ba61cce 100644 --- a/tests/unit/single_controller/test_rollout_pump.py +++ b/tests/unit/single_controller/test_rollout_pump.py @@ -106,6 +106,7 @@ async def generate_and_push( 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() @@ -115,6 +116,7 @@ async def generate_and_push( 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: @@ -167,6 +169,7 @@ 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() @@ -181,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()) @@ -237,6 +241,7 @@ async def _main() -> None: ] 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() @@ -250,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()) 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)) From e6f3d762706f2beaec018f95f267ac4d3f2e4afd Mon Sep 17 00:00:00 2001 From: ruit Date: Thu, 23 Jul 2026 02:13:18 -0700 Subject: [PATCH 17/17] style(sc): format sampler files Signed-off-by: ruit --- .../algorithms/async_utils/staleness_sampler.py | 15 ++++++--------- .../single_controller/test_sampler_interface.py | 8 ++++++-- 2 files changed, 12 insertions(+), 11 deletions(-) diff --git a/nemo_rl/algorithms/async_utils/staleness_sampler.py b/nemo_rl/algorithms/async_utils/staleness_sampler.py index c8d699343b..c729d2f4c9 100644 --- a/nemo_rl/algorithms/async_utils/staleness_sampler.py +++ b/nemo_rl/algorithms/async_utils/staleness_sampler.py @@ -121,8 +121,9 @@ def __init__(self, buffer: TQReplayBuffer) -> None: # ── rollout-pump side ──────────────────────────────────────────────── @abc.abstractmethod - async def admit(self, *, trainer_version_fn: Callable[[], int]) -> Optional[int]: - ... + async def admit( + self, *, trainer_version_fn: Callable[[], int] + ) -> Optional[int]: ... # ── train-pump side ────────────────────────────────────────────────── @abc.abstractmethod @@ -132,8 +133,7 @@ async def select( current_train_weight: int, min_prompt_groups: int, max_prompt_groups: int, - ) -> tuple[Optional[KVBatchMeta], int]: - ... + ) -> tuple[Optional[KVBatchMeta], int]: ... async def evict(self, *, current_train_weight: int) -> int: """Default: drop *ready* groups below the weight window. @@ -437,15 +437,12 @@ def create_sampler( buffer, max_staleness_versions=cfg.max_staleness_versions ) if isinstance(cfg, InOrderSamplerConfig): - return InOrderSampler( - buffer, max_lookahead_versions=cfg.max_lookahead_versions - ) + 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 " - f"{cfg.target!r}" + 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 {})) diff --git a/tests/unit/single_controller/test_sampler_interface.py b/tests/unit/single_controller/test_sampler_interface.py index 386798386e..8e0ef3e20c 100644 --- a/tests/unit/single_controller/test_sampler_interface.py +++ b/tests/unit/single_controller/test_sampler_interface.py @@ -152,7 +152,9 @@ def test_unready_slot_is_never_evicted(self): class TestFactory: def test_windowed_config_builds_windowed(self): - s = create_sampler(FakeBuffer(), WindowedSamplerConfig(max_staleness_versions=3)) + s = create_sampler( + FakeBuffer(), WindowedSamplerConfig(max_staleness_versions=3) + ) assert isinstance(s, WindowedSampler) assert s.max_staleness_versions == 3 @@ -166,7 +168,9 @@ def test_weight_fifo_config_builds_weight_fifo(self): WeightFifoSamplerConfig, ) - s = create_sampler(FakeBuffer(), WeightFifoSamplerConfig(max_staleness_versions=4)) + s = create_sampler( + FakeBuffer(), WeightFifoSamplerConfig(max_staleness_versions=4) + ) assert isinstance(s, WeightFifoSampler) assert s.max_staleness_versions == 4