From 5b7f9dd52c809ed316139e1d8516193f0bc2d64a Mon Sep 17 00:00:00 2001 From: Yuki Huang Date: Sat, 22 Aug 2026 21:40:50 -0700 Subject: [PATCH 01/31] feat(sc): add the driver-side TQValue for the PPO critic Signed-off-by: Yuki Huang --- nemo_rl/data_plane/driver_mixin.py | 97 ++++++++++++++ nemo_rl/data_plane/schema.py | 17 +++ nemo_rl/models/policy/tq_policy.py | 72 +---------- nemo_rl/models/value/__init__.py | 9 +- nemo_rl/models/value/tq_value.py | 201 +++++++++++++++++++++++++++++ pyrefly.toml | 1 + 6 files changed, 326 insertions(+), 71 deletions(-) create mode 100644 nemo_rl/data_plane/driver_mixin.py create mode 100644 nemo_rl/models/value/tq_value.py diff --git a/nemo_rl/data_plane/driver_mixin.py b/nemo_rl/data_plane/driver_mixin.py new file mode 100644 index 00000000000..ce6c2b2d2c1 --- /dev/null +++ b/nemo_rl/data_plane/driver_mixin.py @@ -0,0 +1,97 @@ +# 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. +"""Driver-side TransferQueue helpers shared by TQPolicy and TQValue.""" + +from __future__ import annotations + +from typing import Any, Optional + +from nemo_rl.data_plane.column_io import read_columns, round_up, write_columns +from nemo_rl.data_plane.interfaces import KVBatchMeta +from nemo_rl.data_plane.schema import GLOBAL_FORWARD_PAD_SEQLEN +from nemo_rl.distributed.batched_data_dict import BatchedDataDict + + +class TQDriverMixin: + """Pad-target minting and column read/write against the data plane. + + Hosts must provide cfg, dp_client, and the use_dynamic_batches / + use_sequence_packing attribute pairs that Policy and Value both set. + """ + + def _packing_args( + self, + mb_tokens_key: str, + ) -> tuple[Optional[dict[str, Any]], Optional[dict[str, Any]]]: + """Resolve (sequence_packing_args, dynamic_batching_args) for a given stage. + + The stage is identified by ``mb_tokens_key`` (``"logprob_mb_tokens"`` or + ``"train_mb_tokens"``). + """ + if getattr(self, "use_dynamic_batches", False): + args = dict(self.dynamic_batching_args) + args["max_tokens_per_microbatch"] = self.cfg["dynamic_batching"][ + mb_tokens_key + ] + return None, args + if getattr(self, "use_sequence_packing", False): + args = dict(self.sequence_packing_args) + args["max_tokens_per_microbatch"] = self.cfg["sequence_packing"][ + mb_tokens_key + ] + return args, None + return None, None + + def _stamp_pad_seqlen(self, meta: KVBatchMeta) -> None: + """Mint ``GLOBAL_FORWARD_PAD_SEQLEN`` onto ``meta.extra_info`` (idempotent). + + Cross-DP forward pad target. Preshard shards inherit it via + ``dict(meta.extra_info)`` propagation. + """ + if not meta.sequence_lengths: + return + if GLOBAL_FORWARD_PAD_SEQLEN in meta.extra_info: + return + _, dba = self._packing_args("train_mb_tokens") + seq_round = int(dba["sequence_length_round"]) if dba is not None else 1 + pad_mult = int(meta.extra_info.get("pad_to_multiple", 1)) + meta.extra_info[GLOBAL_FORWARD_PAD_SEQLEN] = round_up( + max(meta.sequence_lengths), max(pad_mult, seq_round) + ) + + def read_from_dataplane( + self, + meta: KVBatchMeta, + *, + select_fields: list[str], + pad_value_dict: Optional[dict[str, Any]] = None, + ) -> BatchedDataDict[Any]: + """Fetch + materialize columns from the data plane (TQ). + + ``read_columns`` pads to ``meta.extra_info[GLOBAL_FORWARD_PAD_SEQLEN]`` + — the same value workers pad to in their forward pass. Driver + and workers thus return columns at one identical seq dim, with + no driver-side knowledge of ``sequence_length_round``. + """ + self._stamp_pad_seqlen(meta) + return read_columns( + self.dp_client, + meta, + select_fields=select_fields, + pad_value_dict=pad_value_dict, + ) + + def write_to_dataplane(self, meta: KVBatchMeta, fields: dict[str, Any]) -> None: + """Write driver-computed columns to the data plane (TQ).""" + write_columns(self.dp_client, meta, fields=fields) diff --git a/nemo_rl/data_plane/schema.py b/nemo_rl/data_plane/schema.py index 49cf79422e7..56464e4c4d1 100644 --- a/nemo_rl/data_plane/schema.py +++ b/nemo_rl/data_plane/schema.py @@ -52,6 +52,23 @@ "sample_mask", ) +# Kept out of DP_TRAIN_FIELDS: a GRPO run writes neither, and a worker fetching +# a column nobody wrote errors out rather than reading zeros. +PPO_VALUE_FIELDS = ( + "values", + "returns", +) + +DP_VALUE_TRAIN_FIELDS = ( + "input_ids", + "input_lengths", + "token_mask", + "sample_mask", + *PPO_VALUE_FIELDS, +) + +VALUE_SEED_FIELDS = LP_SEED_FIELDS + # Fields requested for KV-scale calibration. Positive include-list: # calibration only handles seq-dim tensor inputs, so we name them # explicitly. Train-side deltas (logprobs/advantages/masks) and diff --git a/nemo_rl/models/policy/tq_policy.py b/nemo_rl/models/policy/tq_policy.py index 13f04d7d77c..6fb4becde70 100644 --- a/nemo_rl/models/policy/tq_policy.py +++ b/nemo_rl/models/policy/tq_policy.py @@ -39,15 +39,13 @@ from nemo_rl.algorithms.loss.interfaces import LossFunction from nemo_rl.data_plane import DataPlaneConfig, KVBatchMeta, build_data_plane_client -from nemo_rl.data_plane.column_io import read_columns, round_up, write_columns +from nemo_rl.data_plane.driver_mixin import TQDriverMixin from nemo_rl.data_plane.preshard import shard_meta_for_dp from nemo_rl.data_plane.schema import ( DP_TRAIN_FIELDS, - GLOBAL_FORWARD_PAD_SEQLEN, LP_SEED_FIELDS, fields_with_optional_routed_experts, ) -from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.models.policy.lm_policy import Policy from nemo_rl.utils.flops_tracker import get_theoretical_tflops from nemo_rl.utils.timer import Timer @@ -81,7 +79,7 @@ def _aggregate_train_results(results: list[dict[str, Any]]) -> dict[str, Any]: # dispatcher only waits for completion — no aggregation needed. -class TQPolicy(Policy): +class TQPolicy(TQDriverMixin, Policy): """TQ-mediated counterpart to :class:`Policy`. Constructor accepts an additional ``dp_cfg`` (the @@ -199,74 +197,8 @@ def finish_step(self, meta: KVBatchMeta) -> None: """Drop this step's bulk from TQ. Mirror of :meth:`prepare_step`.""" self.discard_samples(meta.sample_ids, meta.partition_id) - def _stamp_pad_seqlen(self, meta: KVBatchMeta) -> None: - """Mint ``GLOBAL_FORWARD_PAD_SEQLEN`` onto ``meta.extra_info`` (idempotent). - - Cross-DP forward pad target. Preshard shards inherit it via - ``dict(meta.extra_info)`` propagation. - """ - if not meta.sequence_lengths: - return - if GLOBAL_FORWARD_PAD_SEQLEN in meta.extra_info: - return - _, dba = self._packing_args("train_mb_tokens") - seq_round = int(dba["sequence_length_round"]) if dba is not None else 1 - pad_mult = int(meta.extra_info.get("pad_to_multiple", 1)) - meta.extra_info[GLOBAL_FORWARD_PAD_SEQLEN] = round_up( - max(meta.sequence_lengths), max(pad_mult, seq_round) - ) - - def read_from_dataplane( - self, - meta: KVBatchMeta, - *, - select_fields: list[str], - pad_value_dict: Optional[dict[str, Any]] = None, - ) -> BatchedDataDict[Any]: - """Fetch + materialize columns from the data plane (TQ). - - ``read_columns`` pads to ``meta.extra_info[GLOBAL_FORWARD_PAD_SEQLEN]`` - — the same value workers pad to in their forward pass. Driver - and workers thus return columns at one identical seq dim, with - no driver-side knowledge of ``sequence_length_round``. - """ - self._stamp_pad_seqlen(meta) - return read_columns( - self.dp_client, - meta, - select_fields=select_fields, - pad_value_dict=pad_value_dict, - ) - - def write_to_dataplane(self, meta: KVBatchMeta, fields: dict[str, Any]) -> None: - """Write driver-computed columns to the data plane (TQ).""" - write_columns(self.dp_client, meta, fields=fields) - # ── 1-hop entrypoints (KVBatchMeta in, no re-fan-out) ────────────────── - def _packing_args( - self, - mb_tokens_key: str, - ) -> tuple[Optional[dict[str, Any]], Optional[dict[str, Any]]]: - """Resolve (sequence_packing_args, dynamic_batching_args) for a given stage. - - The stage is identified by ``mb_tokens_key`` (``"logprob_mb_tokens"`` or - ``"train_mb_tokens"``). - """ - if getattr(self, "use_dynamic_batches", False): - args = dict(self.dynamic_batching_args) - args["max_tokens_per_microbatch"] = self.cfg["dynamic_batching"][ - mb_tokens_key - ] - return None, args - if getattr(self, "use_sequence_packing", False): - args = dict(self.sequence_packing_args) - args["max_tokens_per_microbatch"] = self.cfg["sequence_packing"][ - mb_tokens_key - ] - return args, None - return None, None - def _logprob_dispatch( self, meta: KVBatchMeta, diff --git a/nemo_rl/models/value/__init__.py b/nemo_rl/models/value/__init__.py index b3e412ee1f6..13cc3099ea8 100644 --- a/nemo_rl/models/value/__init__.py +++ b/nemo_rl/models/value/__init__.py @@ -15,5 +15,12 @@ from nemo_rl.models.value.config import ValueConfig from nemo_rl.models.value.interfaces import ValueInterface, ValueOutputSpec from nemo_rl.models.value.lm_value import Value +from nemo_rl.models.value.tq_value import TQValue -__all__ = ["Value", "ValueConfig", "ValueInterface", "ValueOutputSpec"] +__all__ = [ + "TQValue", + "Value", + "ValueConfig", + "ValueInterface", + "ValueOutputSpec", +] diff --git a/nemo_rl/models/value/tq_value.py b/nemo_rl/models/value/tq_value.py new file mode 100644 index 00000000000..317d155373f --- /dev/null +++ b/nemo_rl/models/value/tq_value.py @@ -0,0 +1,201 @@ +# 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. +"""TQ-mediated Value: meta-driven 1-hop counterpart to Value, mirroring TQPolicy.""" + +from __future__ import annotations + +import warnings +from contextlib import nullcontext +from dataclasses import replace +from typing import Any, Optional + +import ray + +from nemo_rl.algorithms.loss.interfaces import LossFunction +from nemo_rl.data_plane import DataPlaneConfig, KVBatchMeta, build_data_plane_client +from nemo_rl.data_plane.driver_mixin import TQDriverMixin +from nemo_rl.data_plane.preshard import shard_meta_for_dp +from nemo_rl.data_plane.schema import DP_VALUE_TRAIN_FIELDS, VALUE_SEED_FIELDS +from nemo_rl.models.value.lm_value import Value +from nemo_rl.utils.timer import Timer + +_REPLICATED_AXES = ["context_parallel", "tensor_parallel", "pipeline_parallel"] + + +def _aggregate_train_results(results: list[dict[str, Any]]) -> dict[str, Any]: + """Assemble per-rank value-train results into Value.train's return shape.""" + out: dict[str, Any] = { + "loss": results[0]["global_loss"], + "grad_norm": results[0]["grad_norm"], + } + all_mb_metrics: dict[str, list[Any]] = {} + for r in results: + for k, v in r["all_mb_metrics"].items(): + all_mb_metrics.setdefault(k, []).extend(v) + out["all_mb_metrics"] = all_mb_metrics + return out + + +class TQValue(TQDriverMixin, Value): + """TQ-mediated counterpart to Value, taking an extra dp_cfg. + + Attaches to the TQ controller rather than bootstrapping it: the TQPolicy + built alongside this critic already did that. Partition lifecycle stays + with the caller. + + TODO(#2625): the value workers have no split begin/microbatch/finish train + API yet, so one train_from_meta call is one optimizer step and the + SingleController requires a PPO step to be a single streaming chunk. + """ + + def __init__( + self, + *args: Any, + dp_cfg: DataPlaneConfig, + **kwargs: Any, + ) -> None: + super().__init__(*args, **kwargs) + dp_world = self.sharding_annotations.get_axis_size("data_parallel") + if dp_world <= 0: + raise ValueError( + f"TQValue requires data_parallel axis size > 0, got {dp_world}. " + f"Check cluster config (gpus_per_node * num_nodes) vs. " + f"TP/PP/CP sizes." + ) + self.dp_cfg = dp_cfg + self.dp_client = build_data_plane_client(dp_cfg, bootstrap=False) + ray.get( + self.worker_group.run_all_workers_single_data( + "setup_data_plane", cfg=dp_cfg + ) + ) + + # ── lifecycle ────────────────────────────────────────────────────── + + def shutdown(self) -> bool: # type: ignore[override] + """Close the TQ client before shutting down the worker group.""" + try: + self.dp_client.close() + except Exception as e: + warnings.warn(f"Error closing data-plane client: {e}") + return super().shutdown() + + # ── 1-hop entrypoints (KVBatchMeta in, no re-fan-out) ────────────────── + + def get_values_from_meta( + self, + meta: KVBatchMeta, + micro_batch_size: Optional[int] = None, + timer: Optional[Timer] = None, + ) -> None: + """1-hop counterpart to get_values. + + Returns nothing: the per-token prediction lands in TQ under values via + the worker-side leader write-back, so the GAE stage reads it from there + rather than through Ray. + + Args: + meta: Full-step batch metadata consumed by all DP ranks. + micro_batch_size: Inference micro batch size; None uses the config default. + timer: Optional timer for nested get_values measurements. + """ + self._stamp_pad_seqlen(meta) + spa, dba = self._packing_args("logprob_mb_tokens") + value_meta = replace( + meta, + fields=list(VALUE_SEED_FIELDS), + task_name="value_fwd", + ) + with timer.time("get_values/shard_meta") if timer else nullcontext(): + metas, _ = shard_meta_for_dp( + value_meta, + dp_world=self.sharding_annotations.get_axis_size("data_parallel"), + batch_size=None, + sequence_packing_args=spa, + dynamic_batching_args=dba, + ) + with timer.time("get_values/submit_value_futures") if timer else nullcontext(): + futures = self.worker_group.run_all_workers_sharded_data( + "get_values_presharded", + meta=metas, + in_sharded_axes=["data_parallel"], + replicate_on_axes=_REPLICATED_AXES, + output_is_replicated=_REPLICATED_AXES, + common_kwargs={"micro_batch_size": micro_batch_size}, + ) + # Wait for completion; per-rank returns are None. + self.worker_group.get_all_worker_results(futures) + + def train_from_meta( + self, + meta: KVBatchMeta, + loss_fn: LossFunction, + eval_mode: bool = False, + gbs: Optional[int] = None, + mbs: Optional[int] = None, + timer: Optional[Timer] = None, + ) -> dict[str, Any]: + """1-hop counterpart to train. One call is one optimizer step. + + Args: + meta: Full-step batch metadata consumed by all DP ranks. + loss_fn: Value loss; MseValueLossFn in the PPO path. + eval_mode: Run forward only, without an optimizer step. + gbs: Global batch size; defaults to the config's train_global_batch_size. + mbs: Micro batch size; defaults to the config's train_micro_batch_size. + timer: Optional timer for nested value_training measurements. + + Returns: + Aggregated training-step output dict. + """ + batch_size = gbs or self.cfg["train_global_batch_size"] + micro_batch_size = mbs or self.cfg["train_micro_batch_size"] + + self._stamp_pad_seqlen(meta) + spa, dba = self._packing_args("train_mb_tokens") + train_meta = replace( + meta, + fields=list(DP_VALUE_TRAIN_FIELDS), + task_name="value_train", + ) + with timer.time("value_training/shard_meta") if timer else nullcontext(): + dp_metas, _ = shard_meta_for_dp( + train_meta, + dp_world=self.sharding_annotations.get_axis_size("data_parallel"), + batch_size=batch_size, + sequence_packing_args=spa, + dynamic_batching_args=dba, + ) + + with ( + timer.time("value_training/submit_training_futures") + if timer + else nullcontext() + ): + futures = self.worker_group.run_all_workers_sharded_data( + "train_presharded", + meta=dp_metas, + in_sharded_axes=["data_parallel"], + replicate_on_axes=_REPLICATED_AXES, + output_is_replicated=_REPLICATED_AXES, + common_kwargs={ + "loss_fn": loss_fn, + "eval_mode": eval_mode, + "gbs": batch_size, + "mbs": micro_batch_size, + }, + ) + return _aggregate_train_results( + self.worker_group.get_all_worker_results(futures) + ) diff --git a/pyrefly.toml b/pyrefly.toml index fc6fa8a1cc5..de3ddb4018a 100644 --- a/pyrefly.toml +++ b/pyrefly.toml @@ -121,6 +121,7 @@ project-includes = [ "nemo_rl/data_plane/adapters/transfer_queue_env.py", "nemo_rl/data_plane/codec.py", "nemo_rl/data_plane/column_io.py", + "nemo_rl/data_plane/driver_mixin.py", "nemo_rl/data_plane/factory.py", "nemo_rl/data_plane/interfaces.py", "nemo_rl/data_plane/observability.py", From b5ba0eb4d2ded4956e2f802527e86c5607f8af8e Mon Sep 17 00:00:00 2001 From: Yuki Huang Date: Sat, 22 Aug 2026 21:52:39 -0700 Subject: [PATCH 02/31] feat(sc): add the worker-side TQ value forward for the PPO critic, with unit tests Signed-off-by: Yuki Huang --- nemo_rl/data_plane/worker_mixin.py | 25 ++ .../value/workers/megatron_value_worker.py | 48 +++- tests/unit/models/value/test_tq_value.py | 214 ++++++++++++++++++ 3 files changed, 286 insertions(+), 1 deletion(-) create mode 100644 tests/unit/models/value/test_tq_value.py diff --git a/nemo_rl/data_plane/worker_mixin.py b/nemo_rl/data_plane/worker_mixin.py index 1125245c98a..88329e73697 100644 --- a/nemo_rl/data_plane/worker_mixin.py +++ b/nemo_rl/data_plane/worker_mixin.py @@ -525,6 +525,31 @@ def get_reference_policy_logprobs_presharded( ) del result + @wrap_with_nvtx_name("value_worker/get_values_presharded") + def get_values_presharded( + self, + meta: "KVBatchMeta", + micro_batch_size: Optional[int] = None, + ) -> None: + """Per-rank value-forward entrypoint. Fetch → packing prep → run → write back. + + Same contract as get_logprobs_presharded, and only the value workers + mix it in: only the PPO critic implements get_values. + """ + data = self._fetch(meta) + data = self._attach_or_repack_pack_metadata(data, meta) + result: BatchedDataDict[Any] = self.get_values( # type: ignore[attr-defined] + data=data, + micro_batch_size=micro_batch_size, + ) + self._write_back_result_field( + meta, + result, + result_key="values", + tq_field="values", + ) + del result + # ── split-API entrypoints (SC async path) ────────────────────────────── # # The split path lets SingleController drive forward/backward per diff --git a/nemo_rl/models/value/workers/megatron_value_worker.py b/nemo_rl/models/value/workers/megatron_value_worker.py index 637a7c9f8be..8abff32b7aa 100644 --- a/nemo_rl/models/value/workers/megatron_value_worker.py +++ b/nemo_rl/models/value/workers/megatron_value_worker.py @@ -52,6 +52,7 @@ from transformers import PreTrainedTokenizerBase from nemo_rl.algorithms.loss.interfaces import LossFunction +from nemo_rl.data_plane.worker_mixin import TQWorkerMixin from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.distributed.model_utils import allgather_cp_sharded_tensor from nemo_rl.distributed.named_sharding import NamedSharding @@ -216,7 +217,7 @@ def _value_loss_prepare_fn( # Classes with @ray.remote can't be inherited from, so we split the implementation out. # This is useful when using worker extension classes. -class MegatronValueWorkerImpl(AbstractPolicyWorker): +class MegatronValueWorkerImpl(TQWorkerMixin, AbstractPolicyWorker): """Megatron-Core based value function worker for PPO. This worker wraps a Megatron-Core GPT model backbone with a value head @@ -238,6 +239,51 @@ def __repr__(self): else: return f"{self.__class__.__qualname__}" + def _local_coords(self) -> dict[str, int]: + """Axis to local-rank mapping. Deliberate copy of MegatronPolicyWorkerImpl.""" + if not torch.distributed.is_initialized(): + return {} + return { + "tensor_parallel": parallel_state.get_tensor_model_parallel_rank(), + "context_parallel": parallel_state.get_context_parallel_rank(), + "pipeline_parallel": parallel_state.get_pipeline_model_parallel_rank(), + } + + def _get_replica_group(self) -> Optional[Any]: + """Replica group = TP x CP x PP siblings within this DP rank. + + Deliberate copy of MegatronPolicyWorkerImpl._get_replica_group; see + there for why it is never gated on CP > 1 and why new_group has to be + called collectively. + """ + if not torch.distributed.is_initialized(): + return None + cached = getattr(self, "_replica_group_cache", "uninit") + if cached != "uninit": + return cached + + world_size = torch.distributed.get_world_size() + my_dp_rank = parallel_state.get_data_parallel_rank() + my_replica_ranks_t = torch.full( + (world_size,), + -1, + dtype=torch.long, + device="cuda", + ) + my_replica_ranks_t[torch.distributed.get_rank()] = my_dp_rank + torch.distributed.all_reduce( + my_replica_ranks_t, op=torch.distributed.ReduceOp.MAX + ) + all_dp_ranks = my_replica_ranks_t.tolist() + + groups: dict[int, Any] = {} + for dp in sorted(set(all_dp_ranks)): + ranks = [r for r, d in enumerate(all_dp_ranks) if d == dp] + grp = torch.distributed.new_group(ranks=ranks, backend="nccl") + groups[dp] = grp + self._replica_group_cache = groups[my_dp_rank] + return self._replica_group_cache + @staticmethod def configure_worker( num_gpus: int | float, diff --git a/tests/unit/models/value/test_tq_value.py b/tests/unit/models/value/test_tq_value.py new file mode 100644 index 00000000000..7080f0a9583 --- /dev/null +++ b/tests/unit/models/value/test_tq_value.py @@ -0,0 +1,214 @@ +# 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. + +"""CPU tests for the TQ-mediated value model (worker wrapper + driver fan-out). + +Same shape as tests/unit/models/policy/test_split_api_wrappers.py: the +GPU-gated PPO tests never reach these two layers cheaply, and the things +that break here are contract-level — a forward pass whose result is +returned through Ray instead of written to TQ, or a train dispatch that +asks workers for a column no producer wrote. +""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from nemo_rl.data_plane import KVBatchMeta +from nemo_rl.data_plane.schema import DP_VALUE_TRAIN_FIELDS, VALUE_SEED_FIELDS +from nemo_rl.data_plane.worker_mixin import TQWorkerMixin +from nemo_rl.models.value.tq_value import TQValue + + +class _ValueStubWorker(TQWorkerMixin): + """Mixin host recording backend calls; fetch/attach are stubbed.""" + + def __init__(self, is_leader: bool = True, values: torch.Tensor | None = None): + self.calls: list[tuple] = [] + self._leader = is_leader + self._dp_client = MagicMock() + self._values = values if values is not None else torch.ones(2, 3) + + def _fetch(self, meta): + self.calls.append(("fetch", meta)) + return {"data_from": meta} + + def _attach_or_repack_pack_metadata(self, data, meta): + self.calls.append(("attach", meta)) + return data + + def _is_replica_leader(self) -> bool: + return self._leader + + def get_values(self, data, micro_batch_size=None): + self.calls.append(("get_values", data, micro_batch_size)) + return {"values": self._values} + + +def _meta(sample_ids: list[str] | None = None) -> KVBatchMeta: + return KVBatchMeta( + partition_id="rollout_data", + task_name="train", + sample_ids=sample_ids if sample_ids is not None else ["s0", "s1"], + ) + + +class TestGetValuesPresharded: + def test_fetches_attaches_then_writes_values_back(self): + w = _ValueStubWorker() + meta = _meta() + + with patch("nemo_rl.data_plane.column_io.write_columns") as write_columns: + out = w.get_values_presharded(meta=meta, micro_batch_size=4) + + # The [B, S] tensor goes to TQ, not through Ray. + assert out is None + assert [c[0] for c in w.calls] == ["fetch", "attach", "get_values"] + assert w.calls[2][1] == {"data_from": meta} + assert w.calls[2][2] == 4 + written = write_columns.call_args.args[2] + assert torch.equal(written["values"], torch.ones(2, 3)) + + def test_non_leader_twin_does_not_write(self): + """TP/CP/PP twins hold identical copies; a second writer is the + duplicate-write bug the leader gate exists to prevent.""" + w = _ValueStubWorker(is_leader=False) + + with patch("nemo_rl.data_plane.column_io.write_columns") as write_columns: + w.get_values_presharded(meta=_meta()) + + write_columns.assert_not_called() + + def test_rejects_batch_dim_mismatch(self): + w = _ValueStubWorker(values=torch.ones(3, 3)) + + with pytest.raises(ValueError, match="shape mismatch"): + w.get_values_presharded(meta=_meta(["s0", "s1"])) + + +def _make_tq_value() -> tuple[TQValue, MagicMock]: + """Bare TQValue with the attributes the fan-out touches.""" + v = object.__new__(TQValue) + v.cfg = {"train_global_batch_size": 8, "train_micro_batch_size": 2} + wg = MagicMock() + v.worker_group = wg + v.sharding_annotations = MagicMock() + v.sharding_annotations.get_axis_size.return_value = 2 + return v, wg + + +class TestTQValueFanout: + def test_get_values_from_meta_narrows_fields_and_returns_none(self): + v, wg = _make_tq_value() + meta = _meta() + with ( + patch.object(TQValue, "_stamp_pad_seqlen"), + patch.object(TQValue, "_packing_args", return_value=(None, None)), + patch( + "nemo_rl.models.value.tq_value.shard_meta_for_dp", + return_value=([meta, meta], None), + ) as mock_shard, + ): + out = v.get_values_from_meta(meta) + + assert out is None + value_meta = mock_shard.call_args.args[0] + assert value_meta.fields == list(VALUE_SEED_FIELDS) + assert value_meta.task_name == "value_fwd" + assert ( + wg.run_all_workers_sharded_data.call_args.args[0] == "get_values_presharded" + ) + wg.get_all_worker_results.assert_called_once() + + def test_get_values_from_meta_uses_the_logprob_packing_budget(self): + """The forward pass is inference-shaped, so it must size microbatches + off logprob_mb_tokens rather than the (larger) train budget.""" + v, _ = _make_tq_value() + meta = _meta() + with ( + patch.object(TQValue, "_stamp_pad_seqlen"), + patch.object( + TQValue, "_packing_args", return_value=(None, None) + ) as mock_packing, + patch( + "nemo_rl.models.value.tq_value.shard_meta_for_dp", + return_value=([meta, meta], None), + ), + ): + v.get_values_from_meta(meta) + + assert mock_packing.call_args.args[0] == "logprob_mb_tokens" + + def test_train_from_meta_requests_the_value_train_columns(self): + v, wg = _make_tq_value() + meta = _meta() + wg.get_all_worker_results.return_value = [ + { + "global_loss": 1.0, + "grad_norm": 0.5, + "all_mb_metrics": {"loss": [0.1]}, + } + ] + with ( + patch.object(TQValue, "_stamp_pad_seqlen"), + patch.object(TQValue, "_packing_args", return_value=(None, None)), + patch( + "nemo_rl.models.value.tq_value.shard_meta_for_dp", + return_value=([meta, meta], None), + ) as mock_shard, + ): + out = v.train_from_meta(meta, loss_fn="LF") + + train_meta = mock_shard.call_args.args[0] + assert train_meta.fields == list(DP_VALUE_TRAIN_FIELDS) + assert "returns" in train_meta.fields and "values" in train_meta.fields + # advantages / prev_logprobs are the policy's business only. + assert "advantages" not in train_meta.fields + assert train_meta.task_name == "value_train" + assert wg.run_all_workers_sharded_data.call_args.args[0] == "train_presharded" + assert wg.run_all_workers_sharded_data.call_args.kwargs["common_kwargs"] == { + "loss_fn": "LF", + "eval_mode": False, + "gbs": 8, + "mbs": 2, + } + assert out["loss"] == 1.0 + assert out["all_mb_metrics"]["loss"] == [0.1] + + def test_train_from_meta_concatenates_per_rank_metrics(self): + v, wg = _make_tq_value() + meta = _meta() + + def _result(loss: float) -> dict: + return { + "global_loss": 1.0, + "grad_norm": 0.5, + "all_mb_metrics": {"loss": [loss]}, + } + + wg.get_all_worker_results.return_value = [_result(0.1), _result(0.2)] + with ( + patch.object(TQValue, "_stamp_pad_seqlen"), + patch.object(TQValue, "_packing_args", return_value=(None, None)), + patch( + "nemo_rl.models.value.tq_value.shard_meta_for_dp", + return_value=([meta, meta], None), + ), + ): + out = v.train_from_meta(meta, loss_fn="LF") + + assert out["all_mb_metrics"]["loss"] == [0.1, 0.2] From f7e90f3a60353643c44d3a69cda409f013e97579 Mon Sep 17 00:00:00 2001 From: Yuki Huang Date: Sat, 22 Aug 2026 23:08:11 -0700 Subject: [PATCH 03/31] refactor: move the PPO estimator config to a BaseModel, pin estimator names to Literals Signed-off-by: Yuki Huang --- nemo_rl/algorithms/advantage_estimator.py | 38 ++++++--- nemo_rl/algorithms/ppo.py | 35 ++------- tests/unit/algorithms/test_ppo.py | 94 +++++++++-------------- 3 files changed, 67 insertions(+), 100 deletions(-) diff --git a/nemo_rl/algorithms/advantage_estimator.py b/nemo_rl/algorithms/advantage_estimator.py index 491de0d9414..cd39e8981cb 100644 --- a/nemo_rl/algorithms/advantage_estimator.py +++ b/nemo_rl/algorithms/advantage_estimator.py @@ -28,6 +28,8 @@ - MOPD: https://arxiv.org/abs/2601.02780 """ +from typing import Literal, Optional + import torch from pydantic import BaseModel @@ -42,9 +44,9 @@ class AdvEstimatorConfig(BaseModel, extra="allow"): - """Configuration for advantage estimator (GRPO, GDPO, or Reinforce++).""" + """Configuration for advantage estimator (GRPO, GDPO, OPD, or Reinforce++).""" - name: str = "grpo" # "grpo", "gdpo", or "reinforce_plus_plus" + name: Literal["grpo", "gdpo", "opd", "reinforce_plus_plus"] = "grpo" # GRPO specific normalize_rewards: bool = True use_leave_one_out_baseline: bool = True @@ -54,6 +56,20 @@ class AdvEstimatorConfig(BaseModel, extra="allow"): minus_baseline: bool = True +class GAEConfig(BaseModel, extra="allow"): + """Configuration for the value-model advantage estimators (PPO).""" + + name: Literal["gae", "raw_reward"] = "gae" + gae_lambda: float = 0.95 + gae_gamma: float = 1.0 + normalize_advantages: bool = True + # VAPO decoupled GAE (None = standard GAE, no decoupling) + gae_lambda_value: Optional[float] = None + gae_lambda_policy: Optional[float] = None + # Length-adaptive λ_policy = 1 - 1/(α·l). 0 = disabled. + length_adaptive_alpha: float = 0.0 + + class GRPOAdvantageEstimator: """GRPO-style advantage estimator with leave-one-out baseline. @@ -273,8 +289,8 @@ class RawRewardAdvantageEstimator: No value model, no baselines. Optionally normalizes across the batch. """ - def __init__(self, estimator_config: dict, loss_config: ClippedPGLossConfig): - self.normalize_advantages = estimator_config["normalize_advantages"] + def __init__(self, estimator_config: GAEConfig, loss_config: ClippedPGLossConfig): + self.normalize_advantages = estimator_config.normalize_advantages def compute_advantage(self, prompt_ids, rewards, mask, **kwargs): """Compute advantages as raw rewards expanded to token-level shape. @@ -326,17 +342,17 @@ class GeneralizedAdvantageEstimator: normalize_advantages: If True, normalize advantages globally across batch """ - def __init__(self, estimator_config: dict, loss_config: ClippedPGLossConfig): - self.gae_lambda = estimator_config["gae_lambda"] - self.gae_gamma = estimator_config["gae_gamma"] - self.normalize_advantages = estimator_config["normalize_advantages"] + def __init__(self, estimator_config: GAEConfig, loss_config: ClippedPGLossConfig): + self.gae_lambda = estimator_config.gae_lambda + self.gae_gamma = estimator_config.gae_gamma + self.normalize_advantages = estimator_config.normalize_advantages # VAPO decoupled GAE: separate λ for value returns vs policy advantages. # None for both = standard GAE (use gae_lambda everywhere, no decoupling). - self.gae_lambda_value = estimator_config["gae_lambda_value"] - self.gae_lambda_policy = estimator_config["gae_lambda_policy"] + self.gae_lambda_value = estimator_config.gae_lambda_value + self.gae_lambda_policy = estimator_config.gae_lambda_policy # Length-adaptive λ_policy = 1 - 1/(α·l). 0 = disabled (use fixed λ). - self.length_adaptive_alpha = estimator_config["length_adaptive_alpha"] + self.length_adaptive_alpha = estimator_config.length_adaptive_alpha self.use_kl_in_reward = loss_config.use_kl_in_reward self.kl_coef = loss_config.reference_policy_kl_penalty diff --git a/nemo_rl/algorithms/ppo.py b/nemo_rl/algorithms/ppo.py index de83ff0e07f..22148c153b8 100644 --- a/nemo_rl/algorithms/ppo.py +++ b/nemo_rl/algorithms/ppo.py @@ -27,6 +27,7 @@ from transformers.tokenization_utils_base import PreTrainedTokenizerBase from nemo_rl.algorithms.advantage_estimator import ( + GAEConfig, GeneralizedAdvantageEstimator, RawRewardAdvantageEstimator, ) @@ -150,21 +151,6 @@ def resolved_warmup_generation_lead_steps(self) -> int: return self.warmup_generation_lead_steps -class AdvEstimatorConfig(TypedDict): - """Configuration for PPO advantage estimator (GAE or raw_reward).""" - - name: str # "gae" or "raw_reward" - # GAE-specific (only used when name="gae") - gae_lambda: NotRequired[float] - gae_gamma: NotRequired[float] - normalize_advantages: NotRequired[bool] - # VAPO decoupled GAE (None = standard GAE, no decoupling) - gae_lambda_value: NotRequired[Optional[float]] - gae_lambda_policy: NotRequired[Optional[float]] - # Length-adaptive λ_policy = 1 - 1/(α·l). 0 = disabled. - length_adaptive_alpha: NotRequired[float] - - class PPOConfig(BaseModel, extra="allow"): num_prompts_per_step: int = 32 num_generations_per_prompt: int = 16 @@ -193,18 +179,7 @@ class PPOConfig(BaseModel, extra="allow"): ppo_epochs: int = 4 reward_shaping: RewardShapingConfig = Field(default_factory=RewardShapingConfig) reward_scaling: RewardScalingConfig = Field(default_factory=RewardScalingConfig) - # Advantage estimator configuration (gae or raw_reward) - adv_estimator: AdvEstimatorConfig = Field( - default_factory=lambda: AdvEstimatorConfig( - name="gae", - gae_lambda=0.95, - gae_gamma=1.0, - normalize_advantages=True, - gae_lambda_value=None, - gae_lambda_policy=None, - length_adaptive_alpha=0.0, - ) - ) + adv_estimator: GAEConfig = Field(default_factory=GAEConfig) # Number of PPO steps of critic-only warmup before policy training begins. # Value model trains from step 0; policy training is skipped for # total_steps < this value. Default 0 (train from start). @@ -1140,11 +1115,11 @@ def _create_advantage_estimator(master_config: MasterConfig): adv_estimator_config = ppo_config.adv_estimator - adv_estimator_name = adv_estimator_config["name"] + adv_estimator_name = adv_estimator_config.name if adv_estimator_name == "gae": adv_estimator = GeneralizedAdvantageEstimator(adv_estimator_config, loss_config) - gae_lambda = adv_estimator_config["gae_lambda"] - gae_gamma = adv_estimator_config["gae_gamma"] + gae_lambda = adv_estimator_config.gae_lambda + gae_gamma = adv_estimator_config.gae_gamma print(f" ✓ Using GAE advantage estimator (λ={gae_lambda}, γ={gae_gamma})") elif adv_estimator_name == "raw_reward": adv_estimator = RawRewardAdvantageEstimator(adv_estimator_config, loss_config) diff --git a/tests/unit/algorithms/test_ppo.py b/tests/unit/algorithms/test_ppo.py index 46bca0ec802..9defc2e30da 100644 --- a/tests/unit/algorithms/test_ppo.py +++ b/tests/unit/algorithms/test_ppo.py @@ -20,6 +20,7 @@ import torch from nemo_rl.algorithms.advantage_estimator import ( + GAEConfig, GeneralizedAdvantageEstimator, RawRewardAdvantageEstimator, ) @@ -51,32 +52,6 @@ def _make_loss_config( ) -def _make_gae_config( - gae_lambda: float = 0.95, - gae_gamma: float = 1.0, - normalize_advantages: bool = False, - length_adaptive_alpha: float = 0.0, - gae_lambda_value: float | None = None, - gae_lambda_policy: float | None = None, - **overrides, -) -> dict: - """Build an estimator_config dict with all GAE-required keys populated. - - ``GeneralizedAdvantageEstimator.__init__`` requires every field to be - present (no hidden ``.get()`` defaults). VAPO fields default to ``None`` - (standard GAE, no decoupling) and can be overridden via kwargs. - """ - return { - "gae_lambda": gae_lambda, - "gae_gamma": gae_gamma, - "normalize_advantages": normalize_advantages, - "length_adaptive_alpha": length_adaptive_alpha, - "gae_lambda_value": gae_lambda_value, - "gae_lambda_policy": gae_lambda_policy, - **overrides, - } - - # ============================================================================ # Tests for GeneralizedAdvantageEstimator # ============================================================================ @@ -88,7 +63,7 @@ def test_gae_basic_computation(): With gamma=1.0 and lambda=1.0, GAE reduces to Monte Carlo returns minus values, so advantages = cumulative_rewards_from_t - V(s_t). """ - estimator_config = _make_gae_config(gae_lambda=1.0, gae_gamma=1.0) + estimator_config = GAEConfig(gae_lambda=1.0, normalize_advantages=False) loss_config = _make_loss_config(kl_penalty=0.0) estimator = GeneralizedAdvantageEstimator(estimator_config, loss_config) @@ -118,7 +93,7 @@ def test_gae_gamma_lambda_zero(): With lambda=0: A_t = delta_t = r_t + gamma * V(s_{t+1}) - V(s_t). Only immediate TD error, no bootstrapping. """ - estimator_config = _make_gae_config(gae_lambda=0.0, gae_gamma=1.0) + estimator_config = GAEConfig(gae_lambda=0.0, normalize_advantages=False) loss_config = _make_loss_config(kl_penalty=0.0) estimator = GeneralizedAdvantageEstimator(estimator_config, loss_config) @@ -145,7 +120,7 @@ def test_gae_gamma_lambda_zero(): def test_gae_shape_and_masking(): """Test that GAE correctly handles masked (padding) positions.""" - estimator_config = _make_gae_config(gae_lambda=0.95, gae_gamma=1.0) + estimator_config = GAEConfig(normalize_advantages=False) loss_config = _make_loss_config(kl_penalty=0.0) estimator = GeneralizedAdvantageEstimator(estimator_config, loss_config) @@ -176,9 +151,7 @@ def test_gae_shape_and_masking(): def test_gae_normalize_advantages(): """Test that advantage normalization produces zero mean and unit variance.""" - estimator_config = _make_gae_config( - gae_lambda=0.95, gae_gamma=1.0, normalize_advantages=True - ) + estimator_config = GAEConfig(normalize_advantages=True) loss_config = _make_loss_config(kl_penalty=0.0) estimator = GeneralizedAdvantageEstimator(estimator_config, loss_config) @@ -205,7 +178,7 @@ def test_gae_normalize_advantages(): def test_gae_kl_penalty_in_rewards(): """Test KL penalty injection into token-level rewards (gated + applied).""" - estimator_config = _make_gae_config(gae_lambda=1.0, gae_gamma=1.0) + estimator_config = GAEConfig(gae_lambda=1.0, normalize_advantages=False) loss_config = _make_loss_config(kl_penalty=0.1, kl_type="k1", use_kl_in_reward=True) estimator = GeneralizedAdvantageEstimator(estimator_config, loss_config) @@ -274,11 +247,8 @@ def test_gae_kl_penalty_in_rewards(): def test_gae_vapo_decoupled_lambda(): """Test VAPO decoupled GAE: separate lambda for value vs policy.""" - base_config = _make_gae_config( - gae_lambda=0.95, - gae_gamma=1.0, - gae_lambda_value=1.0, - gae_lambda_policy=0.5, + base_config = GAEConfig( + gae_lambda_value=1.0, gae_lambda_policy=0.5, normalize_advantages=False ) loss_config = _make_loss_config(kl_penalty=0.0) estimator = GeneralizedAdvantageEstimator(base_config, loss_config) @@ -308,9 +278,7 @@ def test_gae_vapo_decoupled_lambda(): def test_gae_length_adaptive_lambda(): """Test VAPO length-adaptive lambda: lambda_policy = 1 - 1/(alpha * length).""" - config = _make_gae_config( - gae_lambda=0.95, gae_gamma=1.0, length_adaptive_alpha=0.05 - ) + config = GAEConfig(length_adaptive_alpha=0.05, normalize_advantages=False) loss_config = _make_loss_config(kl_penalty=0.0) estimator = GeneralizedAdvantageEstimator(config, loss_config) @@ -338,7 +306,7 @@ def test_gae_carry_forward_interior_gap(): corrupt the GAE accumulators: the advantages at the valid tokens must match the case where the gap is simply removed from the sequence. """ - estimator_config = _make_gae_config(gae_lambda=0.95, gae_gamma=0.99) + estimator_config = GAEConfig(gae_gamma=0.99, normalize_advantages=False) loss_config = _make_loss_config(kl_penalty=0.0) estimator = GeneralizedAdvantageEstimator(estimator_config, loss_config) @@ -387,7 +355,7 @@ def test_raw_reward_basic_broadcast_and_masking(): masked-out trailing pads — the downstream loss masking is responsible for zeroing those out). ``returns`` is None since there is no value head. """ - estimator_config = {"normalize_advantages": False} + estimator_config = GAEConfig(name="raw_reward", normalize_advantages=False) loss_config = _make_loss_config(kl_penalty=0.0) estimator = RawRewardAdvantageEstimator(estimator_config, loss_config) @@ -634,18 +602,11 @@ def test_create_advantage_estimator_gae(): from nemo_rl.algorithms.ppo import _create_advantage_estimator - # adv_estimator dict needs every GAE-required key (no hidden .get() defaults - # in the estimator __init__); loss_fn must be a real ClippedPGLossConfig - # because the estimator accesses .use_kl_in_reward / .reference_policy_kl_* - # as attributes, not dict keys. + # loss_fn must be a real ClippedPGLossConfig: the estimator reads + # .use_kl_in_reward / .reference_policy_kl_* off it. master_config = SimpleNamespace( ppo=PPOConfig( - adv_estimator={ - "name": "gae", - **_make_gae_config( - gae_lambda=0.95, gae_gamma=1.0, normalize_advantages=True - ), - }, + adv_estimator=GAEConfig(), ), loss_fn=_make_loss_config(kl_penalty=0.0), ) @@ -672,14 +633,32 @@ def test_create_advantage_estimator_raw_reward(): assert isinstance(estimator, RawRewardAdvantageEstimator) +def test_ppo_schema_rejects_unsupported_estimator_name(): + """PPO consumes (advantages, returns), which only GAE/raw_reward return. + + The schema is the real gate, so a group-relative name fails at config load + rather than at the factory. + """ + from pydantic import ValidationError + + with pytest.raises(ValidationError): + PPOConfig(adv_estimator={"name": "grpo"}) + + def test_create_advantage_estimator_rejects_unsupported_name(): - """PPO loop only consumes (advantages, returns) from GAE/raw_reward — others must error.""" + """The factory still guards names that skipped schema validation. + + model_construct bypasses the Literal, so this branch is reachable and is + what turns a bad name into a message instead of an UnboundLocalError. + """ from types import SimpleNamespace from nemo_rl.algorithms.ppo import _create_advantage_estimator master_config = SimpleNamespace( - ppo=PPOConfig(adv_estimator={"name": "grpo"}), + ppo=PPOConfig.model_construct( + adv_estimator=GAEConfig.model_construct(name="grpo") + ), loss_fn={"reference_policy_kl_penalty": 0.0}, ) @@ -2485,10 +2464,7 @@ def test_async_ppo_initial_refit_failure_cleans_up_actors(monkeypatch): config.ppo.num_prompts_per_step = 1 config.ppo.max_rollout_turns = 1 config.ppo.skip_reference_policy_logprobs_calculation = False - config.ppo.adv_estimator = { - "name": "raw_reward", - "normalize_advantages": False, - } + config.ppo.adv_estimator = GAEConfig(name="raw_reward", normalize_advantages=False) config.checkpointing = { "checkpoint_must_save_by": None, "ft_save_period": None, From d5973422ccff4b221cb014a13823f2ce9a52e24e Mon Sep 17 00:00:00 2001 From: Yuki Huang Date: Sun, 23 Aug 2026 00:08:39 -0700 Subject: [PATCH 04/31] feat(sc): add the ppo algorithm block to the SingleController config Signed-off-by: Yuki Huang --- nemo_rl/algorithms/single_controller.py | 32 +++-- .../single_controller_utils/__init__.py | 4 + .../single_controller_utils/config.py | 115 ++++++++++++++++-- .../single_controller_utils/setup.py | 25 ++-- .../single_controller/test_rollout_pump.py | 8 ++ .../test_single_controller.py | 5 + .../single_controller/test_watchdog_pump.py | 1 + 7 files changed, 152 insertions(+), 38 deletions(-) diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 8a96e6c62a0..5a45c28ef37 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -56,6 +56,7 @@ from nemo_rl.algorithms.single_controller_utils.config import ( AdvantageConfig, MasterConfig, + algo_config, validate_sampler_buffer_capacity, validate_single_controller_config, ) @@ -126,14 +127,15 @@ def __init__( self._partition_id: str = actor_args.partition_id self._master_config = master_config + self._algo_cfg = algo_config(master_config) self._async_cfg = master_config.async_rl self._policy_logprobs_required = not ( master_config.loss_fn.force_on_policy_ratio - and master_config.grpo.seq_logprob_error_threshold is None + and self._algo_cfg.seq_logprob_error_threshold is None ) self._reference_logprobs_required = bool( master_config.loss_fn.reference_policy_kl_penalty > 0 - and not master_config.grpo.skip_reference_policy_logprobs_calculation + and not self._algo_cfg.skip_reference_policy_logprobs_calculation ) self._dp_client = actor_args.dp_client self._gen: Generation = actor_args.gen_handle @@ -191,7 +193,7 @@ def __init__( self._train_cluster = actor_args.train_cluster self._inference_cluster = actor_args.inference_cluster - num_prompts_per_step = self._master_config.grpo.num_prompts_per_step + num_prompts_per_step = self._algo_cfg.num_prompts_per_step self._sampler = create_sampler(self._buffer, self._async_cfg.sampler) self._sampler.set_dispatch_index(actor_args.save_state.current_step) required_capacity = self._sampler.required_buffer_capacity(num_prompts_per_step) @@ -382,7 +384,7 @@ async def _maybe_restore_replay_buffer(self) -> None: buffer_state, max_groups=self._async_cfg.max_buffered_rollouts, expected_partition_id=self._partition_id, - expected_group_size=self._master_config.grpo.num_generations_per_prompt, + expected_group_size=self._algo_cfg.num_generations_per_prompt, ) # Each buffered group holds one _buffer_capacity permit; the load # truncation guarantees restored <= capacity, so this never blocks. @@ -582,7 +584,7 @@ async def _launch(prompt: DatumSpec, target_step: Optional[int]) -> None: ) ) - max_epochs = self._master_config.grpo.max_num_epochs + max_epochs = self._algo_cfg.max_num_epochs async with asyncio.TaskGroup() as rollout_tasks: while max_epochs is None or self._current_epoch < max_epochs: for prompt_batch in self._dataloader: @@ -696,7 +698,7 @@ async def _drain_reserve_into_steps( only "replace" ever fills one, so the gate would buy nothing while stranding a pool restored from a checkpoint into a run that has since switched to "shrink". """ - num_prompts_per_step = self._master_config.grpo.num_prompts_per_step + num_prompts_per_step = self._algo_cfg.num_prompts_per_step while len(self._replacement_reserve) >= num_prompts_per_step: # Take the step's prompts out before the first await. A drop resolving # concurrently draws from this same pool, and could otherwise claim one of @@ -822,7 +824,7 @@ def _target_groups_for_step(self, step: int) -> int: ``num_prompts_per_step``. Training a fraction of a batch is a silent change to the gradient estimate, so it is refused rather than absorbed. """ - num_prompts_per_step = self._master_config.grpo.num_prompts_per_step + num_prompts_per_step = self._algo_cfg.num_prompts_per_step dropped = self._batch_shortfall.get(step, 0) target = num_prompts_per_step - dropped fraction = self._async_cfg.rollout_failure.min_step_batch_fraction @@ -854,9 +856,7 @@ async def _train_pump(self) -> None: 5. dp_client.clear_samples on consumed sample_ids; release _buffer_capacity per dropped group, then sync. """ - grpo_cfg = self._master_config.grpo - - while self._train_steps < grpo_cfg.max_num_steps: + while self._train_steps < self._algo_cfg.max_num_steps: version_during_step = self._trainer_version groups_dispatched = 0 evicted_stale_prompt_groups = 0 @@ -1051,7 +1051,7 @@ async def _train_pump(self) -> None: chunks_dispatched, num_groups, groups_dispatched, - grpo_cfg.num_prompts_per_step, + self._algo_cfg.num_prompts_per_step, ) log.info( @@ -1146,7 +1146,7 @@ async def _train_pump(self) -> None: self._total_valid_tokens += step_metrics.get("global_valid_toks", 0) self._timeout.mark_iteration() - is_last_step = self._train_steps >= grpo_cfg.max_num_steps or ( + is_last_step = self._train_steps >= self._algo_cfg.max_num_steps or ( self._rollout_exhausted.is_set() and len(self._buffer) == 0 ) ft_save_period = self._master_config.checkpointing.get("ft_save_period") @@ -1211,7 +1211,7 @@ async def _train_pump(self) -> None: # generated with; lag = training version - oldest sample version. lag = version_during_step - min_sample_version # type: ignore print( - f"train step {self._train_steps}/{grpo_cfg.max_num_steps} " + f"train step {self._train_steps}/{self._algo_cfg.max_num_steps} " f"trainer_v={self._trainer_version} " f"lag={lag} ", flush=True, @@ -1240,7 +1240,7 @@ async def _stall_watchdog_pump(self) -> None: is what is checked instead. """ watchdog_cfg = self._async_cfg.stall_watchdog - max_num_steps = self._master_config.grpo.max_num_steps + max_num_steps = self._algo_cfg.max_num_steps last_progress = (-1, -1) last_progress_at = time.monotonic() @@ -1706,9 +1706,7 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> tuple[KVBatchMeta, bool]: tensor_field(data, adv_cfg.sample_mask_field) ).float() - seq_logprob_error_threshold = ( - self._master_config.grpo.seq_logprob_error_threshold - ) + seq_logprob_error_threshold = self._algo_cfg.seq_logprob_error_threshold # Match the legacy path: whenever real policy logprobs are available, # report sequence-level generation/training mismatch. A threshold adds # masking; leaving it unset keeps this metrics-only. diff --git a/nemo_rl/algorithms/single_controller_utils/__init__.py b/nemo_rl/algorithms/single_controller_utils/__init__.py index 198bea987a1..425c1a9f48f 100644 --- a/nemo_rl/algorithms/single_controller_utils/__init__.py +++ b/nemo_rl/algorithms/single_controller_utils/__init__.py @@ -20,6 +20,8 @@ MasterConfig, RolloutFailureConfig, WatchdogConfig, + algo_config, + is_ppo_run, ) from nemo_rl.algorithms.single_controller_utils.setup import ( SingleControllerActorArgs, @@ -33,5 +35,7 @@ "RolloutFailureConfig", "SingleControllerActorArgs", "WatchdogConfig", + "algo_config", + "is_ppo_run", "setup_single_controller", ] diff --git a/nemo_rl/algorithms/single_controller_utils/config.py b/nemo_rl/algorithms/single_controller_utils/config.py index b69e2d625ee..3a6544c6fcc 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/nemo_rl/algorithms/single_controller_utils/config.py @@ -36,10 +36,13 @@ ) from nemo_rl.algorithms.grpo import GRPOConfig, GRPOLoggerConfig from nemo_rl.algorithms.loss import ClippedPGLossConfig +from nemo_rl.algorithms.loss.loss_functions import MseValueLossConfig +from nemo_rl.algorithms.ppo import PPOConfig from nemo_rl.data import DataConfig from nemo_rl.data_plane.interfaces import DataPlaneConfig from nemo_rl.distributed.virtual_cluster import ClusterConfig from nemo_rl.models.policy import PolicyConfig +from nemo_rl.models.value import ValueConfig from nemo_rl.utils.checkpoint import CheckpointingConfig # ── User-facing SingleController configs ──────────────────────────────────── @@ -534,11 +537,16 @@ def _reject_relocated_keys(self) -> "AsyncRLConfig": class MasterConfig(BaseModel, extra="allow"): + # algo configs + grpo: Optional[GRPOConfig] = None + ppo: Optional[PPOConfig] = None policy: PolicyConfig + value: Optional[ValueConfig] = None # PPO extras loss_fn: ClippedPGLossConfig + value_loss_fn: Optional[MseValueLossConfig] = None # PPO extras + # common configs env: dict[str, Any] data: DataConfig - grpo: GRPOConfig logger: GRPOLoggerConfig cluster: ClusterConfig checkpointing: CheckpointingConfig @@ -546,6 +554,28 @@ class MasterConfig(BaseModel, extra="allow"): async_rl: AsyncRLConfig +def is_ppo_run(master_config: MasterConfig) -> bool: + """Whether this SingleController run trains a PPO critic alongside the policy. + + Single source of truth for the flag: setup reads it to decide whether to + build the value model, and the controller reads it to decide whether the + train pump runs the critic stages. ``model_construct`` skips defaults, so + the attribute can genuinely be missing on a hand-built config. + """ + return getattr(master_config, "ppo", None) is not None + + +def algo_config(master_config: MasterConfig) -> GRPOConfig | PPOConfig: + """The active algorithm block: ``ppo`` on a PPO run, else ``grpo``. + + Exactly one of the two is set; _validate_algo_settings checks that, and it + always runs at setup. + """ + if is_ppo_run(master_config): + return master_config.ppo # type: ignore + return master_config.grpo # type: ignore + + def validate_sampler_buffer_capacity( async_config: AsyncRLConfig, *, @@ -675,19 +705,84 @@ def _validate_failure_settings( ) +def _validate_algo_settings(master_config: MasterConfig) -> None: + """Reject algorithm blocks the SingleController path cannot honour. + + Both directions: a critic the PPO path needs and does not have, and a critic + a GRPO run carries and would never build. + """ + grpo = getattr(master_config, "grpo", None) + ppo = getattr(master_config, "ppo", None) + if grpo is not None and ppo is not None: + raise ValueError("Only one algorithm block can be set, either `grpo` or `ppo`.") + if grpo is None and ppo is None: + raise ValueError( + "At least one algorithm block must be set, either `grpo` or `ppo`." + ) + + algo_cfg = algo_config(master_config) + if not is_ppo_run(master_config): + # A value block without `ppo` is inert -- nothing builds the critic -- + # and a config carrying one is asking for PPO by every reading except + # the one the code uses. Say so rather than training GRPO silently. + for name in ("value", "value_loss_fn"): + if getattr(master_config, name, None) is not None: + raise ValueError( + f"{name} is set but the `ppo` block is absent, so this run " + "trains GRPO and the value model would never be built. Add a " + f"`ppo` block, or remove `{name}`." + ) + return + + for name in ("value", "value_loss_fn"): + if getattr(master_config, name, None) is None: + raise ValueError( + f"the `ppo` block selects the PPO path, which needs `{name}`. " + "See examples/configs/ppo_math_1B_megatron_single_controller.yaml." + ) + + async_config = master_config.async_rl + # Without it the critic steps once per chunk and the policy once per step, + # which is two effective learning rates from one config, and no error. + if async_config.min_groups_for_streaming_train != algo_cfg.num_prompts_per_step: + raise ValueError( + "PPO on the SingleController path requires " + "async_rl.min_groups_for_streaming_train " + f"({async_config.min_groups_for_streaming_train}) == " + f"num_prompts_per_step ({algo_cfg.num_prompts_per_step}) so that each RL " + "step is assembled from a single chunk: the critic steps its " + "optimizer once per chunk and the policy once per step. Streaming " + "PPO needs a split train API on the value workers, which they do " + "not have yet (#2625)." + ) + + rl_step_samples = ( + algo_cfg.num_prompts_per_step * algo_cfg.num_generations_per_prompt + ) + value_global_batch_size = master_config.value["train_global_batch_size"] # type: ignore + if rl_step_samples != value_global_batch_size: + raise ValueError( + "num_prompts_per_step * num_generations_per_prompt " + f"({rl_step_samples}) must equal value.train_global_batch_size " + f"({value_global_batch_size}) so that one RL step maps to exactly one " + "critic optimizer.step." + ) + + def validate_single_controller_config(master_config: MasterConfig) -> None: """Validate cross-section SingleController constraints before setup.""" async_config = master_config.async_rl - num_prompts_per_step = master_config.grpo.num_prompts_per_step - if num_prompts_per_step < async_config.min_groups_for_streaming_train: + algo_cfg = algo_config(master_config) + + if algo_cfg.num_prompts_per_step < async_config.min_groups_for_streaming_train: raise ValueError( - f"grpo.num_prompts_per_step ({num_prompts_per_step}) " + f"grpo.num_prompts_per_step ({algo_cfg.num_prompts_per_step}) " f"must be >= async_rl.min_groups_for_streaming_train " f"({async_config.min_groups_for_streaming_train})" ) rl_step_samples = ( - num_prompts_per_step * master_config.grpo.num_generations_per_prompt + algo_cfg.num_prompts_per_step * algo_cfg.num_generations_per_prompt ) train_global_batch_size = master_config.policy["train_global_batch_size"] if rl_step_samples != train_global_batch_size: @@ -701,7 +796,7 @@ def validate_single_controller_config(master_config: MasterConfig) -> None: required_capacity = required_buffer_capacity_for_config( async_config.sampler, - num_prompts_per_step, + algo_cfg.num_prompts_per_step, ) validate_sampler_buffer_capacity( async_config, @@ -755,7 +850,7 @@ def validate_single_controller_config(master_config: MasterConfig) -> None: if ( reference_policy_kl_penalty - and master_config.grpo.skip_reference_policy_logprobs_calculation + and algo_cfg.skip_reference_policy_logprobs_calculation ): raise ValueError( "loss_fn.reference_policy_kl_penalty=" @@ -768,7 +863,7 @@ def validate_single_controller_config(master_config: MasterConfig) -> None: if ( reference_policy_kl_penalty == 0 - and not master_config.grpo.skip_reference_policy_logprobs_calculation + and not algo_cfg.skip_reference_policy_logprobs_calculation ): print( "Reference policy logprob calculation will be skipped since " @@ -776,7 +871,9 @@ def validate_single_controller_config(master_config: MasterConfig) -> None: "model was initialized." ) - _validate_failure_settings(async_config, num_prompts_per_step) + _validate_algo_settings(master_config) + + _validate_failure_settings(async_config, algo_cfg.num_prompts_per_step) # Nesting says which knob applies to which path, but nothing stops an operator # filling in the block for the path this run is not taking -- and a populated diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py index 5b8cca7bc3d..58d97ab9378 100644 --- a/nemo_rl/algorithms/single_controller_utils/setup.py +++ b/nemo_rl/algorithms/single_controller_utils/setup.py @@ -50,6 +50,7 @@ ) from nemo_rl.algorithms.single_controller_utils.config import ( MasterConfig, + algo_config, validate_single_controller_config, ) from nemo_rl.algorithms.utils import set_seed @@ -370,12 +371,12 @@ def _clamp_max_num_steps( master_config: MasterConfig, dataloader: StatefulDataLoader ) -> None: """Clamp grpo.max_num_steps to max_num_epochs * len(dataloader).""" - grpo_config = master_config.grpo - max_num_epochs = grpo_config.max_num_epochs + algo_cfg = algo_config(master_config) + max_num_epochs = algo_cfg.max_num_epochs if max_num_epochs is None: return - grpo_config.max_num_steps = min( - grpo_config.max_num_steps, + algo_cfg.max_num_steps = min( + algo_cfg.max_num_steps, max_num_epochs * len(dataloader), ) @@ -385,8 +386,8 @@ def _maybe_inject_megatron_train_iters(master_config: MasterConfig) -> None: policy_config = master_config.policy if not policy_config.get("megatron_cfg", {}).get("enabled", False): return - grpo_config = master_config.grpo - policy_config["megatron_cfg"]["train_iters"] = grpo_config.max_num_steps + algo_cfg = algo_config(master_config) + policy_config["megatron_cfg"]["train_iters"] = algo_cfg.max_num_steps def _maybe_attach_fleet_health( @@ -525,13 +526,13 @@ def setup_single_controller( validate_single_controller_config(master_config) # short names for config sections - grpo_config = master_config.grpo + algo_cfg = algo_config(master_config) dp_config = master_config.data_plane policy_config = master_config.policy generation_config = policy_config["generation"] data_config = master_config.data - if grpo_config.val_period > 0 or grpo_config.val_at_start or grpo_config.val_at_end: + if algo_cfg.val_period > 0 or algo_cfg.val_at_start or algo_cfg.val_at_end: raise NotImplementedError( "SingleController doesn't support validation now, will support " "later. Set grpo.val_period=0, val_at_start=false, val_at_end=false." @@ -557,7 +558,7 @@ def setup_single_controller( if checkpointing_pretrained is not None: policy_config["pretrained_checkpoint"] = checkpointing_pretrained - set_seed(grpo_config.seed) + set_seed(algo_cfg.seed) # ========================== # Checkpointing @@ -595,7 +596,7 @@ def setup_single_controller( dataset, _val_dataset, env_handles, _val_env_handles = response_data dataloader = StatefulDataLoader( dataset, - batch_size=grpo_config.num_prompts_per_step, + batch_size=algo_cfg.num_prompts_per_step, shuffle=data_config["shuffle"], collate_fn=rl_collate_fn, drop_last=True, @@ -797,9 +798,9 @@ def _build_generation_then_trainer( rollout_manager = RolloutManager( tokenizer=tokenizer, task_to_env=env_handles, - num_generations_per_prompt=grpo_config.num_generations_per_prompt, + num_generations_per_prompt=algo_cfg.num_generations_per_prompt, max_seq_len=_generation_max_seq_len(generation_config), - max_rollout_turns=grpo_config.max_rollout_turns, + max_rollout_turns=algo_cfg.max_rollout_turns, policy_generation=generation, generation_config=generation_config, use_nemo_gym=use_nemo_gym, diff --git a/tests/unit/single_controller/test_rollout_pump.py b/tests/unit/single_controller/test_rollout_pump.py index 3859af8dcf3..94fb220bb6e 100644 --- a/tests/unit/single_controller/test_rollout_pump.py +++ b/tests/unit/single_controller/test_rollout_pump.py @@ -164,6 +164,7 @@ def test_rollout_pump_stamps_target_steps( ctrl._master_config = SimpleNamespace( grpo=GRPOConfig.model_construct(max_num_epochs=1) ) + ctrl._algo_cfg = ctrl._master_config.grpo ctrl._rollout_manager = _RecordingRolloutManager(buffer) # The sampler owns admission + target_step stamping (the dispatch counter # lives on the sampler, not the actor). @@ -223,6 +224,7 @@ async def generate_and_push( ctrl._master_config = SimpleNamespace( grpo=GRPOConfig.model_construct(max_num_epochs=1) ) + ctrl._algo_cfg = ctrl._master_config.grpo ctrl._rollout_manager = _OutcomeRolloutManager() ctrl._sampler = WindowedSampler(None, max_staleness_versions=1) ctrl._dataloader = [ @@ -277,6 +279,7 @@ def test_rollout_pump_tops_up_restored_target_step( ctrl._master_config = SimpleNamespace( grpo=GRPOConfig.model_construct(max_num_epochs=1) ) + ctrl._algo_cfg = ctrl._master_config.grpo ctrl._rollout_manager = _RecordingRolloutManager(buffer) # lookahead=0 keeps the single batch on target_step 0. ctrl._sampler = InOrderSampler(buffer, max_lookahead_versions=0) @@ -354,6 +357,7 @@ def test_rollout_pump_credits_shortfall_only_for_stamped_prompts( ctrl._master_config = SimpleNamespace( grpo=GRPOConfig.model_construct(max_num_epochs=1) ) + ctrl._algo_cfg = ctrl._master_config.grpo ctrl._rollout_manager = _SkippingRolloutManager() ctrl._sampler = make_sampler(buffer) prompt_batch = BatchedDataDict( @@ -433,6 +437,7 @@ def _pump_controller( max_num_epochs=1, num_prompts_per_step=num_prompts_per_step ) ) + ctrl._algo_cfg = ctrl._master_config.grpo ctrl._rollout_manager = manager ctrl._sampler = InOrderSampler(buffer, max_lookahead_versions=1) ctrl._dataloader = dataloader @@ -716,6 +721,7 @@ def _controller( num_prompts_per_step=num_prompts_per_step, ) ) + ctrl._algo_cfg = ctrl._master_config.grpo ctrl._async_cfg = SimpleNamespace( rollout_failure=SimpleNamespace(min_step_batch_fraction=fraction) ) @@ -850,6 +856,7 @@ async def _main() -> None: ctrl._master_config = SimpleNamespace( grpo=GRPOConfig.model_construct(max_num_epochs=1) ) + ctrl._algo_cfg = ctrl._master_config.grpo ctrl._rollout_manager = manager # Over-sampled windowed policy: admit never gates (buffer unused here). ctrl._sampler = WindowedSampler(None, max_staleness_versions=1) @@ -938,6 +945,7 @@ async def _main() -> None: ctrl._master_config = SimpleNamespace( grpo=GRPOConfig.model_construct(max_num_epochs=1) ) + ctrl._algo_cfg = ctrl._master_config.grpo ctrl._rollout_manager = _NeverCalledRolloutManager() # Over-sampled windowed policy: admit never gates (buffer unused here). ctrl._sampler = WindowedSampler(None, max_staleness_versions=1) diff --git a/tests/unit/single_controller/test_single_controller.py b/tests/unit/single_controller/test_single_controller.py index d14ec2a5e23..1be65731a01 100644 --- a/tests/unit/single_controller/test_single_controller.py +++ b/tests/unit/single_controller/test_single_controller.py @@ -413,6 +413,7 @@ def test_advantage_stage_applies_seq_logprob_error_mask_before_streaming_train( ctrl._master_config = SimpleNamespace( grpo=SimpleNamespace(seq_logprob_error_threshold=2.0) ) + ctrl._algo_cfg = ctrl._master_config.grpo ctrl._step_log_dict = { "rewards": [], "masked_advantages": [], @@ -479,6 +480,7 @@ def test_advantage_stage_reports_seq_logprob_metrics_without_masking() -> None: ctrl._master_config = SimpleNamespace( grpo=SimpleNamespace(seq_logprob_error_threshold=None) ) + ctrl._algo_cfg = ctrl._master_config.grpo ctrl._step_log_dict = { "rewards": [], "masked_advantages": [], @@ -539,6 +541,7 @@ def test_advantage_stage_skips_estimator_when_seq_mask_removes_whole_chunk( ctrl._master_config = SimpleNamespace( grpo=SimpleNamespace(seq_logprob_error_threshold=2.0) ) + ctrl._algo_cfg = ctrl._master_config.grpo ctrl._step_log_dict = { "rewards": [], "masked_advantages": [], @@ -592,6 +595,7 @@ def test_advantage_stage_skips_preexisting_empty_mask_without_seq_threshold() -> ctrl._master_config = SimpleNamespace( grpo=SimpleNamespace(seq_logprob_error_threshold=None) ) + ctrl._algo_cfg = ctrl._master_config.grpo ctrl._step_log_dict = { "rewards": [], "masked_advantages": [], @@ -737,6 +741,7 @@ def _train_pump_controller(*, sampler) -> object: # is disabled. checkpointing={"enabled": False, "save_period": 10}, ) + ctrl._algo_cfg = ctrl._master_config.grpo ctrl._async_cfg = SimpleNamespace( min_groups_for_streaming_train=1, rollout_failure=SimpleNamespace(min_step_batch_fraction=0.9), diff --git a/tests/unit/single_controller/test_watchdog_pump.py b/tests/unit/single_controller/test_watchdog_pump.py index 55cb59a7369..fe8ff37b023 100644 --- a/tests/unit/single_controller/test_watchdog_pump.py +++ b/tests/unit/single_controller/test_watchdog_pump.py @@ -81,6 +81,7 @@ def _make_controller( ctrl._master_config = SimpleNamespace( grpo=GRPOConfig.model_construct(max_num_steps=max_num_steps) ) + ctrl._algo_cfg = ctrl._master_config.grpo ctrl._rollout_manager = SimpleNamespace(stats=stats) ctrl._inflight_rollouts = inflight ctrl._train_steps = train_steps From ec4ac13e26790e3f03dfe16cbc6f421171f8b2c0 Mon Sep 17 00:00:00 2001 From: Yuki Huang Date: Sun, 23 Aug 2026 01:04:22 -0700 Subject: [PATCH 05/31] feat(sc): build the PPO value model in SingleController setup, validate the algo block first Signed-off-by: Yuki Huang --- nemo_rl/algorithms/metric_utils.py | 5 + nemo_rl/algorithms/single_controller.py | 10 +- .../single_controller_utils/config.py | 6 +- .../single_controller_utils/setup.py | 171 ++++++-- tests/unit/algorithms/test_metric_utils.py | 11 + .../single_controller/test_sc_ppo_setup.py | 369 ++++++++++++++++++ .../test_single_controller.py | 105 +++-- .../test_single_controller_setup.py | 5 +- 8 files changed, 613 insertions(+), 69 deletions(-) create mode 100644 tests/unit/single_controller/test_sc_ppo_setup.py diff --git a/nemo_rl/algorithms/metric_utils.py b/nemo_rl/algorithms/metric_utils.py index 2384521db98..2425eed0826 100644 --- a/nemo_rl/algorithms/metric_utils.py +++ b/nemo_rl/algorithms/metric_utils.py @@ -29,6 +29,8 @@ class SetupTimingMetrics: generation_init_load_time_s: Optional[float] = None policy_init_time_s: Optional[float] = None + # PPO only: the critic shares the training GPUs, so it is built after the policy. + value_init_time_s: Optional[float] = None nemo_gym_init_time_s: Optional[float] = None collective_init_time_s: Optional[float] = None # Non-colocated megatron's post-init weight sync into the engine. @@ -82,6 +84,9 @@ def print_setup_timing_summary(metrics: SetupTimingMetrics) -> None: print(f" Policy init: {metrics.policy_init_time_s:.1f}s") + if metrics.value_init_time_s: + print(f" Value init: {metrics.value_init_time_s:.1f}s") + if metrics.nemo_gym_init_time_s: print(f" NeMo-Gym init: {metrics.nemo_gym_init_time_s:.1f}s") diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 5a45c28ef37..713a33d1290 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -78,6 +78,7 @@ from nemo_rl.models.generation.sglang.sglang_generation import SGLangGeneration from nemo_rl.models.generation.vllm import VllmGeneration from nemo_rl.models.policy.tq_policy import TQPolicy +from nemo_rl.models.value.tq_value import TQValue from nemo_rl.utils.checkpoint import CheckpointManager, PathLike from nemo_rl.utils.logger import Logger from nemo_rl.utils.timer import TimeoutChecker, Timer @@ -140,12 +141,18 @@ def __init__( self._dp_client = actor_args.dp_client self._gen: Generation = actor_args.gen_handle self._trainer: TQPolicy = actor_args.trainer_handle + self._value: Optional[TQValue] = getattr(actor_args, "value_handle", None) self._dataloader = actor_args.dataloader self._weight_synchronizer = actor_args.weight_synchronizer self._advantage_estimator = actor_args.advantage_estimator self._loss_fn = actor_args.loss_fn + self._value_loss_fn = getattr(actor_args, "value_loss_fn", None) self._buffer = actor_args.tq_buffer self._rollout_manager = actor_args.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 + # Direct access, deliberately. A getattr default here reads as defensive but # buys a silent failure mode: rename or drop the field and # watchdog.gym_subprocess_check: true degrades to a health check that iterates @@ -158,9 +165,6 @@ def __init__( # therefore degrades to the documented off state rather than to a broken one. self._gen_fleet = getattr(actor_args, "fleet_monitor", None) self._generation_router = getattr(actor_args, "generation_router", None) - # 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. diff --git a/nemo_rl/algorithms/single_controller_utils/config.py b/nemo_rl/algorithms/single_controller_utils/config.py index 3a6544c6fcc..971961f5b00 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/nemo_rl/algorithms/single_controller_utils/config.py @@ -771,6 +771,10 @@ def _validate_algo_settings(master_config: MasterConfig) -> None: def validate_single_controller_config(master_config: MasterConfig) -> None: """Validate cross-section SingleController constraints before setup.""" + # First: everything below reads the active algorithm block, which only exists + # once this has confirmed exactly one is set. + _validate_algo_settings(master_config) + async_config = master_config.async_rl algo_cfg = algo_config(master_config) @@ -871,8 +875,6 @@ def validate_single_controller_config(master_config: MasterConfig) -> None: "model was initialized." ) - _validate_algo_settings(master_config) - _validate_failure_settings(async_config, algo_cfg.num_prompts_per_step) # Nesting says which knob applies to which path, but nothing stops an operator diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py index 58d97ab9378..5dd6ba5d91c 100644 --- a/nemo_rl/algorithms/single_controller_utils/setup.py +++ b/nemo_rl/algorithms/single_controller_utils/setup.py @@ -37,20 +37,22 @@ from nemo_rl.algorithms.async_utils.replay_buffer import TQReplayBuffer from nemo_rl.algorithms.grpo import ( GRPOSaveState, - _create_advantage_estimator, _get_effort_config, _get_grpo_save_state, ) -from nemo_rl.algorithms.grpo import MasterConfig as GrpoMasterConfig +from nemo_rl.algorithms.grpo import MasterConfig as GRPOMasterConfig from nemo_rl.algorithms.loss import ClippedPGLossFn from nemo_rl.algorithms.loss.interfaces import LossFunction +from nemo_rl.algorithms.loss.loss_functions import MseValueLossFn from nemo_rl.algorithms.metric_utils import ( SetupTimingMetrics, print_setup_timing_summary, ) +from nemo_rl.algorithms.ppo import MasterConfig as PPOMasterConfig from nemo_rl.algorithms.single_controller_utils.config import ( MasterConfig, algo_config, + is_ppo_run, validate_single_controller_config, ) from nemo_rl.algorithms.utils import set_seed @@ -91,6 +93,7 @@ router_replay_enabled, ) from nemo_rl.models.policy.tq_policy import TQPolicy +from nemo_rl.models.value.tq_value import TQValue from nemo_rl.utils.checkpoint import CheckpointManager from nemo_rl.weight_sync import WeightSynchronizer, create_weight_synchronizer @@ -126,6 +129,10 @@ class SingleControllerActorArgs: # serving backend set to it. Parameterized with the Impl class because the decorated # GenerationRouterActor name is an ActorClass instance, not a type. generation_router: Optional[ray.actor.ActorHandle[GenerationRouterImpl]] = None + # None on a GRPO run. Both are set together on the PPO path: the critic and + # the MSE loss it trains under. + value_handle: Optional[TQValue] = None + value_loss_fn: Optional[LossFunction] = None def _build_clusters( @@ -313,6 +320,40 @@ def _build_trainer( return trainer, time.perf_counter() - t0 +def _build_value( + train_cluster: RayVirtualCluster, + master_config: MasterConfig, + tokenizer: PreTrainedTokenizerBase, + *, + weights_path: Optional[Path], + optimizer_path: Optional[Path], +) -> tuple[TQValue, float]: + """Build the TQ-mediated PPO critic (driver-side TQValue). + + Args: + train_cluster: Ray virtual cluster the critic shares with the trainer. + master_config: SC MasterConfig. + tokenizer: Tokenizer used by the value model. + weights_path: Checkpointed value weights to resume from, or None. + optimizer_path: Checkpointed value optimizer state to resume from, or None. + + Returns: + A tuple of (TQValue critic, wall time spent in this call). + """ + t0 = time.perf_counter() + value = TQValue( + cluster=train_cluster, + config=master_config.value, + tokenizer=tokenizer, + name_prefix="lm_value", + weights_path=weights_path, + optimizer_path=optimizer_path, + init_optimizer=True, + dp_cfg=master_config.data_plane, + ) + return value, time.perf_counter() - t0 + + def _spinup_gym( master_config: MasterConfig, base_urls: list[str], @@ -383,11 +424,19 @@ def _clamp_max_num_steps( def _maybe_inject_megatron_train_iters(master_config: MasterConfig) -> None: """Set train_iters from max_num_steps after its dataloader clamp.""" + max_num_steps = algo_config(master_config).max_num_steps + + # policy policy_config = master_config.policy - if not policy_config.get("megatron_cfg", {}).get("enabled", False): + if policy_config.get("megatron_cfg", {}).get("enabled", False): + policy_config["megatron_cfg"]["train_iters"] = max_num_steps + + # value + if not is_ppo_run(master_config): return - algo_cfg = algo_config(master_config) - policy_config["megatron_cfg"]["train_iters"] = algo_cfg.max_num_steps + value_config = master_config.value + if value_config.get("megatron_cfg", {}).get("enabled", False): + value_config["megatron_cfg"]["train_iters"] = max_num_steps # type: ignore def _maybe_attach_fleet_health( @@ -490,6 +539,20 @@ def _maybe_start_generation_router(generation: Any, master_config: MasterConfig) return router +def _build_advantage_estimator(master_config: MasterConfig) -> Any: + """Build the advantage estimator from whichever algorithm's factory applies.""" + if is_ppo_run(master_config): + # TODO(#2625): raw_reward passes this factory but yields no returns, so + # the critic train would then fetch a column nobody wrote. + from nemo_rl.algorithms.ppo import _create_advantage_estimator + + return _create_advantage_estimator(cast(PPOMasterConfig, master_config)) + else: + from nemo_rl.algorithms.grpo import _create_advantage_estimator + + return _create_advantage_estimator(cast(GRPOMasterConfig, master_config)) + + def _build_retry_policy(master_config: MasterConfig) -> RolloutRetryPolicy: """Translate ``async_rl.rollout_failure`` into the rollout layer's policy object.""" failure_config = master_config.async_rl.rollout_failure @@ -570,6 +633,10 @@ def setup_single_controller( ) save_state = _get_grpo_save_state(loaded_state) weights_path, optimizer_path = checkpointer.get_resume_paths(last_checkpoint_path) + value_weights_path, value_optimizer_path = checkpointer.get_resume_paths( + last_checkpoint_path, + model_component="value", + ) # ========================== # Setup Dataset & Environments @@ -630,10 +697,47 @@ def setup_single_controller( # is disabled or NeMo-Gym is not in play -- it is Gym that needs one stable URL. generation_router = None + def _build_trainer_and_value() -> tuple[Any, Optional[TQValue], dict[str, float]]: + """Build the trainer, then the critic when this is a PPO run. + + Serial, and with the trainer offloaded in between, because both worker + groups live on the same training GPUs: leaving the policy resident + while the critic loads is what OOMs a tight fit. The trainer comes back + to GPU before returning so callers see the same state GRPO leaves them. + + Returns: + A tuple of (TQPolicy trainer, TQValue critic or None, per-phase wall + times keyed as "trainer_time" and "value_time"). + """ + time_metrics: dict[str, float] = {} + trainer, time_metrics["trainer_time"] = _build_trainer( + train_cluster, + master_config, + tokenizer, + processor, + weights_path=weights_path, + optimizer_path=optimizer_path, + ) + if not is_ppo_run(master_config): + return trainer, None, time_metrics + + trainer.offload_to_cpu() + value, time_metrics["value_time"] = _build_value( + train_cluster, + master_config, + tokenizer, + weights_path=value_weights_path, + optimizer_path=value_optimizer_path, + ) + # Blocks on the critic's async Ray __init__, then parks it on CPU. + value.finish_training() + trainer.prepare_for_training() + return trainer, value, time_metrics + def _build_generation_then_trainer( defer_generation_model_load: bool, generation=None - ) -> tuple[Any, Any, dict[str, float]]: - """Build generation then trainer serially. + ) -> tuple[Any, Any, Optional[TQValue], dict[str, float]]: + """Build generation then trainer (and critic) serially. Args: defer_generation_model_load: If True, generation is a pre-reserved handle and this call @@ -641,8 +745,9 @@ def _build_generation_then_trainer( generation: Pre-reserved generation handle when defer_generation_model_load=True; None otherwise. Returns: - A tuple of (finalized generation object, TQPolicy trainer, - per-phase wall times keyed as "gen_time" and "trainer_time"). + A tuple of (finalized generation object, TQPolicy trainer, TQValue + critic or None, per-phase wall times keyed as "gen_time", + "trainer_time" and "value_time"). """ time_metrics = {} @@ -656,17 +761,11 @@ def _build_generation_then_trainer( inference_cluster, master_config ) - # trainer - trainer, time_metrics["trainer_time"] = _build_trainer( - train_cluster, - master_config, - tokenizer, - processor, - weights_path=weights_path, - optimizer_path=optimizer_path, - ) + # trainer (+ critic when PPO) + trainer, value, train_side_metrics = _build_trainer_and_value() + time_metrics.update(train_side_metrics) - return generation, trainer, time_metrics + return generation, trainer, value, time_metrics if not use_nemo_gym and master_config.async_rl.generation_router.enabled: # The router exists to hand NeMo-Gym one URL; the native path calls generation @@ -722,15 +821,7 @@ def _build_generation_then_trainer( inference_cluster=inference_cluster, master_config=master_config, ) - build_tasks["trainer"] = partial( - _build_trainer, - train_cluster=train_cluster, - master_config=master_config, - tokenizer=tokenizer, - processor=processor, - weights_path=weights_path, - optimizer_path=optimizer_path, - ) + build_tasks["trainer"] = _build_trainer_and_value # Submit build tasks and get results with ThreadPoolExecutor(max_workers=len(build_tasks)) as executor: @@ -738,15 +829,17 @@ def _build_generation_then_trainer( results = {k: f.result() for k, f in submitted.items()} if colocated: - generation, trainer, time_metrics = results["generation_trainer"] + generation, trainer, value, time_metrics = results["generation_trainer"] gen_load_time = time_metrics["gen_time"] - setup_timing_metrics.policy_init_time_s = time_metrics["trainer_time"] else: generation, gen_load_time = results["generation"] - trainer, trainer_time = results["trainer"] - setup_timing_metrics.policy_init_time_s = trainer_time + trainer, value, time_metrics = results["trainer"] setup_timing_metrics.generation_init_time_s = gen_reserve_time + gen_load_time + setup_timing_metrics.policy_init_time_s = time_metrics["trainer_time"] + if "value_time" in time_metrics: + setup_timing_metrics.value_init_time_s = time_metrics["value_time"] + if use_nemo_gym: env_handles["nemo_gym"], gym_time = results["nemo_gym"] setup_timing_metrics.nemo_gym_init_time_s = gym_time @@ -783,10 +876,13 @@ def _build_generation_then_trainer( # ========================== # Setup Algorithm + Rollout Wiring # ========================== - advantage_estimator = _create_advantage_estimator( - cast(GrpoMasterConfig, master_config) - ) + advantage_estimator = _build_advantage_estimator(master_config) loss_fn: LossFunction = ClippedPGLossFn(master_config.loss_fn) + value_loss_fn: Optional[LossFunction] = ( + MseValueLossFn(master_config.value_loss_fn) # type: ignore + if is_ppo_run(master_config) + else None + ) pad_id = int(getattr(tokenizer, "pad_token_id", 0) or 0) tq_buffer = TQReplayBuffer( @@ -812,7 +908,7 @@ def _build_generation_then_trainer( env_s=master_config.async_rl.rollout_failure.native.env_timeout_s, ), retry_policy=_build_retry_policy(master_config), - effort_config=_get_effort_config(cast(GrpoMasterConfig, master_config)), + effort_config=_get_effort_config(cast(GRPOMasterConfig, master_config)), ) # Print setup timing metrics @@ -840,5 +936,8 @@ def _build_generation_then_trainer( last_checkpoint_path=last_checkpoint_path, fleet_monitor=fleet_monitor, generation_router=generation_router, + # PPO extras + value_handle=value, + value_loss_fn=value_loss_fn, ) return actor_args, setup_timing_metrics diff --git a/tests/unit/algorithms/test_metric_utils.py b/tests/unit/algorithms/test_metric_utils.py index 15359689b9d..fcf460d8c55 100644 --- a/tests/unit/algorithms/test_metric_utils.py +++ b/tests/unit/algorithms/test_metric_utils.py @@ -76,6 +76,17 @@ def test_optional_nemo_gym_and_teacher_lines(self, capsys): assert "NeMo-Gym init: 8.0s" in out assert "Teacher init: 6.0s" in out + def test_value_init_line_only_on_a_ppo_run(self, capsys): + """Without it the summary does not add up: the critic's time is missing + between Policy init and Total setup.""" + metrics = self._common_setup(generation_init_time_s=15.0) + print_setup_timing_summary(metrics) + assert "Value init" not in capsys.readouterr().out + + metrics = self._common_setup(generation_init_time_s=15.0, value_init_time_s=7.0) + print_setup_timing_summary(metrics) + assert "Value init: 7.0s" in capsys.readouterr().out + class TestSetupTimingMetricsToDict: """to_metrics_dict serializes into a dict for Logger.log_metrics.""" diff --git a/tests/unit/single_controller/test_sc_ppo_setup.py b/tests/unit/single_controller/test_sc_ppo_setup.py new file mode 100644 index 00000000000..074e88414b0 --- /dev/null +++ b/tests/unit/single_controller/test_sc_ppo_setup.py @@ -0,0 +1,369 @@ +# 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 the PPO half of the SingleController setup path. + +Covers what selects the PPO path (`ppo:` present), what that selection then +requires of the rest of the config, and the products setup has to hand the +controller: a GAE estimator, an MSE value loss, and a critic built on the +training cluster after the policy has stepped off it. +""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +import pytest + +import nemo_rl.algorithms.single_controller_utils.setup as sc_setup_mod +from nemo_rl.algorithms.advantage_estimator import ( + GAEConfig, + GeneralizedAdvantageEstimator, +) +from nemo_rl.algorithms.grpo import GRPOConfig +from nemo_rl.algorithms.loss import ClippedPGLossConfig +from nemo_rl.algorithms.loss.loss_functions import MseValueLossConfig, MseValueLossFn +from nemo_rl.algorithms.ppo import PPOConfig +from nemo_rl.algorithms.single_controller_utils import ( + AsyncRLConfig, + MasterConfig, + is_ppo_run, + setup_single_controller, +) +from nemo_rl.algorithms.single_controller_utils.config import ( + algo_config, + validate_single_controller_config, +) + +_NUM_PROMPTS_PER_STEP = 4 +_NUM_GENERATIONS_PER_PROMPT = 2 +_GLOBAL_BATCH_SIZE = _NUM_PROMPTS_PER_STEP * _NUM_GENERATIONS_PER_PROMPT + + +def _value_config( + *, + megatron_enabled: bool = False, + train_global_batch_size: int = _GLOBAL_BATCH_SIZE, +) -> dict: + return { + "model_name": "Qwen/Qwen3-0.6B", + "tokenizer": {"name": "Qwen/Qwen3-0.6B"}, + "train_global_batch_size": train_global_batch_size, + "train_micro_batch_size": 1, + "max_total_sequence_length": 32, + "megatron_cfg": {"enabled": megatron_enabled}, + "dtensor_cfg": {"enabled": not megatron_enabled, "_v2": True}, + } + + +_STEP_CONFIG = dict( + seed=42, + max_num_epochs=1, + num_prompts_per_step=_NUM_PROMPTS_PER_STEP, + num_generations_per_prompt=_NUM_GENERATIONS_PER_PROMPT, + max_rollout_turns=1, + val_period=0, + val_at_start=False, + val_at_end=False, + skip_reference_policy_logprobs_calculation=True, +) + + +def _make_master_config( + *, + ppo: PPOConfig | None = None, + value: dict | None = None, + value_loss_fn: MseValueLossConfig | None = None, + min_groups_for_streaming_train: int = _NUM_PROMPTS_PER_STEP, + megatron_enabled: bool = False, + max_num_steps: int = 100, +) -> MasterConfig: + """An SC config; model_construct skips the fields setup never reads. + + ``grpo`` and ``ppo`` are alternatives, so the step config goes in whichever + block is active and the other one stays None. + """ + return MasterConfig.model_construct( + data_plane={"enabled": True, "impl": "transfer_queue"}, + data={ + "use_multiple_dataloader": False, + "shuffle": False, + "num_workers": 0, + "train": [{"env_name": "math"}], + }, + grpo=None + if ppo is not None + else GRPOConfig.model_construct(max_num_steps=max_num_steps, **_STEP_CONFIG), + policy={ + "train_global_batch_size": _GLOBAL_BATCH_SIZE, + "max_total_sequence_length": 32, + "tokenizer": {"use_fastokens": False}, + "megatron_cfg": {"enabled": megatron_enabled}, + "generation": { + "backend": "vllm", + "colocated": {"enabled": True, "resources": {}}, + }, + }, + checkpointing={ + "enabled": False, + "checkpoint_dir": "results/_sc_ppo_setup_test_ckpt", + "metric_name": None, + "higher_is_better": False, + "keep_top_k": None, + "save_period": 10, + "save_optimizer": False, + }, + loss_fn=ClippedPGLossConfig(reference_policy_kl_penalty=0.0), + env={}, + async_rl=AsyncRLConfig( + min_groups_for_streaming_train=min_groups_for_streaming_train, + max_buffered_rollouts=_NUM_PROMPTS_PER_STEP * 2, + ), + ppo=ppo, + value=value, + value_loss_fn=value_loss_fn, + ) + + +def _ppo_master_config(**kwargs) -> MasterConfig: + kwargs.setdefault( + "ppo", + PPOConfig.model_construct( + max_num_steps=kwargs.get("max_num_steps", 100), **_STEP_CONFIG + ), + ) + kwargs.setdefault( + "value", _value_config(megatron_enabled=kwargs.get("megatron_enabled", False)) + ) + kwargs.setdefault("value_loss_fn", MseValueLossConfig()) + return _make_master_config(**kwargs) + + +class TestIsPPORun: + def test_absent_ppo_block_is_grpo(self): + assert is_ppo_run(_make_master_config()) is False + + def test_present_ppo_block_is_ppo(self): + assert is_ppo_run(_ppo_master_config()) is True + + def test_missing_attribute_is_grpo(self): + """model_construct can omit defaulted fields entirely.""" + assert is_ppo_run(MasterConfig.model_construct()) is False + + +class TestAlgoConfigSelection: + """grpo and ppo are alternatives; SC reads its step config off the live one.""" + + def test_returns_the_ppo_block_on_a_ppo_run(self): + mc = _ppo_master_config() + + assert algo_config(mc) is mc.ppo + + def test_returns_the_grpo_block_otherwise(self): + mc = _make_master_config() + + assert algo_config(mc) is mc.grpo + + def test_rejects_a_config_with_neither_block(self): + """Also caught at setup; algo_config only asserts the invariant.""" + mc = _make_master_config() + mc.grpo = None + + with pytest.raises(ValueError, match="At least one algorithm block"): + validate_single_controller_config(mc) + + def test_rejects_a_config_with_both_blocks(self): + """Caught at setup rather than in algo_config, which runs on every access.""" + mc = _ppo_master_config() + mc.grpo = GRPOConfig.model_construct(**_STEP_CONFIG) + + with pytest.raises(ValueError, match="Only one algorithm block"): + validate_single_controller_config(mc) + + +class TestPPOValidation: + def test_accepts_a_well_formed_ppo_config(self): + validate_single_controller_config(_ppo_master_config()) + + @pytest.mark.parametrize("missing", ["value", "value_loss_fn"]) + def test_rejects_ppo_without_its_critic_blocks(self, missing): + mc = _ppo_master_config(**{missing: None}) + + with pytest.raises(ValueError, match=f"needs `{missing}`"): + validate_single_controller_config(mc) + + def test_rejects_value_block_without_ppo_block(self): + """A value model nothing builds is a silently-downgraded PPO run.""" + mc = _make_master_config(value=_value_config()) + + with pytest.raises(ValueError, match="the `ppo` block is absent"): + validate_single_controller_config(mc) + + def test_rejects_multi_chunk_streaming(self): + """The critic has no split API, so it would step once per chunk while + the policy steps once per RL step.""" + mc = _ppo_master_config(min_groups_for_streaming_train=1) + + with pytest.raises( + ValueError, + match=r"min_groups_for_streaming_train \(1\) == " + rf"num_prompts_per_step \({_NUM_PROMPTS_PER_STEP}\)", + ): + validate_single_controller_config(mc) + + def test_rejects_value_global_batch_size_mismatch(self): + mc = _ppo_master_config(value=_value_config(train_global_batch_size=4)) + + with pytest.raises( + ValueError, + match=r"must equal value.train_global_batch_size \(4\)", + ): + validate_single_controller_config(mc) + + def test_rejects_non_gae_estimator(self): + with pytest.raises(ValueError): + PPOConfig(adv_estimator={"name": "grpo"}) + + +class TestAdvantageEstimatorSelection: + def test_ppo_gets_gae_with_the_configured_lambdas(self): + mc = _ppo_master_config( + ppo=PPOConfig.model_construct( + adv_estimator=GAEConfig(gae_lambda=0.9, gae_gamma=0.99), + **_STEP_CONFIG, + ) + ) + + estimator = sc_setup_mod._build_advantage_estimator(mc) + + assert isinstance(estimator, GeneralizedAdvantageEstimator) + assert estimator.gae_lambda == 0.9 + assert estimator.gae_gamma == 0.99 + + def test_grpo_delegates_to_the_group_relative_factory(self): + mc = _make_master_config() + sentinel = MagicMock(name="grpo_estimator") + + with patch( + "nemo_rl.algorithms.grpo._create_advantage_estimator", + return_value=sentinel, + ) as mock_factory: + assert sc_setup_mod._build_advantage_estimator(mc) is sentinel + + mock_factory.assert_called_once_with(mc) + + +class TestMegatronTrainIters: + def test_injects_into_both_policy_and_value(self): + mc = _ppo_master_config(megatron_enabled=True, max_num_steps=7) + + sc_setup_mod._maybe_inject_megatron_train_iters(mc) + + assert mc.policy["megatron_cfg"]["train_iters"] == 7 + assert mc.value["megatron_cfg"]["train_iters"] == 7 + + def test_skips_a_critic_on_a_non_megatron_backend(self): + mc = _ppo_master_config(megatron_enabled=False, max_num_steps=7) + + sc_setup_mod._maybe_inject_megatron_train_iters(mc) + + assert "train_iters" not in mc.value["megatron_cfg"] + + +@pytest.fixture +def patched_ppo_factories(): + """Patch every external factory the PPO setup path calls.""" + fake_dataloader = MagicMock(name="dataloader") + fake_dataloader.__len__ = MagicMock(return_value=4) + fake_policy = MagicMock(name="policy") + fake_value = MagicMock(name="value") + + with ( + patch.object( + sc_setup_mod, + "setup_response_data", + return_value=(list(range(8)), None, {"math": MagicMock()}, {}), + ), + patch.object(sc_setup_mod, "StatefulDataLoader", return_value=fake_dataloader), + patch.object( + sc_setup_mod, + "_build_clusters", + return_value=(MagicMock(name="train"), MagicMock(name="inference")), + ), + patch.object( + sc_setup_mod, "_build_generation", return_value=(MagicMock(), 0.0) + ), + patch.object( + sc_setup_mod, "_build_trainer", return_value=(fake_policy, 1.0) + ) as mock_trainer, + patch.object( + sc_setup_mod, "_build_value", return_value=(fake_value, 2.0) + ) as mock_value, + patch.object(sc_setup_mod, "build_data_plane_client", return_value=MagicMock()), + patch.object( + sc_setup_mod, "create_weight_synchronizer", return_value=MagicMock() + ), + patch.object(sc_setup_mod, "ClippedPGLossFn", return_value=MagicMock()), + patch( + "nemo_rl.algorithms.grpo._create_advantage_estimator", + return_value=MagicMock(), + ), + patch.object(sc_setup_mod, "_generation_max_seq_len", return_value=32), + ): + yield { + "_build_trainer": mock_trainer, + "_build_value": mock_value, + "policy": fake_policy, + "value": fake_value, + } + + +class TestSetupBuildsTheCritic: + def test_actor_args_carry_the_critic_and_its_loss(self, patched_ppo_factories): + mc = _ppo_master_config() + + actor_args, timing = setup_single_controller( + mc, tokenizer=MagicMock(pad_token_id=0) + ) + + assert actor_args.value_handle is patched_ppo_factories["value"] + assert isinstance(actor_args.value_loss_fn, MseValueLossFn) + assert timing.value_init_time_s == 2.0 + + def test_policy_steps_off_the_gpu_while_the_critic_loads( + self, patched_ppo_factories + ): + """Both worker groups sit on the training cluster; leaving the policy + resident through the critic's init is what OOMs a tight fit.""" + policy = patched_ppo_factories["policy"] + value = patched_ppo_factories["value"] + + setup_single_controller( + _ppo_master_config(), tokenizer=MagicMock(pad_token_id=0) + ) + + policy.offload_to_cpu.assert_called_once_with() + value.finish_training.assert_called_once_with() + policy.prepare_for_training.assert_called_once_with() + + def test_grpo_run_builds_no_critic(self, patched_ppo_factories): + actor_args, timing = setup_single_controller( + _make_master_config(), tokenizer=MagicMock(pad_token_id=0) + ) + + patched_ppo_factories["_build_value"].assert_not_called() + assert actor_args.value_handle is None + assert actor_args.value_loss_fn is None + assert timing.value_init_time_s is None + patched_ppo_factories["policy"].offload_to_cpu.assert_not_called() diff --git a/tests/unit/single_controller/test_single_controller.py b/tests/unit/single_controller/test_single_controller.py index 1be65731a01..7df75df59a5 100644 --- a/tests/unit/single_controller/test_single_controller.py +++ b/tests/unit/single_controller/test_single_controller.py @@ -57,6 +57,59 @@ def _checkpointing_config(tmp_path) -> dict: } +def _grpo_master_config(tmp_path) -> MasterConfig: + """A minimal GRPO MasterConfig the real __init__ accepts.""" + return MasterConfig.model_construct( + policy={"train_global_batch_size": 8}, + grpo=GRPOConfig.model_construct( + num_prompts_per_step=2, + num_generations_per_prompt=4, + ), + loss_fn=ClippedPGLossConfig(force_on_policy_ratio=False), + async_rl=AsyncRLConfig( + min_groups_for_streaming_train=1, + max_buffered_rollouts=4, + ), + logger={}, + env={}, + checkpointing=_checkpointing_config(tmp_path), + ) + + +def _actor_args_for_init(**overrides) -> SimpleNamespace: + """Minimal actor args for a controller built through the real __init__.""" + args = dict( + partition_id="rollout_data", + dp_client=None, + gen_handle=None, + trainer_handle=None, + dataloader=None, + weight_synchronizer=FakeWeightSynchronizer(), + advantage_estimator=None, + loss_fn=None, + tq_buffer=None, + rollout_manager=SimpleNamespace(_tq_buffer=None), + env_handles={}, + fleet_monitor=None, + generation_router=None, + train_cluster=None, + inference_cluster=None, + save_state=_initial_grpo_save_state(), + last_checkpoint_path=None, + ) + args.update(overrides) + return SimpleNamespace(**args) + + +def _init_controller(master_config, actor_args): + controller_cls = SingleControllerActor.__ray_metadata__.modified_class + return controller_cls( + master_config=master_config, + actor_args=actor_args, + setup_timing_metrics=SetupTimingMetrics(), + ) + + def test_rejects_multiple_optimizer_steps_per_rl_step(monkeypatch) -> None: monkeypatch.setattr(single_controller, "Logger", lambda _: object()) master_config = MasterConfig.model_construct( @@ -194,36 +247,38 @@ def test_reference_logprobs_required_only_when_kl_enabled( env={}, checkpointing=_checkpointing_config(tmp_path), ) - actor_args = SimpleNamespace( - partition_id="rollout_data", - dp_client=None, - gen_handle=None, - trainer_handle=None, - dataloader=None, - weight_synchronizer=FakeWeightSynchronizer(), - advantage_estimator=None, - loss_fn=None, - tq_buffer=None, - rollout_manager=SimpleNamespace(_tq_buffer=None), - env_handles={}, - fleet_monitor=None, - generation_router=None, - train_cluster=None, - inference_cluster=None, - save_state=_initial_grpo_save_state(), - last_checkpoint_path=None, - ) - controller_cls = SingleControllerActor.__ray_metadata__.modified_class - controller = controller_cls( - master_config=master_config, - actor_args=actor_args, - setup_timing_metrics=SetupTimingMetrics(), - ) + controller = _init_controller(master_config, _actor_args_for_init()) assert controller._reference_logprobs_required is expected_required +def test_init_picks_up_the_critic_handles(monkeypatch, tmp_path) -> None: + """The PPO path hands the critic and its loss in through actor_args.""" + monkeypatch.setattr(single_controller, "Logger", lambda _: MagicMock()) + value, value_loss_fn = MagicMock(name="value"), MagicMock(name="value_loss_fn") + + ctrl = _init_controller( + _grpo_master_config(tmp_path), + _actor_args_for_init(value_handle=value, value_loss_fn=value_loss_fn), + ) + + assert ctrl._value is value + assert ctrl._value_loss_fn is value_loss_fn + + +def test_init_leaves_the_critic_handles_unset_on_a_grpo_run( + monkeypatch, tmp_path +) -> None: + """actor_args defaults them to None, and older args may omit them entirely.""" + monkeypatch.setattr(single_controller, "Logger", lambda _: MagicMock()) + + ctrl = _init_controller(_grpo_master_config(tmp_path), _actor_args_for_init()) + + assert ctrl._value is None + assert ctrl._value_loss_fn is None + + def test_logs_setup_timing_metrics(monkeypatch, tmp_path) -> None: """setup_timing_metrics is forwarded to Logger.log_metrics under timing/setup.""" logger = MagicMock() diff --git a/tests/unit/single_controller/test_single_controller_setup.py b/tests/unit/single_controller/test_single_controller_setup.py index ea5e994e3c6..efec93d0a93 100644 --- a/tests/unit/single_controller/test_single_controller_setup.py +++ b/tests/unit/single_controller/test_single_controller_setup.py @@ -158,9 +158,8 @@ def patched_factories(): "create_weight_synchronizer", return_value=MagicMock(name="weight_sync"), ) as mock_weight_sync, - patch.object( - sc_setup_mod, - "_create_advantage_estimator", + patch( + "nemo_rl.algorithms.grpo._create_advantage_estimator", return_value=MagicMock(name="adv"), ) as mock_adv, patch.object( From f60ec3249ed396cbd530b7cf4bc6408ac3e472c1 Mon Sep 17 00:00:00 2001 From: Yuki Huang Date: Sun, 23 Aug 2026 01:11:00 -0700 Subject: [PATCH 06/31] refactor: name the SingleController test modules after what they test Signed-off-by: Yuki Huang --- .../{test_sc_checkpointing.py => test_checkpointing.py} | 4 ++-- ...{test_run_grpo_single_controller.py => test_entrypoint.py} | 1 + .../{test_sc_ppo_setup.py => test_ppo_setup.py} | 0 .../{test_single_controller_setup.py => test_setup.py} | 0 ...t_single_controller.py => test_single_controller_actor.py} | 0 .../{test_train_pump.py => test_train_pump_e2e.py} | 0 .../{test_sc_utils_helpers.py => test_utils.py} | 0 7 files changed, 3 insertions(+), 2 deletions(-) rename tests/unit/single_controller/{test_sc_checkpointing.py => test_checkpointing.py} (99%) rename tests/unit/single_controller/{test_run_grpo_single_controller.py => test_entrypoint.py} (99%) rename tests/unit/single_controller/{test_sc_ppo_setup.py => test_ppo_setup.py} (100%) rename tests/unit/single_controller/{test_single_controller_setup.py => test_setup.py} (100%) rename tests/unit/single_controller/{test_single_controller.py => test_single_controller_actor.py} (100%) rename tests/unit/single_controller/{test_train_pump.py => test_train_pump_e2e.py} (100%) rename tests/unit/single_controller/{test_sc_utils_helpers.py => test_utils.py} (100%) diff --git a/tests/unit/single_controller/test_sc_checkpointing.py b/tests/unit/single_controller/test_checkpointing.py similarity index 99% rename from tests/unit/single_controller/test_sc_checkpointing.py rename to tests/unit/single_controller/test_checkpointing.py index a6a0ff2409a..260cafefd2d 100644 --- a/tests/unit/single_controller/test_sc_checkpointing.py +++ b/tests/unit/single_controller/test_checkpointing.py @@ -71,7 +71,7 @@ # Reuse the factory patches from the setup tests (same cross-module fixture # import pattern as test_rollout_pump.py). -from tests.unit.single_controller.test_single_controller_setup import ( +from tests.unit.single_controller.test_setup import ( patched_factories, # noqa: F401 ) @@ -861,7 +861,7 @@ def _write_checkpoint( def _setup_master_config(checkpoint_dir: str) -> MasterConfig: """Partially-populated MasterConfig for setup_single_controller tests. - Same shape as test_single_controller_setup._make_master_config, plus the + Same shape as test_setup._make_master_config, plus the checkpointing block setup now reads. """ return MasterConfig.model_construct( diff --git a/tests/unit/single_controller/test_run_grpo_single_controller.py b/tests/unit/single_controller/test_entrypoint.py similarity index 99% rename from tests/unit/single_controller/test_run_grpo_single_controller.py rename to tests/unit/single_controller/test_entrypoint.py index 5b5d0330276..c4436936db8 100644 --- a/tests/unit/single_controller/test_run_grpo_single_controller.py +++ b/tests/unit/single_controller/test_entrypoint.py @@ -50,6 +50,7 @@ def main_context(monkeypatch: pytest.MonkeyPatch) -> SimpleNamespace: env_handles={}, gen_handle=SimpleNamespace(shutdown=MagicMock()), trainer_handle=SimpleNamespace(shutdown=MagicMock()), + value_handle=None, ) ray_get = MagicMock(return_value={}) # The driver now polls ping() around the run. Report the run as ready on the first diff --git a/tests/unit/single_controller/test_sc_ppo_setup.py b/tests/unit/single_controller/test_ppo_setup.py similarity index 100% rename from tests/unit/single_controller/test_sc_ppo_setup.py rename to tests/unit/single_controller/test_ppo_setup.py diff --git a/tests/unit/single_controller/test_single_controller_setup.py b/tests/unit/single_controller/test_setup.py similarity index 100% rename from tests/unit/single_controller/test_single_controller_setup.py rename to tests/unit/single_controller/test_setup.py diff --git a/tests/unit/single_controller/test_single_controller.py b/tests/unit/single_controller/test_single_controller_actor.py similarity index 100% rename from tests/unit/single_controller/test_single_controller.py rename to tests/unit/single_controller/test_single_controller_actor.py diff --git a/tests/unit/single_controller/test_train_pump.py b/tests/unit/single_controller/test_train_pump_e2e.py similarity index 100% rename from tests/unit/single_controller/test_train_pump.py rename to tests/unit/single_controller/test_train_pump_e2e.py diff --git a/tests/unit/single_controller/test_sc_utils_helpers.py b/tests/unit/single_controller/test_utils.py similarity index 100% rename from tests/unit/single_controller/test_sc_utils_helpers.py rename to tests/unit/single_controller/test_utils.py From a638ccc0479fb0e0f2df0a6d83eacaa8e0328f3f Mon Sep 17 00:00:00 2001 From: Yuki Huang Date: Sun, 23 Aug 2026 01:21:08 -0700 Subject: [PATCH 07/31] feat(sc): accept PPO runs in the SingleController entrypoint, make ppo.async_ppo nullable Signed-off-by: Yuki Huang --- examples/run_grpo_single_controller.py | 19 ++++++++++++++----- .../async_utils/trajectory_collector.py | 6 ++---- nemo_rl/algorithms/ppo.py | 6 ++++-- 3 files changed, 20 insertions(+), 11 deletions(-) diff --git a/examples/run_grpo_single_controller.py b/examples/run_grpo_single_controller.py index 8a440b4832d..b54cb898e15 100644 --- a/examples/run_grpo_single_controller.py +++ b/examples/run_grpo_single_controller.py @@ -12,11 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Async GRPO launcher driven by the SingleController actor. +"""Async GRPO / PPO launcher driven by the SingleController actor. Builds the full SC actor args driver-side via setup_single_controller and hands them to SingleControllerActor. Mirrors run_grpo.py for config loading so the same YAML -files apply. data_plane.enabled=true is mandatory. +files apply. data_plane.enabled=true is mandatory. A config carrying a `ppo:` block +additionally brings up the PPO critic and trains it alongside the policy. """ import argparse @@ -33,6 +34,7 @@ from nemo_rl.algorithms.single_controller_utils import ( MasterConfig, WatchdogConfig, + is_ppo_run, setup_single_controller, ) from nemo_rl.algorithms.utils import get_tokenizer @@ -61,7 +63,7 @@ def parse_args() -> tuple[argparse.Namespace, list[str]]: """Parse command line arguments.""" parser = argparse.ArgumentParser( - description="Run async GRPO training via SingleController" + description="Run async GRPO / PPO training via SingleController" ) parser.add_argument( "--config", type=str, default=None, help="Path to YAML config file" @@ -93,9 +95,13 @@ def main() -> None: config = MasterConfig(**config) print("Applied CLI overrides") - if config.grpo.async_grpo is not None: + if is_ppo_run(config): + legacy_async_block, legacy_async = "ppo.async_ppo", config.ppo.async_ppo + else: + legacy_async_block, legacy_async = "grpo.async_grpo", config.grpo.async_grpo + if legacy_async is not None: raise ValueError( - "SC requires `grpo.async_grpo: null`; use `async_rl.*` instead. " + f"SC requires `{legacy_async_block}: null`; use `async_rl.*` instead. " "See docs/guides/single-controller.md#migrating-a-legacy-async-config." ) @@ -164,7 +170,10 @@ def main() -> None: for resource_name, resource in ( ("Generation", actor_args.gen_handle), ("Trainer", actor_args.trainer_handle), + ("Value", actor_args.value_handle), ): + if resource is None: + continue try: resource.shutdown() except Exception as e: diff --git a/nemo_rl/algorithms/async_utils/trajectory_collector.py b/nemo_rl/algorithms/async_utils/trajectory_collector.py index 4668bdb4b5a..65881c89d11 100644 --- a/nemo_rl/algorithms/async_utils/trajectory_collector.py +++ b/nemo_rl/algorithms/async_utils/trajectory_collector.py @@ -136,9 +136,7 @@ def __init__( async_config: AsyncGRPOConfig | AsyncPPOConfig if isinstance(master_config, GRPOMasterConfig): algorithm_config = master_config.grpo - grpo_async_config = algorithm_config.async_grpo - assert grpo_async_config is not None - async_config = grpo_async_config + async_config = algorithm_config.async_grpo # type: ignore self._deduplicate_multimodal_data = ( algorithm_config.deduplicate_multimodal_data ) @@ -146,7 +144,7 @@ def __init__( self._max_generation_failures = async_config.max_generation_failures elif isinstance(master_config, PPOMasterConfig): algorithm_config = master_config.ppo - async_config = algorithm_config.async_ppo + async_config = algorithm_config.async_ppo # type: ignore self._deduplicate_multimodal_data = False self._debug_payload_metrics = False self._max_generation_failures = 0 diff --git a/nemo_rl/algorithms/ppo.py b/nemo_rl/algorithms/ppo.py index 22148c153b8..bd404a16edf 100644 --- a/nemo_rl/algorithms/ppo.py +++ b/nemo_rl/algorithms/ppo.py @@ -188,12 +188,14 @@ class PPOConfig(BaseModel, extra="allow"): # None logs metrics without masking; values above the threshold are excluded. seq_logprob_error_threshold: float | None = None # Asynchronous PPO uses a replay buffer with non-colocated generation. - async_ppo: AsyncPPOConfig = Field(default_factory=AsyncPPOConfig) + # Legacy async config block; SC reads its async knobs from `async_rl` instead. + async_ppo: AsyncPPOConfig | None = Field(default_factory=AsyncPPOConfig) @model_validator(mode="after") def validate_async_warmup_settings(self) -> "PPOConfig": if ( - self.async_ppo.enabled + self.async_ppo is not None + and self.async_ppo.enabled and self.policy_training_start_step == 0 and self.async_ppo.warmup_generation_lead_steps is not None ): From 55c73ace436b8fe860d2df68634c3fcd9339d0a2 Mon Sep 17 00:00:00 2001 From: Yuki Huang Date: Sun, 23 Aug 2026 04:20:24 -0700 Subject: [PATCH 08/31] feat(sc): size the train cluster for the critic, rename the GAE logprob kwargs to match GRPO's Signed-off-by: Yuki Huang --- nemo_rl/algorithms/advantage_estimator.py | 22 +++--- nemo_rl/algorithms/ppo.py | 8 +- .../single_controller_utils/setup.py | 13 +++- tests/unit/algorithms/test_ppo.py | 12 +-- .../unit/single_controller/test_ppo_setup.py | 76 +++++++++++++++++++ 5 files changed, 107 insertions(+), 24 deletions(-) diff --git a/nemo_rl/algorithms/advantage_estimator.py b/nemo_rl/algorithms/advantage_estimator.py index cd39e8981cb..46e855eee85 100644 --- a/nemo_rl/algorithms/advantage_estimator.py +++ b/nemo_rl/algorithms/advantage_estimator.py @@ -377,8 +377,8 @@ def _build_token_level_rewards( self, rewards: torch.Tensor, mask: torch.Tensor, - logprobs: torch.Tensor | None = None, - reference_logprobs: torch.Tensor | None = None, + logprobs_policy: torch.Tensor | None = None, + logprobs_reference: torch.Tensor | None = None, ) -> torch.Tensor: """Build per-token reward tensor with optional KL penalty. @@ -388,8 +388,8 @@ def _build_token_level_rewards( Args: rewards: Scalar reward per sample, shape [batch_size]. mask: Response token mask, shape [batch_size, seq_len]. - logprobs: Current policy log probs, shape [batch_size, seq_len]. - reference_logprobs: Reference policy log probs, shape [batch_size, seq_len]. + logprobs_policy: Current policy log probs, shape [batch_size, seq_len]. + logprobs_reference: Reference policy log probs, shape [batch_size, seq_len]. Returns: token_level_rewards: shape [batch_size, seq_len]. @@ -403,10 +403,10 @@ def _build_token_level_rewards( if ( self.use_kl_in_reward and self.kl_coef > 0 - and logprobs is not None - and reference_logprobs is not None + and logprobs_policy is not None + and logprobs_reference is not None ): - kl = calculate_kl(logprobs, reference_logprobs, self.kl_type) + kl = calculate_kl(logprobs_policy, logprobs_reference, self.kl_type) token_level_rewards = token_level_rewards - self.kl_coef * kl # Place terminal reward at the last response token (last mask=1 @@ -453,8 +453,8 @@ def compute_advantage( rewards, mask, values, - reference_logprobs=None, - logprobs=None, + logprobs_policy=None, + logprobs_reference=None, **kwargs, ): """Compute GAE advantages with temporal bootstrapping. @@ -472,8 +472,8 @@ def compute_advantage( token_level_rewards = self._build_token_level_rewards( rewards, mask, - logprobs, - reference_logprobs, + logprobs_policy, + logprobs_reference, ) lam_value = self._resolve_lambda_value() diff --git a/nemo_rl/algorithms/ppo.py b/nemo_rl/algorithms/ppo.py index bd404a16edf..9741eb3d62c 100644 --- a/nemo_rl/algorithms/ppo.py +++ b/nemo_rl/algorithms/ppo.py @@ -1561,8 +1561,8 @@ def ppo_train( prompt_ids=prompt_ids_for_adv, rewards=train_data["rewards"], mask=advantage_mask, - reference_logprobs=train_data.get("reference_policy_logprobs"), - logprobs=train_data["prev_logprobs"], + logprobs_policy=train_data["prev_logprobs"], + logprobs_reference=train_data.get("reference_policy_logprobs"), ) if "values" in train_data: adv_kwargs["values"] = train_data["values"] @@ -2551,8 +2551,8 @@ def _raise_if_collector_stopped(waiting_for: str) -> None: prompt_ids=prompt_ids_for_adv, rewards=train_data["rewards"], mask=advantage_mask, - reference_logprobs=train_data.get("reference_policy_logprobs"), - logprobs=train_data["prev_logprobs"], + logprobs_policy=train_data["prev_logprobs"], + logprobs_reference=train_data.get("reference_policy_logprobs"), ) if "values" in train_data: adv_kwargs["values"] = train_data["values"] diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py index 5dd6ba5d91c..e1950ac7dd0 100644 --- a/nemo_rl/algorithms/single_controller_utils/setup.py +++ b/nemo_rl/algorithms/single_controller_utils/setup.py @@ -147,15 +147,22 @@ def _build_clusters( gpus_per_node = cluster_config["gpus_per_node"] port_range_low = cluster_config.get("master_port_range_low") port_range_high = cluster_config.get("master_port_range_high") + # Worker groups sharing the training GPUs: the policy, plus the critic on + # the PPO path. + train_worker_groups = 2 if is_ppo_run(master_config) else 1 if colocated: - # Policy + generation share GPUs — one cluster. + # Policy (+ critic) + generation share GPUs — one cluster. cluster = RayVirtualCluster( name="sc_policy_cluster", bundle_ct_per_node_list=[gpus_per_node] * num_nodes, use_gpus=True, num_gpus_per_node=gpus_per_node, - max_colocated_worker_groups=1 if backend == "megatron" else 2, + max_colocated_worker_groups=( + train_worker_groups + if backend == "megatron" + else train_worker_groups + 1 + ), port_range_low=port_range_low, port_range_high=port_range_high, ) @@ -192,7 +199,7 @@ def _build_clusters( bundle_ct_per_node_list=[train_gpus_per_node] * train_nodes, use_gpus=True, num_gpus_per_node=train_gpus_per_node, - max_colocated_worker_groups=1, + max_colocated_worker_groups=train_worker_groups, port_range_low=port_range_low, port_range_high=port_range_high, ) diff --git a/tests/unit/algorithms/test_ppo.py b/tests/unit/algorithms/test_ppo.py index 9defc2e30da..fcbca49d9b8 100644 --- a/tests/unit/algorithms/test_ppo.py +++ b/tests/unit/algorithms/test_ppo.py @@ -197,8 +197,8 @@ def test_gae_kl_penalty_in_rewards(): mask=mask, lengths=lengths, values=values, - logprobs=logprobs, - reference_logprobs=reference_logprobs_same, + logprobs_policy=logprobs, + logprobs_reference=reference_logprobs_same, ) adv_no_kl, _ = estimator.compute_advantage( prompt_ids=torch.tensor([[0]]), @@ -218,8 +218,8 @@ def test_gae_kl_penalty_in_rewards(): mask=mask, lengths=lengths, values=values, - logprobs=logprobs, - reference_logprobs=reference_logprobs_divergent, + logprobs_policy=logprobs, + logprobs_reference=reference_logprobs_divergent, ) assert not torch.allclose(adv_kl_positive, adv_no_kl), ( "Non-zero KL should change advantages relative to the no-KL baseline" @@ -239,8 +239,8 @@ def test_gae_kl_penalty_in_rewards(): mask=mask, lengths=lengths, values=values, - logprobs=logprobs, - reference_logprobs=reference_logprobs_divergent, + logprobs_policy=logprobs, + logprobs_reference=reference_logprobs_divergent, ) torch.testing.assert_close(adv_gate_off, adv_no_kl) diff --git a/tests/unit/single_controller/test_ppo_setup.py b/tests/unit/single_controller/test_ppo_setup.py index 074e88414b0..ec932f23f53 100644 --- a/tests/unit/single_controller/test_ppo_setup.py +++ b/tests/unit/single_controller/test_ppo_setup.py @@ -367,3 +367,79 @@ def test_grpo_run_builds_no_critic(self, patched_ppo_factories): assert actor_args.value_loss_fn is None assert timing.value_init_time_s is None patched_ppo_factories["policy"].offload_to_cpu.assert_not_called() + + +def _cluster_config(mc: MasterConfig, *, colocated: bool, backend: str) -> MasterConfig: + """Fill in the cluster / generation keys _build_clusters reads.""" + mc.cluster = { + "num_nodes": 1, + "gpus_per_node": 8, + "master_port_range_low": None, + "master_port_range_high": None, + } + mc.policy["generation"] = { + "backend": backend, + "colocated": { + "enabled": colocated, + "resources": {"gpus_per_node": None if colocated else 2, "num_nodes": None}, + }, + } + return mc + + +class TestTrainClusterSizesForTheCritic: + """The critic shares the training GPUs, so it needs its own worker-group slot.""" + + @pytest.fixture + def fake_cluster(self): + with patch.object(sc_setup_mod, "RayVirtualCluster") as cls: + cls.side_effect = lambda **kwargs: MagicMock(kwargs=kwargs) + yield cls + + def _groups(self, mc): + train, inference = sc_setup_mod._build_clusters(mc) + return train.kwargs["max_colocated_worker_groups"], inference.kwargs[ + "max_colocated_worker_groups" + ] + + def test_noncolocated_ppo_leaves_a_slot_for_the_critic(self, fake_cluster): + mc = _cluster_config(_ppo_master_config(), colocated=False, backend="vllm") + + train_groups, inference_groups = self._groups(mc) + + assert train_groups == 2 + # The critic never lands on the inference cluster. + assert inference_groups == 1 + + def test_noncolocated_grpo_is_unchanged(self, fake_cluster): + mc = _cluster_config(_make_master_config(), colocated=False, backend="vllm") + + train_groups, inference_groups = self._groups(mc) + + assert train_groups == 1 + assert inference_groups == 1 + + def test_colocated_ppo_adds_the_critic_beside_policy_and_generation( + self, fake_cluster + ): + mc = _cluster_config(_ppo_master_config(), colocated=True, backend="vllm") + + train, inference = sc_setup_mod._build_clusters(mc) + + assert train is inference + assert train.kwargs["max_colocated_worker_groups"] == 3 + + def test_colocated_grpo_is_unchanged(self, fake_cluster): + mc = _cluster_config(_make_master_config(), colocated=True, backend="vllm") + + train, _ = sc_setup_mod._build_clusters(mc) + + assert train.kwargs["max_colocated_worker_groups"] == 2 + + def test_colocated_megatron_generation_needs_no_extra_slot(self, fake_cluster): + """The megatron backend generates from the policy's own workers.""" + mc = _cluster_config(_ppo_master_config(), colocated=True, backend="megatron") + + train, _ = sc_setup_mod._build_clusters(mc) + + assert train.kwargs["max_colocated_worker_groups"] == 2 From 596a4bda5d17993f063766735fafbdcbca865e87 Mon Sep 17 00:00:00 2001 From: Yuki Huang Date: Sun, 23 Aug 2026 04:20:31 -0700 Subject: [PATCH 09/31] feat(sc): run the PPO critic stages in the SingleController train pump Signed-off-by: Yuki Huang --- nemo_rl/algorithms/single_controller.py | 200 +++++++++-- .../single_controller_utils/config.py | 4 + .../single_controller/test_checkpointing.py | 101 +++++- .../test_single_controller_actor.py | 316 ++++++++++++++++++ 4 files changed, 585 insertions(+), 36 deletions(-) diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 713a33d1290..ddf11147f74 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -53,10 +53,12 @@ compute_and_apply_seq_logprob_error_masking, ) from nemo_rl.algorithms.metric_utils import SetupTimingMetrics +from nemo_rl.algorithms.ppo import _compute_critic_metrics from nemo_rl.algorithms.single_controller_utils.config import ( AdvantageConfig, MasterConfig, algo_config, + is_ppo_run, validate_sampler_buffer_capacity, validate_single_controller_config, ) @@ -130,6 +132,8 @@ def __init__( self._master_config = master_config self._algo_cfg = algo_config(master_config) self._async_cfg = master_config.async_rl + self._is_ppo: bool = is_ppo_run(master_config) + self._policy_logprobs_required = not ( master_config.loss_fn.force_on_policy_ratio and self._algo_cfg.seq_logprob_error_threshold is None @@ -850,16 +854,35 @@ def _target_groups_for_step(self, step: int) -> int: async def _train_pump(self) -> None: """Per-prompt-group streaming train loop. - Per step: - 1. sampler.evict drops stale groups from the buffer and clears their TQ rows. - 2. sampler.select returns K prompt groups (or None) and drops them from the - buffer; DP rows survive so the trainer can read them. Already trainable — - buffer wrote training-shaped rows at rollout time. - 3. _advantage_stage(train_meta). - 4. trainer.train_microbatches_from_meta + finish_train_step. - 5. dp_client.clear_samples on consumed sample_ids; release _buffer_capacity - per dropped group, then sync. + Per step, with 1-4 running per streaming chunk and 5-6 once the chunk + loop closes: + 1. Select the rollouts to train on. + a. sampler.evict drops stale groups from the buffer and clears their + TQ rows. + b. sampler.select returns K prompt groups, or None, and removes them from + the buffer. The DP rows survive, already training-shaped because the + buffer wrote them that way at rollout time. + c. One _buffer_capacity permit is released per group that left the buffer. + 2. Prepare the batch. + a. Policy and reference logprobs. + b. Value model forward (PPO only), with the policy parked on CPU + so the critic never shares the training GPUs with it. + c. _advantage_stage. + 3. Train on the chunk. + a. Value model (PPO only): train_from_meta, which is a whole + optimizer step. That is why a PPO step is pinned to a single + chunk -- the value workers have no split train API yet (#2625). + b. Policy model: train_microbatches_from_meta, which only + accumulates gradients. + 4. Clear the batch. dp_client.clear_samples on the consumed sample_ids. + 5. Train the policy model -- finish_train_step all_reduces the + accumulated gradients, rescales, and runs optimizer.step. + 6. Refit the model. Sync the new policy weights to generation. """ + policy_training_start_step = ( + self._algo_cfg.policy_training_start_step if self._is_ppo else 0 + ) + while self._train_steps < self._algo_cfg.max_num_steps: version_during_step = self._trainer_version groups_dispatched = 0 @@ -868,6 +891,10 @@ async def _train_pump(self) -> None: step_open = False chunks_dispatched = 0 calibration_batches: list[BatchedDataDict[Any]] = [] + # One chunk per step on the PPO path, so this is the step's value model update. + value_result: Optional[dict[str, Any]] = None + # Always True off the PPO path: the start step is pinned to 0 there. + is_policy_training_step = self._train_steps >= policy_training_start_step with self._timer.time("total_step_time"): # Re-read on every iteration rather than once: a prompt stamped for this @@ -876,7 +903,7 @@ async def _train_pump(self) -> None: while groups_dispatched < self._target_groups_for_step( version_during_step ): - # Wait for a selectable batch + # ---- 1. Select the rollouts to train on ---- with self._timer.time("exposed_generation"): await asyncio.sleep(0) @@ -945,6 +972,7 @@ async def _train_pump(self) -> None: for _ in range(num_groups): self._buffer_capacity.release() + # ---- 2. Prepare the batch ---- # Compute prev_logprobs / ref_logprobs if ( self._policy_logprobs_required @@ -972,6 +1000,12 @@ async def _train_pump(self) -> None: train_meta, ) + # Value model forward + if self._is_ppo: + with self._timer.time("value_inference"): + await asyncio.to_thread(self._trainer.finish_inference) + train_meta = await self._value_stage(train_meta) + # Compute advantages with self._timer.time("advantage_calculation"): ( @@ -979,24 +1013,34 @@ async def _train_pump(self) -> None: has_valid_training_tokens, ) = await self._advantage_stage(train_meta) + # ---- 3. Train the model -- train_microbatches_from_meta ---- # Filtering can leave a streaming chunk with no training tokens. # Consume that chunk without F/B, then continue the same optimizer - # step with the next chunk. Always restore training mode because - # log-prob inference may have switched the model to inference mode. - with self._timer.time("training_prep"): - await asyncio.to_thread(self._trainer.prepare_for_training) - if has_valid_training_tokens: - with self._timer.time("policy_training"): - if not step_open: + # step with the next chunk. + + # Value model first, then policy, as in the legacy PPO epoch loop. + if self._is_ppo and has_valid_training_tokens: + with self._timer.time("value_training"): + value_result = await self._value_train(train_meta) + + if is_policy_training_step: + # Always restore training mode because log-prob inference may have + # switched the model to inference mode. + with self._timer.time("training_prep"): + await asyncio.to_thread(self._trainer.prepare_for_training) + + if has_valid_training_tokens: + with self._timer.time("policy_training"): + if not step_open: + await asyncio.to_thread( + self._trainer.begin_train_step, + self._loss_fn, + ) + step_open = True await asyncio.to_thread( - self._trainer.begin_train_step, - self._loss_fn, + self._trainer.train_microbatches_from_meta, + train_meta, ) - step_open = True - await asyncio.to_thread( - self._trainer.train_microbatches_from_meta, - train_meta, - ) if train_meta.sequence_lengths: self._step_log_dict["sequence_lengths"].extend( @@ -1017,6 +1061,7 @@ async def _train_pump(self) -> None: ) ) + # ---- 4. Clear the batch ---- # Refresh min_sample_version curr_min_sample_version = min( t["weight_version"] @@ -1058,23 +1103,34 @@ async def _train_pump(self) -> None: self._algo_cfg.num_prompts_per_step, ) + # ---- 5. Train the policy model -- finish_train_step ---- log.info( "train_pump: step %d closing on %d chunk(s), %d group(s)", version_during_step, chunks_dispatched, groups_dispatched, ) - if not step_open: - raise RuntimeError( - "SingleController has no valid response tokens after " - "filtering. Check grpo.seq_logprob_error_threshold to " - "avoid an optimizer step with an empty batch." - ) - with self._timer.time("policy_training"): - result = await asyncio.to_thread(self._trainer.finish_train_step) + policy_result: Optional[dict[str, Any]] = None + if is_policy_training_step: + if not step_open: + raise RuntimeError( + "SingleController has no valid response tokens after " + "filtering. Check grpo.seq_logprob_error_threshold to " + "avoid an optimizer step with an empty batch." + ) - step_metrics = aggregate_step_metrics(result) + with self._timer.time("policy_training"): + policy_result = await asyncio.to_thread( + self._trainer.finish_train_step + ) + + # Aggregate step metrics + step_metrics = {} + if policy_result is not None: + step_metrics.update(aggregate_step_metrics(policy_result)) + if value_result is not None: + step_metrics.update(_compute_critic_metrics(value_result)) step_metrics.update( reduce_advantage_pump_metrics(**self._step_log_dict) ) @@ -1109,6 +1165,8 @@ async def _train_pump(self) -> None: for step, promoted in self._batch_promotions.items() if step > version_during_step } + + # ---- 6. Refit the model ---- with self._timer.time("weight_sync"): calibration_data = ( BatchedDataDict.from_batches(calibration_batches) @@ -1567,6 +1625,28 @@ async def _save_checkpoint(self, step_metrics: dict[str, Any]) -> None: vars(save_state), self._master_config, ) + if self._is_ppo: + # The critic shares the training GPUs, so the two weight saves are + # serialized. The critic goes first because offloading the policy runs + # finalize_async_save, which would block on the policy's own write if + # that had already been staged. + await asyncio.to_thread(self._trainer.offload_to_cpu) + # The value model writes synchronously, so unlike the policy below it is + # fully on disk when this returns and needs no finalize wait. + await asyncio.to_thread(self._value.prepare_for_training) + await asyncio.to_thread( + self._value.save_checkpoint, + weights_path=os.path.join(checkpoint_path, "value", "weights"), + optimizer_path=os.path.join(checkpoint_path, "value", "optimizer") + if self._checkpointer.save_optimizer + else None, + tokenizer_path=os.path.join(checkpoint_path, "value", "tokenizer"), + checkpointing_cfg=self._master_config.checkpointing, + ) + await asyncio.to_thread(self._value.finish_training) + # Also covers a warmup step, which never ran prepare_for_training in + # the pump and would otherwise save from CPU-resident params. + await asyncio.to_thread(self._trainer.prepare_for_training) # With async_save this returns after D2H staging; disk writes finish # in the background. await asyncio.to_thread( @@ -1677,6 +1757,38 @@ async def _sync_weights( self._rollout_permitted.set() return aborted_stale_inflight_groups + async def _value_stage(self, meta: KVBatchMeta) -> KVBatchMeta: + """Run the PPO value model's forward pass over the selected chunk. + + Tensors never touch SC: workers fetch the sequence columns from + DataPlane and commit the per-token prediction back under ``values``, + which the advantage stage then reads alongside the rewards. The value model + is loaded and offloaded around the call, so it holds the training GPUs + only for the duration of the forward. + + Returns: + The batch metadata with the ``values`` column recorded on it. + """ + await asyncio.to_thread(self._value.prepare_for_inference) + await asyncio.to_thread(self._value.get_values_from_meta, meta) + await asyncio.to_thread(self._value.finish_inference) + return meta.with_fields([self._advantage_cfg.values_field]) + + async def _value_train(self, meta: KVBatchMeta) -> dict[str, Any]: + """Run one value model optimizer step against this chunk's GAE returns. + + Returns: + The aggregated value model train result, shaped like Value.train's. + """ + await asyncio.to_thread(self._value.prepare_for_training) + result = await asyncio.to_thread( + self._value.train_from_meta, + meta, + self._value_loss_fn, # pyrefly: ignore + ) + await asyncio.to_thread(self._value.finish_training) + return result + async def _advantage_stage(self, meta: KVBatchMeta) -> tuple[KVBatchMeta, bool]: """Fetch advantage inputs, compute advantages, and write them back. @@ -1773,20 +1885,32 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> tuple[KVBatchMeta, bool]: data, adv_cfg.reference_logprobs_field, ) + if self._is_ppo: + kwargs["values"] = tensor_field(data, adv_cfg.values_field) # Training predicts token t from position t - 1, so token_mask[:, 1:] # is the exact mask used when global_valid_toks and the loss are built. has_valid_training_tokens = bool(mask[:, 1:].bool().any().item()) + # Value-model estimators (GAE) hand back the regression target alongside + # the advantages; the group-relative ones return a bare tensor. + returns: Optional[torch.Tensor] = None if has_valid_training_tokens: - advantages = self._advantage_estimator.compute_advantage( + result = self._advantage_estimator.compute_advantage( prompt_ids=prompt_ids, rewards=rewards, mask=mask, repeated_batch=repeated_batch, **kwargs, ) + if self._is_ppo: + advantages, returns = result + else: + advantages = result else: advantages = torch.zeros_like(mask) + if self._is_ppo: + returns = torch.zeros_like(mask) + response_advantages = torch.masked_select(advantages, mask.bool()) self._step_log_dict["rewards"].append(rewards.detach().cpu()) self._step_log_dict["masked_advantages"].append( @@ -1796,6 +1920,10 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> tuple[KVBatchMeta, bool]: fields_to_put = {adv_cfg.output_field: advantages} if seq_logprob_error_threshold is not None: fields_to_put[adv_cfg.sample_mask_field] = sample_mask + new_fields = [adv_cfg.output_field] + if returns is not None: + fields_to_put[adv_cfg.returns_field] = returns + new_fields.append(adv_cfg.returns_field) await self._call_dp( "put_samples", @@ -1804,7 +1932,7 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> tuple[KVBatchMeta, bool]: fields=fields_for_put(meta, fields_to_put), ) return ( - meta.with_fields([adv_cfg.output_field]), + meta.with_fields(new_fields), has_valid_training_tokens, ) @@ -1825,4 +1953,6 @@ def _advantage_input_fields(self) -> list[str]: fields.append(adv_cfg.generation_logprobs_field) if self._reference_logprobs_required: fields.append(adv_cfg.reference_logprobs_field) + if self._is_ppo: + fields.append(adv_cfg.values_field) return list(dict.fromkeys(fields)) diff --git a/nemo_rl/algorithms/single_controller_utils/config.py b/nemo_rl/algorithms/single_controller_utils/config.py index 971961f5b00..88d6d7898a6 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/nemo_rl/algorithms/single_controller_utils/config.py @@ -926,3 +926,7 @@ class AdvantageConfig: policy_logprobs_field: str = "prev_logprobs" generation_logprobs_field: str = "generation_logprobs" reference_logprobs_field: str = "reference_policy_logprobs" + # PPO only: the critic's pre-update prediction (input) and GAE's + # regression target for it (output). + values_field: str = "values" + returns_field: str = "returns" diff --git a/tests/unit/single_controller/test_checkpointing.py b/tests/unit/single_controller/test_checkpointing.py index 260cafefd2d..1c026b5f6ad 100644 --- a/tests/unit/single_controller/test_checkpointing.py +++ b/tests/unit/single_controller/test_checkpointing.py @@ -41,8 +41,9 @@ import threading from collections.abc import Callable from pathlib import Path +from types import SimpleNamespace from typing import Any, Optional, Union -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest import torch @@ -776,6 +777,104 @@ def test_failed_background_finalization_raises_at_shutdown(self, tmp_path): actor._checkpointer.shutdown() +# ── PPO save ordering ──────────────────────────────────────────────────────── + + +class _OrderRecordingPolicy: + """Policy stand-in logging the residency calls _save_checkpoint drives.""" + + def __init__(self, calls: list[str]) -> None: + self.calls = calls + + def offload_to_cpu(self) -> None: + self.calls.append("policy.offload_to_cpu") + + def prepare_for_training(self) -> None: + self.calls.append("policy.prepare_for_training") + + def save_checkpoint(self, **kwargs: Any) -> None: + del kwargs + self.calls.append("policy.save_checkpoint") + + def finalize_async_save(self) -> None: + pass + + +class _OrderRecordingCritic: + """Critic stand-in sharing the policy's call log.""" + + def __init__(self, calls: list[str]) -> None: + self.calls = calls + + def prepare_for_training(self) -> None: + self.calls.append("critic.prepare_for_training") + + def save_checkpoint(self, **kwargs: Any) -> None: + del kwargs + self.calls.append("critic.save_checkpoint") + + def finish_training(self) -> None: + self.calls.append("critic.finish_training") + + +def _ppo_save_actor(tmp_path: Path, calls: list[str]): + """Bare actor carrying only what _save_checkpoint reads.""" + actor = object.__new__(_ACTOR_CLS) + checkpoint_path = tmp_path / "tmp_step_1" + checkpoint_path.mkdir(parents=True, exist_ok=True) + + actor._save_state = SimpleNamespace() + actor._train_steps = 1 + actor._current_epoch = 0 + actor._consumed_samples = 0 + actor._total_valid_tokens = 0 + actor._replacement_reserve = [] + actor._async_cfg = SimpleNamespace( + sampler=SimpleNamespace(name="in_order"), + max_buffered_rollouts=4, + ) + actor._master_config = SimpleNamespace(checkpointing={"metric_name": None}) + actor._dataloader = SimpleNamespace(state_dict=lambda: {}) + actor._buffer = SimpleNamespace(state_dict=AsyncMock(return_value={})) + actor._checkpointer = MagicMock() + actor._checkpointer.save_optimizer = True + actor._checkpointer.init_tmp_checkpoint.return_value = str(checkpoint_path) + actor._is_ppo = True + actor._trainer = _OrderRecordingPolicy(calls) + actor._value = _OrderRecordingCritic(calls) + return actor + + +class TestPPOSaveOrder: + def test_the_policy_is_offloaded_across_the_critic_save( + self, tmp_path, monkeypatch + ): + """The critic shares the training GPUs, so the two saves serialize. + + The critic goes first because offloading the policy runs + finalize_async_save, which would block on the policy's own write if that + had already been staged. The onload before the policy save is also what a + critic-warmup step needs, having skipped prepare_for_training in the pump. + """ + monkeypatch.setattr( + "nemo_rl.algorithms.single_controller._write_latest_checkpoint_status", + lambda *args, **kwargs: None, + ) + calls: list[str] = [] + actor = _ppo_save_actor(tmp_path, calls) + + asyncio.run(actor._save_checkpoint({})) + + assert calls == [ + "policy.offload_to_cpu", + "critic.prepare_for_training", + "critic.save_checkpoint", + "critic.finish_training", + "policy.prepare_for_training", + "policy.save_checkpoint", + ] + + # ── metric_name behavior ───────────────────────────────────────────────────── diff --git a/tests/unit/single_controller/test_single_controller_actor.py b/tests/unit/single_controller/test_single_controller_actor.py index 7df75df59a5..a6094b81fe0 100644 --- a/tests/unit/single_controller/test_single_controller_actor.py +++ b/tests/unit/single_controller/test_single_controller_actor.py @@ -28,6 +28,7 @@ from nemo_rl.algorithms.grpo import GRPOConfig, _initial_grpo_save_state from nemo_rl.algorithms.loss import ClippedPGLossConfig from nemo_rl.algorithms.metric_utils import SetupTimingMetrics +from nemo_rl.algorithms.ppo import PPOConfig from nemo_rl.algorithms.single_controller import SingleControllerActor from nemo_rl.algorithms.single_controller_utils.config import ( AdvantageConfig, @@ -465,6 +466,7 @@ def test_advantage_stage_applies_seq_logprob_error_mask_before_streaming_train( ctrl._advantage_estimator = estimator ctrl._policy_logprobs_required = True ctrl._reference_logprobs_required = False + ctrl._is_ppo = False ctrl._master_config = SimpleNamespace( grpo=SimpleNamespace(seq_logprob_error_threshold=2.0) ) @@ -532,6 +534,7 @@ def test_advantage_stage_reports_seq_logprob_metrics_without_masking() -> None: ctrl._advantage_estimator = estimator ctrl._policy_logprobs_required = True ctrl._reference_logprobs_required = False + ctrl._is_ppo = False ctrl._master_config = SimpleNamespace( grpo=SimpleNamespace(seq_logprob_error_threshold=None) ) @@ -593,6 +596,7 @@ def test_advantage_stage_skips_estimator_when_seq_mask_removes_whole_chunk( ctrl._advantage_estimator = estimator ctrl._policy_logprobs_required = True ctrl._reference_logprobs_required = False + ctrl._is_ppo = False ctrl._master_config = SimpleNamespace( grpo=SimpleNamespace(seq_logprob_error_threshold=2.0) ) @@ -647,6 +651,7 @@ def test_advantage_stage_skips_preexisting_empty_mask_without_seq_threshold() -> ctrl._advantage_estimator = estimator ctrl._policy_logprobs_required = False ctrl._reference_logprobs_required = False + ctrl._is_ppo = False ctrl._master_config = SimpleNamespace( grpo=SimpleNamespace(seq_logprob_error_threshold=None) ) @@ -753,6 +758,9 @@ class _NoOpTrainer: def prepare_for_lp_inference(self, keep_train_buffers: bool = False) -> None: del keep_train_buffers + def finish_inference(self) -> None: + pass + def prepare_for_training(self) -> None: pass @@ -779,6 +787,27 @@ def get_logprobs_from_meta(self, meta: KVBatchMeta) -> None: del meta +class _OrderRecordingTrainer(_NoOpTrainer): + """Records the policy lifecycle into a log shared with the critic double.""" + + def __init__(self, calls: list[str]) -> None: + self.calls = calls + + def prepare_for_lp_inference(self, keep_train_buffers: bool = False) -> None: + del keep_train_buffers + self.calls.append("policy.prepare_for_lp_inference") + + def get_logprobs_from_meta(self, meta: KVBatchMeta) -> None: + del meta + self.calls.append("policy.get_logprobs_from_meta") + + def finish_inference(self) -> None: + self.calls.append("policy.finish_inference") + + def prepare_for_training(self) -> None: + self.calls.append("policy.prepare_for_training") + + class _NoOpDataPlane: def clear_samples(self, **kwargs) -> None: del kwargs @@ -816,6 +845,9 @@ def _train_pump_controller(*, sampler) -> object: ctrl._rollout_exhausted = asyncio.Event() ctrl._rollout_exhausted.set() ctrl._trainer = _NoOpTrainer() + ctrl._is_ppo = False + ctrl._value = None + ctrl._value_loss_fn = None ctrl._gen = SimpleNamespace(requires_kv_scale_sync=False) ctrl._loss_fn = None ctrl._dp_client = _NoOpDataPlane() @@ -1107,3 +1139,287 @@ def test_train_pump_keeps_train_buffers_once_the_step_is_open(monkeypatch) -> No assert ctrl._train_steps == 1 assert trainer.keep_train_buffers_calls == [False, True] + + +# ── PPO ──────────────────────────────────────────────────────────────────── + + +class _NoOpValue: + """Records the critic lifecycle the pump drives around each stage. + + Pass a shared list and a prefix to interleave this log with the policy's. + """ + + def __init__(self, calls: list[str] | None = None, prefix: str = "") -> None: + self.calls: list[str] = [] if calls is None else calls + self._prefix = prefix + + def _record(self, name: str) -> None: + self.calls.append(f"{self._prefix}{name}") + + def prepare_for_inference(self) -> None: + self._record("prepare_for_inference") + + def get_values_from_meta(self, meta: KVBatchMeta) -> None: + del meta + self._record("get_values_from_meta") + + def finish_inference(self) -> None: + self._record("finish_inference") + + def prepare_for_training(self) -> None: + self._record("prepare_for_training") + + def train_from_meta(self, meta: KVBatchMeta, loss_fn) -> dict: + del meta, loss_fn + self._record("train_from_meta") + return { + "loss": torch.tensor([0.25]), + "grad_norm": torch.tensor([1.5]), + "all_mb_metrics": {"vf_clipfrac": [0.0], "values_min": [-1.0]}, + } + + def finish_training(self) -> None: + self._record("finish_training") + + +def _ppo_train_pump_controller( + *, + sampler, + policy_training_start_step: int = 0, + value: _NoOpValue | None = None, +) -> tuple[object, _NoOpValue]: + ctrl = _train_pump_controller(sampler=sampler) + value = _NoOpValue() if value is None else value + ctrl._is_ppo = True + ctrl._value = value + ctrl._value_loss_fn = MagicMock(name="value_loss_fn") + ctrl._master_config.grpo = None + ctrl._master_config.ppo = PPOConfig.model_construct( + num_prompts_per_step=1, + max_num_steps=1, + policy_training_start_step=policy_training_start_step, + seq_logprob_error_threshold=None, + ) + ctrl._algo_cfg = ctrl._master_config.ppo + ctrl._sync_weights = AsyncMock(return_value=0) + ctrl._logger = MagicMock() + return ctrl, value + + +def _single_group_meta() -> KVBatchMeta: + return KVBatchMeta( + partition_id="rollout_data", + task_name="train", + sample_ids=["sample-0"], + fields=[], + sequence_lengths=[1], + tags=[{"weight_version": 0}], + ) + + +def test_train_pump_loads_and_offloads_the_critic_around_each_stage( + monkeypatch, +) -> None: + """The critic holds the training GPUs only inside its forward and its train.""" + meta = _single_group_meta() + ctrl, value = _ppo_train_pump_controller(sampler=_OneThenEmptySampler(meta)) + ctrl._policy_logprobs_required = True + trainer = _LpRecordingTrainer() + ctrl._trainer = trainer + ctrl._advantage_stage = AsyncMock(return_value=(meta, True)) + monkeypatch.setattr(single_controller.ray, "cluster_resources", lambda: {}) + + asyncio.run(asyncio.wait_for(ctrl._train_pump(), timeout=1.0)) + + assert value.calls == [ + "prepare_for_inference", + "get_values_from_meta", + "finish_inference", + # Critic before actor, mirroring the legacy PPO epoch loop. + "prepare_for_training", + "train_from_meta", + "finish_training", + ] + assert ctrl._train_steps == 1 + + +def test_train_pump_parks_the_policy_on_cpu_across_the_critic_stages( + monkeypatch, +) -> None: + """The critic shares the training GPUs, so the two models never overlap. + + The critic forward runs after the log-prob pass rather than before it, as + ppo.py does, so the policy reaches finish_inference with its grad buffers + already freed. + """ + meta = _single_group_meta() + calls: list[str] = [] + ctrl, _ = _ppo_train_pump_controller( + sampler=_OneThenEmptySampler(meta), + value=_NoOpValue(calls=calls, prefix="critic."), + ) + ctrl._policy_logprobs_required = True + ctrl._trainer = _OrderRecordingTrainer(calls) + ctrl._advantage_stage = AsyncMock(return_value=(meta, True)) + monkeypatch.setattr(single_controller.ray, "cluster_resources", lambda: {}) + + asyncio.run(asyncio.wait_for(ctrl._train_pump(), timeout=1.0)) + + assert calls == [ + "policy.prepare_for_lp_inference", + "policy.get_logprobs_from_meta", + "policy.finish_inference", + "critic.prepare_for_inference", + "critic.get_values_from_meta", + "critic.finish_inference", + "critic.prepare_for_training", + "critic.train_from_meta", + "critic.finish_training", + "policy.prepare_for_training", + ] + + +def test_train_pump_logs_critic_metrics(monkeypatch) -> None: + meta = _single_group_meta() + ctrl, _ = _ppo_train_pump_controller(sampler=_OneThenEmptySampler(meta)) + ctrl._advantage_stage = AsyncMock(return_value=(meta, True)) + monkeypatch.setattr(single_controller.ray, "cluster_resources", lambda: {}) + + asyncio.run(asyncio.wait_for(ctrl._train_pump(), timeout=1.0)) + + train_metrics = ctrl._logger.log_metrics.call_args_list[0].args[0] + assert train_metrics["critic/loss"].item() == pytest.approx(0.25) + assert train_metrics["critic/grad_norm"].item() == pytest.approx(1.5) + assert "critic/explained_var" in train_metrics + + +def test_train_pump_skips_the_critic_on_an_empty_chunk(monkeypatch) -> None: + """No training tokens means no GAE returns to regress against.""" + meta = _single_group_meta() + ctrl, value = _ppo_train_pump_controller(sampler=_OneThenEmptySampler(meta)) + ctrl._advantage_stage = AsyncMock(return_value=(meta, False)) + monkeypatch.setattr(single_controller.ray, "cluster_resources", lambda: {}) + + with pytest.raises(RuntimeError, match="no valid response tokens after filtering"): + asyncio.run(asyncio.wait_for(ctrl._train_pump(), timeout=1.0)) + + assert "train_from_meta" not in value.calls + # The forward still ran -- it is what the advantage stage consumes. + assert "get_values_from_meta" in value.calls + + +def test_train_pump_freezes_the_policy_during_critic_warmup(monkeypatch) -> None: + """Below policy_training_start_step the critic trains alone: no optimizer + step, and no weight transfer to generation either.""" + meta = _single_group_meta() + ctrl, value = _ppo_train_pump_controller( + sampler=_OneThenEmptySampler(meta), + policy_training_start_step=1, + ) + trainer = MagicMock(spec=_NoOpTrainer) + ctrl._trainer = trainer + ctrl._advantage_stage = AsyncMock(return_value=(meta, True)) + monkeypatch.setattr(single_controller.ray, "cluster_resources", lambda: {}) + + asyncio.run(asyncio.wait_for(ctrl._train_pump(), timeout=1.0)) + + assert "train_from_meta" in value.calls + trainer.prepare_for_training.assert_not_called() + trainer.begin_train_step.assert_not_called() + trainer.finish_train_step.assert_not_called() + ctrl._sync_weights.assert_awaited_once_with(calibration_data=None) + # The step still closed and advanced the version, so staleness accounting + # keeps working through the warmup. + assert ctrl._train_steps == 1 + assert ctrl._trainer_version == 1 + + +def test_train_pump_trains_the_policy_once_warmup_is_over(monkeypatch) -> None: + meta = _single_group_meta() + ctrl, _ = _ppo_train_pump_controller( + sampler=_OneThenEmptySampler(meta), + policy_training_start_step=1, + ) + ctrl._train_steps = 1 + ctrl._trainer_version = 1 + ctrl._algo_cfg.max_num_steps = 2 + trainer = MagicMock(spec=_NoOpTrainer) + trainer.finish_train_step.return_value = {} + ctrl._trainer = trainer + ctrl._advantage_stage = AsyncMock(return_value=(meta, True)) + monkeypatch.setattr(single_controller.ray, "cluster_resources", lambda: {}) + + asyncio.run(asyncio.wait_for(ctrl._train_pump(), timeout=1.0)) + + trainer.begin_train_step.assert_called_once() + trainer.finish_train_step.assert_called_once_with() + ctrl._sync_weights.assert_awaited_once_with(calibration_data=None) + + +def test_advantage_stage_writes_gae_returns_alongside_advantages() -> None: + """The critic's regression target has to reach TQ, or the value train step + fetches a column nobody wrote.""" + batch_size, sequence_length = 2, 4 + data = TensorDict( + { + "prompt_ids_for_adv": torch.zeros( + batch_size, sequence_length, dtype=torch.long + ), + "total_reward": torch.tensor([1.0, 0.0]), + "token_mask": torch.ones(batch_size, sequence_length), + "sample_mask": torch.ones(batch_size), + "values": torch.zeros(batch_size, sequence_length), + }, + batch_size=[batch_size], + ) + data_plane = _AdvantageDataPlane(data) + + class _GaeLikeEstimator: + def __init__(self) -> None: + self.kwargs: dict | None = None + + def compute_advantage(self, *, rewards, mask, **kwargs): + self.kwargs = kwargs + adv = rewards.unsqueeze(-1).expand_as(mask).clone() + return adv, adv + 1.0 + + estimator = _GaeLikeEstimator() + controller_cls = SingleControllerActor.__ray_metadata__.modified_class + ctrl = object.__new__(controller_cls) + ctrl._dp_client = data_plane + ctrl._advantage_cfg = AdvantageConfig() + ctrl._advantage_estimator = estimator + ctrl._policy_logprobs_required = False + ctrl._reference_logprobs_required = False + ctrl._is_ppo = True + ctrl._master_config = SimpleNamespace( + ppo=SimpleNamespace(seq_logprob_error_threshold=None) + ) + ctrl._algo_cfg = ctrl._master_config.ppo + ctrl._step_log_dict = { + "rewards": [], + "masked_advantages": [], + "sequence_lengths": [], + "seq_logprob_error_metrics": [], + } + meta = KVBatchMeta( + partition_id="rollout_data", + task_name="train", + sample_ids=[f"sample-{i}" for i in range(batch_size)], + fields=list(data.keys()), + ) + + result_meta, has_valid_training_tokens = asyncio.run(ctrl._advantage_stage(meta)) + + assert has_valid_training_tokens + assert "values" in (data_plane.selected_fields or []) + assert estimator.kwargs is not None + assert torch.equal(estimator.kwargs["values"], torch.zeros(2, 4)) + assert data_plane.written_fields is not None + assert torch.equal( + data_plane.written_fields["returns"], + torch.tensor([[2.0] * 4, [1.0] * 4]), + ) + assert "returns" in (result_meta.fields or []) + assert "advantages" in (result_meta.fields or []) From 74738e59cbd3991bf4f5b5a205102cfabf221913 Mon Sep 17 00:00:00 2001 From: Yuki Huang Date: Sun, 23 Aug 2026 04:41:31 -0700 Subject: [PATCH 10/31] feat(sc): skip the weight sync during PPO critic warmup, require the in_order sampler Signed-off-by: Yuki Huang --- nemo_rl/algorithms/single_controller.py | 12 +++++--- .../single_controller_utils/config.py | 9 ++++++ .../unit/single_controller/test_ppo_setup.py | 30 +++++++++++++++++++ 3 files changed, 47 insertions(+), 4 deletions(-) diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index ddf11147f74..4d42eba2379 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -289,6 +289,7 @@ async def run(self) -> dict[str, Any]: """Main entry point. Runs until max_train_steps is reached.""" # Synchronize weights before starting the pumps await self._sync_weights() + self._rollout_manager.set_weight_version(self._trainer_version) await self._maybe_restore_replay_buffer() await self._maybe_restore_replacement_reserve() @@ -1173,9 +1174,13 @@ async def _train_pump(self) -> None: if calibration_batches else None ) - aborted_stale_inflight_groups = await self._sync_weights( - calibration_data=calibration_data - ) + # Critic warmup doesn't need refit, and the version still advances. + aborted_stale_inflight_groups = 0 + if is_policy_training_step: + aborted_stale_inflight_groups = await self._sync_weights( + calibration_data=calibration_data + ) + self._rollout_manager.set_weight_version(self._trainer_version) step_metrics.update( { "evicted_stale_prompt_groups": evicted_stale_prompt_groups, @@ -1753,7 +1758,6 @@ async def _sync_weights( elapsed = time.monotonic() - t0 print(f" _sync_weights: sync done in {elapsed:.3f}s", flush=True) - self._rollout_manager.set_weight_version(self._trainer_version) self._rollout_permitted.set() return aborted_stale_inflight_groups diff --git a/nemo_rl/algorithms/single_controller_utils/config.py b/nemo_rl/algorithms/single_controller_utils/config.py index 88d6d7898a6..98ffc5d7be1 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/nemo_rl/algorithms/single_controller_utils/config.py @@ -756,6 +756,15 @@ def _validate_algo_settings(master_config: MasterConfig) -> None: "not have yet (#2625)." ) + sampler_name = async_config.sampler.name + if sampler_name != "in_order": + raise ValueError( + "PPO on the SingleController path only supports " + f"async_rl.sampler.name='in_order', but got '{sampler_name}'. " + "Other samplers are not supported yet (in particular during critic " + "warmup) (#2625)." + ) + rl_step_samples = ( algo_cfg.num_prompts_per_step * algo_cfg.num_generations_per_prompt ) diff --git a/tests/unit/single_controller/test_ppo_setup.py b/tests/unit/single_controller/test_ppo_setup.py index ec932f23f53..e700d6615ab 100644 --- a/tests/unit/single_controller/test_ppo_setup.py +++ b/tests/unit/single_controller/test_ppo_setup.py @@ -31,6 +31,11 @@ GAEConfig, GeneralizedAdvantageEstimator, ) +from nemo_rl.algorithms.async_utils.staleness_sampler import ( + ReadyFirstSamplerConfig, + WeightFifoSamplerConfig, + WindowedSamplerConfig, +) from nemo_rl.algorithms.grpo import GRPOConfig from nemo_rl.algorithms.loss import ClippedPGLossConfig from nemo_rl.algorithms.loss.loss_functions import MseValueLossConfig, MseValueLossFn @@ -235,6 +240,31 @@ def test_rejects_non_gae_estimator(self): with pytest.raises(ValueError): PPOConfig(adv_estimator={"name": "grpo"}) + @pytest.mark.parametrize( + "sampler_config", + [WindowedSamplerConfig(), ReadyFirstSamplerConfig(), WeightFifoSamplerConfig()], + ids=lambda cfg: cfg.name, + ) + def test_rejects_samplers_that_drop_rollouts_by_weight_version( + self, sampler_config + ): + """Critic warmup advances the version while the policy is frozen, so a + version-based sampler would evict rollouts that are not actually stale.""" + mc = _ppo_master_config() + mc.async_rl.sampler = sampler_config + + with pytest.raises( + ValueError, + match=rf"sampler.name='in_order', but got '{sampler_config.name}'", + ): + validate_single_controller_config(mc) + + def test_grpo_is_free_to_use_any_sampler(self): + mc = _make_master_config() + mc.async_rl.sampler = WindowedSamplerConfig() + + validate_single_controller_config(mc) + class TestAdvantageEstimatorSelection: def test_ppo_gets_gae_with_the_configured_lambdas(self): From 8edca3d9c30aa2842d500abf9813911e8ec331d4 Mon Sep 17 00:00:00 2001 From: Yuki Huang Date: Sun, 23 Aug 2026 06:42:29 -0700 Subject: [PATCH 11/31] feat(sc): run ppo.ppo_epochs optimizer steps per RL step, fix the stale weight-sync test assertions Signed-off-by: Yuki Huang --- nemo_rl/algorithms/single_controller.py | 93 ++++++++++---- .../single_controller_utils/config.py | 3 + .../single_controller_utils/setup.py | 13 +- .../unit/single_controller/test_ppo_setup.py | 31 ++++- .../test_single_controller_actor.py | 120 ++++++++++++++++-- 5 files changed, 219 insertions(+), 41 deletions(-) diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 4d42eba2379..6316071808f 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -133,6 +133,8 @@ def __init__( self._algo_cfg = algo_config(master_config) self._async_cfg = master_config.async_rl self._is_ppo: bool = is_ppo_run(master_config) + # GRPO has no epoch knob: it makes one optimizer step per RL step. + self._ppo_epochs: int = self._algo_cfg.ppo_epochs if self._is_ppo else 1 self._policy_logprobs_required = not ( master_config.loss_fn.force_on_policy_ratio @@ -875,8 +877,11 @@ async def _train_pump(self) -> None: chunk -- the value workers have no split train API yet (#2625). b. Policy model: train_microbatches_from_meta, which only accumulates gradients. + c. PPO only: 3a-3b repeat ppo.ppo_epochs times, and the policy's + optimizer step closes here rather than in 5 -- a PPO step is + one chunk, so there is nothing to accumulate across chunks. 4. Clear the batch. dp_client.clear_samples on the consumed sample_ids. - 5. Train the policy model -- finish_train_step all_reduces the + 5. Train the policy model (GRPO) -- finish_train_step all_reduces the accumulated gradients, rescales, and runs optimizer.step. 6. Refit the model. Sync the new policy weights to generation. """ @@ -892,7 +897,9 @@ async def _train_pump(self) -> None: step_open = False chunks_dispatched = 0 calibration_batches: list[BatchedDataDict[Any]] = [] - # One chunk per step on the PPO path, so this is the step's value model update. + # One chunk per step on the PPO path, so these are the step's own + # model updates -- the last epoch's, when there is more than one. + policy_result: Optional[dict[str, Any]] = None value_result: Optional[dict[str, Any]] = None # Always True off the PPO path: the start step is pinned to 0 there. is_policy_training_step = self._train_steps >= policy_training_start_step @@ -1014,34 +1021,63 @@ async def _train_pump(self) -> None: has_valid_training_tokens, ) = await self._advantage_stage(train_meta) - # ---- 3. Train the model -- train_microbatches_from_meta ---- - # Filtering can leave a streaming chunk with no training tokens. - # Consume that chunk without F/B, then continue the same optimizer - # step with the next chunk. + # A PPO step is this one chunk, so a chunk with nothing left + # after filtering is a step that trains neither model. + if self._is_ppo and not has_valid_training_tokens: + raise RuntimeError( + "SingleController has no valid response tokens after " + "filtering. Check ppo.seq_logprob_error_threshold to " + "avoid an optimizer step with an empty batch." + ) - # Value model first, then policy, as in the legacy PPO epoch loop. - if self._is_ppo and has_valid_training_tokens: - with self._timer.time("value_training"): - value_result = await self._value_train(train_meta) + # ---- 3. Train the model -- train_microbatches_from_meta ---- + # Filtering can leave a GRPO streaming chunk with no training + # tokens. Consume that chunk without F/B, then continue the same + # optimizer step with the next chunk. + + # GRPO runs one iteration: F/B only, its optimizer step is in 5. + # For PPO, each epoch is a full optimizer step for both models. + # TODO(#2625): value_result, policy_result only record the last epoch's metrics. + # That matches ppo.py for the losses; total_flops is additive and undercounted. + for epoch in range(self._ppo_epochs): + # Value model first, then policy, as in the legacy PPO epoch loop. + if self._is_ppo: + with self._timer.time("value_training"): + value_result = await self._value_train(train_meta) + + if is_policy_training_step: + # Always restore training mode because log-prob inference may have + # switched the model to inference mode. + with self._timer.time("training_prep"): + await asyncio.to_thread( + self._trainer.prepare_for_training + ) - if is_policy_training_step: - # Always restore training mode because log-prob inference may have - # switched the model to inference mode. - with self._timer.time("training_prep"): - await asyncio.to_thread(self._trainer.prepare_for_training) - - if has_valid_training_tokens: - with self._timer.time("policy_training"): - if not step_open: + if has_valid_training_tokens: + with self._timer.time("policy_training"): + if not step_open: + await asyncio.to_thread( + self._trainer.begin_train_step, + self._loss_fn, + ) + step_open = True await asyncio.to_thread( - self._trainer.begin_train_step, - self._loss_fn, + self._trainer.train_microbatches_from_meta, + train_meta, ) - step_open = True - await asyncio.to_thread( - self._trainer.train_microbatches_from_meta, - train_meta, - ) + # A PPO step is one chunk: nothing to + # accumulate, so close every epoch here. + if self._is_ppo: + policy_result = await asyncio.to_thread( + self._trainer.finish_train_step + ) + step_open = False + if epoch < self._ppo_epochs - 1: + # The next epoch's critic train must not + # share the training GPUs with the policy. + await asyncio.to_thread( + self._trainer.offload_to_cpu + ) if train_meta.sequence_lengths: self._step_log_dict["sequence_lengths"].extend( @@ -1112,8 +1148,9 @@ async def _train_pump(self) -> None: groups_dispatched, ) - policy_result: Optional[dict[str, Any]] = None - if is_policy_training_step: + # Only the streaming path has anything left open: a PPO step is one + # chunk, so each epoch already closed its own optimizer step above. + if not self._is_ppo: if not step_open: raise RuntimeError( "SingleController has no valid response tokens after " diff --git a/nemo_rl/algorithms/single_controller_utils/config.py b/nemo_rl/algorithms/single_controller_utils/config.py index 98ffc5d7be1..e772887a41b 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/nemo_rl/algorithms/single_controller_utils/config.py @@ -741,6 +741,9 @@ def _validate_algo_settings(master_config: MasterConfig) -> None: "See examples/configs/ppo_math_1B_megatron_single_controller.yaml." ) + if algo_cfg.ppo_epochs < 1: + raise ValueError("ppo.ppo_epochs must be at least 1") + async_config = master_config.async_rl # Without it the critic steps once per chunk and the policy once per step, # which is two effective learning rates from one config, and no error. diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py index e1950ac7dd0..723c80bad4c 100644 --- a/nemo_rl/algorithms/single_controller_utils/setup.py +++ b/nemo_rl/algorithms/single_controller_utils/setup.py @@ -431,19 +431,24 @@ def _clamp_max_num_steps( def _maybe_inject_megatron_train_iters(master_config: MasterConfig) -> None: """Set train_iters from max_num_steps after its dataloader clamp.""" - max_num_steps = algo_config(master_config).max_num_steps + algo_cfg = algo_config(master_config) + is_ppo = is_ppo_run(master_config) + # train_iters is a scheduler-tick budget, and each PPO epoch steps both + # optimizers once, so the configured warmup/decay horizon has to be scaled. + ppo_epochs = algo_cfg.ppo_epochs if is_ppo else 1 + train_iters = algo_cfg.max_num_steps * ppo_epochs # policy policy_config = master_config.policy if policy_config.get("megatron_cfg", {}).get("enabled", False): - policy_config["megatron_cfg"]["train_iters"] = max_num_steps + policy_config["megatron_cfg"]["train_iters"] = train_iters # value - if not is_ppo_run(master_config): + if not is_ppo: return value_config = master_config.value if value_config.get("megatron_cfg", {}).get("enabled", False): - value_config["megatron_cfg"]["train_iters"] = max_num_steps # type: ignore + value_config["megatron_cfg"]["train_iters"] = train_iters # type: ignore[index] def _maybe_attach_fleet_health( diff --git a/tests/unit/single_controller/test_ppo_setup.py b/tests/unit/single_controller/test_ppo_setup.py index e700d6615ab..4eb4949aaee 100644 --- a/tests/unit/single_controller/test_ppo_setup.py +++ b/tests/unit/single_controller/test_ppo_setup.py @@ -236,6 +236,16 @@ def test_rejects_value_global_batch_size_mismatch(self): ): validate_single_controller_config(mc) + def test_rejects_a_ppo_epoch_count_below_one(self): + mc = _ppo_master_config( + ppo=PPOConfig.model_construct( + max_num_steps=100, ppo_epochs=0, **_STEP_CONFIG + ) + ) + + with pytest.raises(ValueError, match="ppo_epochs must be at least 1"): + validate_single_controller_config(mc) + def test_rejects_non_gae_estimator(self): with pytest.raises(ValueError): PPOConfig(adv_estimator={"name": "grpo"}) @@ -296,13 +306,32 @@ def test_grpo_delegates_to_the_group_relative_factory(self): class TestMegatronTrainIters: def test_injects_into_both_policy_and_value(self): - mc = _ppo_master_config(megatron_enabled=True, max_num_steps=7) + mc = _ppo_master_config( + megatron_enabled=True, + ppo=PPOConfig.model_construct( + max_num_steps=7, ppo_epochs=1, **_STEP_CONFIG + ), + ) sc_setup_mod._maybe_inject_megatron_train_iters(mc) assert mc.policy["megatron_cfg"]["train_iters"] == 7 assert mc.value["megatron_cfg"]["train_iters"] == 7 + def test_scales_the_tick_budget_by_ppo_epochs(self): + """Each epoch steps both optimizers, so each is a scheduler tick.""" + mc = _ppo_master_config( + megatron_enabled=True, + ppo=PPOConfig.model_construct( + max_num_steps=7, ppo_epochs=3, **_STEP_CONFIG + ), + ) + + sc_setup_mod._maybe_inject_megatron_train_iters(mc) + + assert mc.policy["megatron_cfg"]["train_iters"] == 21 + assert mc.value["megatron_cfg"]["train_iters"] == 21 + def test_skips_a_critic_on_a_non_megatron_backend(self): mc = _ppo_master_config(megatron_enabled=False, max_num_steps=7) diff --git a/tests/unit/single_controller/test_single_controller_actor.py b/tests/unit/single_controller/test_single_controller_actor.py index a6094b81fe0..463b3ce42e3 100644 --- a/tests/unit/single_controller/test_single_controller_actor.py +++ b/tests/unit/single_controller/test_single_controller_actor.py @@ -356,8 +356,6 @@ def test_sync_weights_honors_recompute_kv_cache_config( invalidate_kv_cache=MagicMock(), requires_kv_scale_sync=False, ) - ctrl._rollout_manager = SimpleNamespace(set_weight_version=MagicMock()) - ctrl._trainer_version = 3 ctrl._inflight_by_group_id = {} # env={} -> should_use_nemo_gym is False, so _sync_weights takes the native # abort path (empty registry -> no-op) instead of the gym gate. @@ -367,7 +365,6 @@ def test_sync_weights_honors_recompute_kv_cache_config( ctrl._weight_synchronizer.sync_weights.assert_called_once_with(kv_scales=None) assert ctrl._gen.invalidate_kv_cache.call_count == expected_invalidation_calls - ctrl._rollout_manager.set_weight_version.assert_called_once_with(3) assert ctrl._rollout_permitted.is_set() @@ -385,8 +382,6 @@ def test_sync_weights_calibrates_and_forwards_fp8_kv_scales() -> None: ctrl._trainer = SimpleNamespace( calibrate_qkv_fp8_scales=MagicMock(return_value={"layers": {"layer.0": 0.5}}) ) - ctrl._rollout_manager = SimpleNamespace(set_weight_version=MagicMock()) - ctrl._trainer_version = 3 ctrl._inflight_by_group_id = {} # env={} -> should_use_nemo_gym is False, so _sync_weights takes the native # abort path (empty registry -> no-op) instead of the gym gate. @@ -773,6 +768,9 @@ def train_microbatches_from_meta(self, meta: KVBatchMeta) -> None: def finish_train_step(self) -> dict: return {} + def offload_to_cpu(self) -> None: + pass + class _LpRecordingTrainer(_NoOpTrainer): """Records the ``keep_train_buffers`` flag the pump passes on each chunk.""" @@ -808,6 +806,25 @@ def prepare_for_training(self) -> None: self.calls.append("policy.prepare_for_training") +class _EpochRecordingTrainer(_OrderRecordingTrainer): + """Also records the optimizer-step lifecycle, which the epoch loop repeats.""" + + def begin_train_step(self, loss_fn) -> None: + del loss_fn + self.calls.append("policy.begin_train_step") + + def train_microbatches_from_meta(self, meta: KVBatchMeta) -> None: + del meta + self.calls.append("policy.train_microbatches_from_meta") + + def finish_train_step(self) -> dict: + self.calls.append("policy.finish_train_step") + return {} + + def offload_to_cpu(self) -> None: + self.calls.append("policy.offload_to_cpu") + + class _NoOpDataPlane: def clear_samples(self, **kwargs) -> None: del kwargs @@ -846,9 +863,11 @@ def _train_pump_controller(*, sampler) -> object: ctrl._rollout_exhausted.set() ctrl._trainer = _NoOpTrainer() ctrl._is_ppo = False + ctrl._ppo_epochs = 1 ctrl._value = None ctrl._value_loss_fn = None ctrl._gen = SimpleNamespace(requires_kv_scale_sync=False) + ctrl._rollout_manager = SimpleNamespace(set_weight_version=MagicMock()) ctrl._loss_fn = None ctrl._dp_client = _NoOpDataPlane() ctrl._timer = Timer() @@ -1188,10 +1207,12 @@ def _ppo_train_pump_controller( sampler, policy_training_start_step: int = 0, value: _NoOpValue | None = None, + ppo_epochs: int = 1, ) -> tuple[object, _NoOpValue]: ctrl = _train_pump_controller(sampler=sampler) value = _NoOpValue() if value is None else value ctrl._is_ppo = True + ctrl._ppo_epochs = ppo_epochs ctrl._value = value ctrl._value_loss_fn = MagicMock(name="value_loss_fn") ctrl._master_config.grpo = None @@ -1328,11 +1349,12 @@ def test_train_pump_freezes_the_policy_during_critic_warmup(monkeypatch) -> None trainer.prepare_for_training.assert_not_called() trainer.begin_train_step.assert_not_called() trainer.finish_train_step.assert_not_called() - ctrl._sync_weights.assert_awaited_once_with(calibration_data=None) - # The step still closed and advanced the version, so staleness accounting - # keeps working through the warmup. + ctrl._sync_weights.assert_not_awaited() + # The step still closed and published the new version, so staleness + # accounting keeps working through the warmup. assert ctrl._train_steps == 1 assert ctrl._trainer_version == 1 + ctrl._rollout_manager.set_weight_version.assert_called_once_with(1) def test_train_pump_trains_the_policy_once_warmup_is_over(monkeypatch) -> None: @@ -1357,6 +1379,88 @@ def test_train_pump_trains_the_policy_once_warmup_is_over(monkeypatch) -> None: ctrl._sync_weights.assert_awaited_once_with(calibration_data=None) +def test_train_pump_steps_both_optimizers_once_per_ppo_epoch(monkeypatch) -> None: + """ppo_epochs repeats the whole train stage over the step's own batch.""" + meta = _single_group_meta() + ctrl, value = _ppo_train_pump_controller( + sampler=_OneThenEmptySampler(meta), ppo_epochs=2 + ) + trainer = MagicMock(spec=_NoOpTrainer) + trainer.finish_train_step.return_value = {} + ctrl._trainer = trainer + ctrl._advantage_stage = AsyncMock(return_value=(meta, True)) + monkeypatch.setattr(single_controller.ray, "cluster_resources", lambda: {}) + + asyncio.run(asyncio.wait_for(ctrl._train_pump(), timeout=1.0)) + + assert value.calls.count("train_from_meta") == 2 + assert trainer.begin_train_step.call_count == 2 + assert trainer.finish_train_step.call_count == 2 + # Still one RL step, so one refit and one version bump. + ctrl._sync_weights.assert_awaited_once_with(calibration_data=None) + assert ctrl._trainer_version == 1 + + +def test_train_pump_runs_every_critic_epoch_during_warmup(monkeypatch) -> None: + """The frozen policy does not shorten the critic's own epoch loop.""" + meta = _single_group_meta() + ctrl, value = _ppo_train_pump_controller( + sampler=_OneThenEmptySampler(meta), + policy_training_start_step=1, + ppo_epochs=2, + ) + trainer = MagicMock(spec=_NoOpTrainer) + ctrl._trainer = trainer + ctrl._advantage_stage = AsyncMock(return_value=(meta, True)) + monkeypatch.setattr(single_controller.ray, "cluster_resources", lambda: {}) + + asyncio.run(asyncio.wait_for(ctrl._train_pump(), timeout=1.0)) + + assert value.calls.count("train_from_meta") == 2 + trainer.prepare_for_training.assert_not_called() + trainer.finish_train_step.assert_not_called() + + +def test_train_pump_offloads_the_policy_between_ppo_epochs(monkeypatch) -> None: + """The two models share the training GPUs, so every critic train runs with + the policy on CPU -- including the ones after the first epoch.""" + meta = _single_group_meta() + calls: list[str] = [] + ctrl, _ = _ppo_train_pump_controller( + sampler=_OneThenEmptySampler(meta), + value=_NoOpValue(calls=calls, prefix="critic."), + ppo_epochs=2, + ) + ctrl._trainer = _EpochRecordingTrainer(calls) + ctrl._advantage_stage = AsyncMock(return_value=(meta, True)) + monkeypatch.setattr(single_controller.ray, "cluster_resources", lambda: {}) + + asyncio.run(asyncio.wait_for(ctrl._train_pump(), timeout=1.0)) + + assert calls == [ + "policy.finish_inference", + "critic.prepare_for_inference", + "critic.get_values_from_meta", + "critic.finish_inference", + "critic.prepare_for_training", + "critic.train_from_meta", + "critic.finish_training", + "policy.prepare_for_training", + "policy.begin_train_step", + "policy.train_microbatches_from_meta", + "policy.finish_train_step", + "policy.offload_to_cpu", + "critic.prepare_for_training", + "critic.train_from_meta", + "critic.finish_training", + "policy.prepare_for_training", + "policy.begin_train_step", + "policy.train_microbatches_from_meta", + # No offload after the last epoch: the refit needs the policy resident. + "policy.finish_train_step", + ] + + def test_advantage_stage_writes_gae_returns_alongside_advantages() -> None: """The critic's regression target has to reach TQ, or the value train step fetches a column nobody wrote.""" From eab2fbe29d921eccd88ef48b39486f465d5b22fd Mon Sep 17 00:00:00 2001 From: Yuki Huang Date: Sun, 23 Aug 2026 07:43:04 -0700 Subject: [PATCH 12/31] feat(sc): widen the sampler lookahead during PPO critic warmup, port the remaining warmup behaviours from ppo.py Signed-off-by: Yuki Huang --- .../async_utils/staleness_sampler.py | 53 +++++++- nemo_rl/algorithms/single_controller.py | 75 +++++++++++- .../single_controller_utils/config.py | 27 ++++ .../single_controller/test_checkpointing.py | 49 +++++++- .../unit/single_controller/test_ppo_setup.py | 24 ++++ .../test_sampler_interface.py | 71 +++++++++++ .../test_single_controller_actor.py | 115 +++++++++++++++++- 7 files changed, 400 insertions(+), 14 deletions(-) diff --git a/nemo_rl/algorithms/async_utils/staleness_sampler.py b/nemo_rl/algorithms/async_utils/staleness_sampler.py index 2ff01a0aaed..4945d54b2f4 100644 --- a/nemo_rl/algorithms/async_utils/staleness_sampler.py +++ b/nemo_rl/algorithms/async_utils/staleness_sampler.py @@ -51,7 +51,7 @@ runtime_checkable, ) -from pydantic import BaseModel, Field, NonNegativeInt +from pydantic import BaseModel, Field, NonNegativeInt, model_validator from nemo_rl.algorithms.async_utils.replay_buffer import TQReplayBuffer from nemo_rl.data_plane import KVBatchMeta @@ -327,6 +327,10 @@ class _GatedSampler(BaseSampler): def __init__(self, buffer: TQReplayBuffer, *, gate_window: int) -> None: super().__init__(buffer) + self._gate_window = 0 + self.set_gate_window(gate_window) + + def set_gate_window(self, gate_window: int) -> None: if gate_window < 0: raise ValueError(f"gate_window must be non-negative, got {gate_window}") self._gate_window = gate_window @@ -441,9 +445,25 @@ class InOrderSampler(_GatedSampler): upcoming is never dropped early, and evict/select can't disagree. """ - def __init__(self, buffer: TQReplayBuffer, *, max_lookahead_versions: int) -> None: + def __init__( + self, + buffer: TQReplayBuffer, + *, + max_lookahead_versions: int, + warmup_lookahead_versions: Optional[int] = None, + ) -> None: super().__init__(buffer, gate_window=max_lookahead_versions) self.max_lookahead_versions = max_lookahead_versions + self.warmup_lookahead_versions = warmup_lookahead_versions + + def required_buffer_capacity(self, groups_per_step: int) -> Optional[int]: + # Sized for the peak window: otherwise the buffer, not the gate, bounds + # the lookahead and widening it during warmup does nothing. + if self.warmup_lookahead_versions is None: + peak = self.max_lookahead_versions + else: + peak = self.warmup_lookahead_versions + return _gated_required_buffer_capacity(groups_per_step, gate_window=peak) def _stamp(self) -> Optional[int]: return self._dispatch_index @@ -508,6 +528,27 @@ 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: NonNegativeInt = 1 + # Widened lookahead while the PPO policy is frozen; None keeps the steady value. + warmup_lookahead_versions: Optional[NonNegativeInt] = None + + @model_validator(mode="after") + def validate_warmup_lookahead(self) -> "InOrderSamplerConfig": + if ( + self.warmup_lookahead_versions is not None + and self.warmup_lookahead_versions < self.max_lookahead_versions + ): + raise ValueError( + "warmup_lookahead_versions must be greater than or equal " + "to max_lookahead_versions" + ) + return self + + @property + def peak_lookahead_versions(self) -> int: + """Widest window the run can reach; what buffer capacity must cover.""" + if self.warmup_lookahead_versions is None: + return self.max_lookahead_versions + return self.warmup_lookahead_versions class CustomSamplerConfig(BaseModel, extra="allow"): @@ -550,7 +591,7 @@ def required_buffer_capacity_for_config( if isinstance(cfg, InOrderSamplerConfig): return _gated_required_buffer_capacity( groups_per_step, - gate_window=cfg.max_lookahead_versions, + gate_window=cfg.peak_lookahead_versions, ) return None @@ -576,7 +617,11 @@ 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, + warmup_lookahead_versions=cfg.warmup_lookahead_versions, + ) if isinstance(cfg, CustomSamplerConfig): module_name, sep, class_name = cfg.target.partition(":") if not sep: diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 6316071808f..0de75a38416 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -39,6 +39,7 @@ import math import os import time +import warnings from collections import deque from functools import partial from typing import Any, Awaitable, Callable, Optional, Union @@ -276,6 +277,11 @@ def __init__( "seq_logprob_error_metrics": [], } + # Seeded here rather than in run(): on resume _trainer_version is the + # checkpoint's step, so a run resuming mid-warmup needs the widened + # window before the first dispatch. + self._retune_lookahead_versions() + print( f"SingleControllerActor: " f"sampler={self._async_cfg.sampler.name} " @@ -1046,6 +1052,17 @@ async def _train_pump(self) -> None: value_result = await self._value_train(train_meta) if is_policy_training_step: + if ( + self._is_ppo + and self._train_steps == policy_training_start_step + and policy_training_start_step > 0 + and epoch == 0 + ): + print( + f" ✓ Critic warmup complete ({policy_training_start_step} " + "steps). Starting policy training.", + flush=True, + ) # Always restore training mode because log-prob inference may have # switched the model to inference mode. with self._timer.time("training_prep"): @@ -1217,6 +1234,7 @@ async def _train_pump(self) -> None: aborted_stale_inflight_groups = await self._sync_weights( calibration_data=calibration_data ) + self._retune_lookahead_versions() self._rollout_manager.set_weight_version(self._trainer_version) step_metrics.update( { @@ -1272,7 +1290,10 @@ async def _train_pump(self) -> None: should_save_by_step or should_save_by_timeout ): with self._timer.time("checkpointing"): - await self._save_checkpoint(step_metrics) + await self._save_checkpoint( + step_metrics, + is_policy_training_step=is_policy_training_step, + ) timing_metrics: dict[str, float] = self._timer.get_timing_metrics( reduction_op="sum" @@ -1615,11 +1636,18 @@ async def _abort_stale_inflight(self) -> int: ) return len(stale_tasks) - async def _save_checkpoint(self, step_metrics: dict[str, Any]) -> None: + async def _save_checkpoint( + self, + step_metrics: dict[str, Any], + *, + is_policy_training_step: bool, + ) -> None: """Write a full checkpoint for the just-finished train step. Everything except the (possibly async) policy weight write must be on disk before begin_finalization; rollouts keep running throughout. + The policy optimizer is skipped during critic warmup -- it has never + stepped. """ save_state = self._save_state save_state.current_step = self._train_steps @@ -1652,9 +1680,19 @@ async def _save_checkpoint(self, step_metrics: dict[str, Any]) -> None: full_metric_name = self._master_config.checkpointing["metric_name"] if full_metric_name is not None: metric_name = full_metric_name.split(":", 1)[1] - if metric_name not in step_metrics: + if not is_policy_training_step: + warnings.warn( + f"checkpointing.metric_name={full_metric_name!r} is not " + "available during PPO critic warmup; this checkpoint will " + "not be saved as top-k.", + stacklevel=2, + ) + if hasattr(save_state, full_metric_name): + delattr(save_state, full_metric_name) + elif metric_name not in step_metrics: raise ValueError(f"Metric {metric_name} not found in train metrics") - setattr(save_state, full_metric_name, step_metrics[metric_name]) + else: + setattr(save_state, full_metric_name, step_metrics[metric_name]) # Flush the previous checkpoint's background finalization first; # re-raises a failure from it. @@ -1667,6 +1705,8 @@ async def _save_checkpoint(self, step_metrics: dict[str, Any]) -> None: vars(save_state), self._master_config, ) + + # Save value model if self._is_ppo: # The critic shares the training GPUs, so the two weight saves are # serialized. The critic goes first because offloading the policy runs @@ -1689,17 +1729,24 @@ async def _save_checkpoint(self, step_metrics: dict[str, Any]) -> None: # Also covers a warmup step, which never ran prepare_for_training in # the pump and would otherwise save from CPU-resident params. await asyncio.to_thread(self._trainer.prepare_for_training) + + # Save policy model # With async_save this returns after D2H staging; disk writes finish # in the background. await asyncio.to_thread( self._trainer.save_checkpoint, weights_path=os.path.join(checkpoint_path, "policy", "weights"), + # Always save policy weights so every PPO checkpoint has + # the same component layout. Before the first real policy + # update, omit optimizer and scheduler state because their + # lazily initialized state is not yet safe to checkpoint. optimizer_path=os.path.join(checkpoint_path, "policy", "optimizer") - if self._checkpointer.save_optimizer + if self._checkpointer.save_optimizer and is_policy_training_step else None, tokenizer_path=os.path.join(checkpoint_path, "policy", "tokenizer"), checkpointing_cfg=self._master_config.checkpointing, ) + await asyncio.to_thread( torch.save, dataloader_state, @@ -1719,6 +1766,7 @@ async def _save_checkpoint(self, step_metrics: dict[str, Any]) -> None: buffer_state, os.path.join(checkpoint_path, "replay_buffer.pt"), ) + # Rename happens in the background once the async weight writes # finish; flushed at the next save or on exit. self._checkpointer.begin_finalization( @@ -1997,3 +2045,20 @@ def _advantage_input_fields(self) -> list[str]: if self._is_ppo: fields.append(adv_cfg.values_field) return list(dict.fromkeys(fields)) + + def _retune_lookahead_versions(self) -> None: + """Widen the sampler's lookahead while the policy is frozen, then shrink it back. + + Port of ppo.py's _async_ppo_generation_lead_steps. + """ + if not self._is_ppo: + return + steady = self._async_cfg.sampler.max_lookahead_versions + warmup = self._async_cfg.sampler.warmup_lookahead_versions + start = self._algo_cfg.policy_training_start_step + if warmup is None or self._trainer_version >= start: + window = steady + else: + remaining_to_frontier = start + steady - self._trainer_version + window = max(steady, min(warmup, remaining_to_frontier)) + self._sampler.set_gate_window(window) diff --git a/nemo_rl/algorithms/single_controller_utils/config.py b/nemo_rl/algorithms/single_controller_utils/config.py index e772887a41b..b1cbf4b2099 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/nemo_rl/algorithms/single_controller_utils/config.py @@ -759,6 +759,33 @@ def _validate_algo_settings(master_config: MasterConfig) -> None: "not have yet (#2625)." ) + policy_megatron_cfg = master_config.policy.get("megatron_cfg", {}) # type: ignore + if ( + getattr(algo_cfg, "policy_training_start_step", 0) > 0 + and master_config.checkpointing["enabled"] + and master_config.checkpointing["save_optimizer"] + and policy_megatron_cfg.get("enabled") + and policy_megatron_cfg.get("checkpoint", {}).get( + "ckpt_assume_constant_structure" + ) + ): + raise ValueError( + "policy.megatron_cfg.checkpoint.ckpt_assume_constant_structure=true " + "is incompatible with PPO critic warmup when optimizer checkpointing " + "is enabled. Set ckpt_assume_constant_structure=false, " + "ppo.policy_training_start_step=0, or checkpointing.save_optimizer=false." + ) + + if ( + getattr(async_config.sampler, "warmup_lookahead_versions", None) is not None + and getattr(algo_cfg, "policy_training_start_step", 0) == 0 + ): + raise ValueError( + "async_rl.sampler.warmup_lookahead_versions requires " + "ppo.policy_training_start_step > 0; without critic warmup there is " + "no frozen-policy window to widen." + ) + sampler_name = async_config.sampler.name if sampler_name != "in_order": raise ValueError( diff --git a/tests/unit/single_controller/test_checkpointing.py b/tests/unit/single_controller/test_checkpointing.py index 1c026b5f6ad..f0a4472bb21 100644 --- a/tests/unit/single_controller/test_checkpointing.py +++ b/tests/unit/single_controller/test_checkpointing.py @@ -196,6 +196,9 @@ def is_on_policy(self) -> bool: def required_buffer_capacity(self, groups_per_step: int) -> Optional[int]: return None + def set_gate_window(self, gate_window: int) -> None: + self.gate_window = gate_window + def set_dispatch_index(self, resume_from_step: int) -> None: pass @@ -793,7 +796,7 @@ def prepare_for_training(self) -> None: self.calls.append("policy.prepare_for_training") def save_checkpoint(self, **kwargs: Any) -> None: - del kwargs + self.save_kwargs = kwargs self.calls.append("policy.save_checkpoint") def finalize_async_save(self) -> None: @@ -810,7 +813,7 @@ def prepare_for_training(self) -> None: self.calls.append("critic.prepare_for_training") def save_checkpoint(self, **kwargs: Any) -> None: - del kwargs + self.save_kwargs = kwargs self.calls.append("critic.save_checkpoint") def finish_training(self) -> None: @@ -863,7 +866,7 @@ def test_the_policy_is_offloaded_across_the_critic_save( calls: list[str] = [] actor = _ppo_save_actor(tmp_path, calls) - asyncio.run(actor._save_checkpoint({})) + asyncio.run(actor._save_checkpoint({}, is_policy_training_step=True)) assert calls == [ "policy.offload_to_cpu", @@ -875,6 +878,46 @@ def test_the_policy_is_offloaded_across_the_critic_save( ] +class TestPPOWarmupCheckpoint: + """During critic warmup the policy optimizer has never stepped.""" + + @pytest.fixture + def actor(self, tmp_path, monkeypatch): + monkeypatch.setattr( + "nemo_rl.algorithms.single_controller._write_latest_checkpoint_status", + lambda *args, **kwargs: None, + ) + return _ppo_save_actor(tmp_path, []) + + def test_warmup_step_writes_no_policy_optimizer(self, actor): + asyncio.run(actor._save_checkpoint({}, is_policy_training_step=False)) + + assert actor._trainer.save_kwargs["optimizer_path"] is None + # The critic trains from step 0, so its optimizer is still written. + assert actor._value.save_kwargs["optimizer_path"] is not None + + def test_training_step_writes_the_policy_optimizer(self, actor): + asyncio.run(actor._save_checkpoint({}, is_policy_training_step=True)) + + assert actor._trainer.save_kwargs["optimizer_path"] is not None + + def test_warmup_step_skips_the_top_k_metric(self, actor): + """No policy metrics exist yet, so the checkpoint just is not a candidate.""" + actor._master_config.checkpointing["metric_name"] = "train:loss" + + with pytest.warns(UserWarning, match="not available during PPO critic warmup"): + asyncio.run(actor._save_checkpoint({}, is_policy_training_step=False)) + + assert not hasattr(actor._save_state, "train:loss") + + def test_a_training_step_still_raises_on_a_missing_metric(self, actor): + """The warmup branch must not soften the misconfiguration error.""" + actor._master_config.checkpointing["metric_name"] = "train:loss" + + with pytest.raises(ValueError, match="not found in train metrics"): + asyncio.run(actor._save_checkpoint({}, is_policy_training_step=True)) + + # ── metric_name behavior ───────────────────────────────────────────────────── diff --git a/tests/unit/single_controller/test_ppo_setup.py b/tests/unit/single_controller/test_ppo_setup.py index 4eb4949aaee..c360062edf4 100644 --- a/tests/unit/single_controller/test_ppo_setup.py +++ b/tests/unit/single_controller/test_ppo_setup.py @@ -269,6 +269,30 @@ def test_rejects_samplers_that_drop_rollouts_by_weight_version( ): validate_single_controller_config(mc) + @staticmethod + def _warmup_ckpt_config(*, constant_structure: bool) -> MasterConfig: + mc = _ppo_master_config(megatron_enabled=True) + mc.ppo.policy_training_start_step = 2 + mc.policy["megatron_cfg"]["checkpoint"] = { + "ckpt_assume_constant_structure": constant_structure + } + mc.checkpointing["enabled"] = True + mc.checkpointing["save_optimizer"] = True + return mc + + def test_rejects_constant_ckpt_structure_with_warmup(self): + """The policy optimizer first appears when warmup ends, so one cached + layout cannot describe both states.""" + mc = self._warmup_ckpt_config(constant_structure=True) + + with pytest.raises(ValueError, match="ckpt_assume_constant_structure=true"): + validate_single_controller_config(mc) + + def test_accepts_varying_ckpt_structure_with_warmup(self): + validate_single_controller_config( + self._warmup_ckpt_config(constant_structure=False) + ) + def test_grpo_is_free_to_use_any_sampler(self): mc = _make_master_config() mc.async_rl.sampler = WindowedSamplerConfig() diff --git a/tests/unit/single_controller/test_sampler_interface.py b/tests/unit/single_controller/test_sampler_interface.py index c62e1116a48..973983a5ca7 100644 --- a/tests/unit/single_controller/test_sampler_interface.py +++ b/tests/unit/single_controller/test_sampler_interface.py @@ -232,6 +232,77 @@ def test_required_capacity_covers_live_and_lookahead_batches(self): assert sampler.required_buffer_capacity(groups_per_step=4) == 12 +class TestWarmupLookaheadWindow: + """The PPO critic warmup widens the gate, so capacity must cover the peak.""" + + def test_capacity_is_sized_from_the_warmup_window(self): + cfg = InOrderSamplerConfig( + max_lookahead_versions=1, warmup_lookahead_versions=3 + ) + + # Steady state alone would be 4*(1+1)=8; the warmup peak needs 4*(3+1)=16. + assert required_buffer_capacity_for_config(cfg, groups_per_step=4) == 16 + sampler = create_sampler(FakeBuffer(), cfg) + assert sampler.required_buffer_capacity(groups_per_step=4) == 16 + + def test_capacity_is_unchanged_without_a_warmup_window(self): + cfg = InOrderSamplerConfig(max_lookahead_versions=1) + + assert required_buffer_capacity_for_config(cfg, groups_per_step=4) == 8 + sampler = create_sampler(FakeBuffer(), cfg) + assert sampler.required_buffer_capacity(groups_per_step=4) == 8 + + def test_capacity_does_not_shrink_when_the_gate_is_retuned(self): + """Retuning must not let the reported requirement follow the live window.""" + sampler = InOrderSampler( + FakeBuffer(), max_lookahead_versions=1, warmup_lookahead_versions=3 + ) + + sampler.set_gate_window(1) + + assert sampler.required_buffer_capacity(groups_per_step=4) == 16 + + def test_set_gate_window_retunes_admission(self): + sampler = InOrderSampler(FakeBuffer(), max_lookahead_versions=1) + + sampler.set_gate_window(3) + + assert sampler._gate_window == 3 + + def test_set_gate_window_rejects_a_negative_window(self): + sampler = InOrderSampler(FakeBuffer(), max_lookahead_versions=1) + + with pytest.raises(ValueError, match="gate_window must be non-negative"): + sampler.set_gate_window(-1) + + def test_only_gated_samplers_can_be_retuned(self): + """WindowedSampler has no gate, so it deliberately has no setter. + + SC PPO is validated to run under in_order, so the driver never reaches + a sampler that lacks it. + """ + assert not hasattr( + WindowedSampler(FakeBuffer(), max_staleness_versions=1), + "set_gate_window", + ) + + def test_warmup_window_below_the_steady_window_is_rejected(self): + with pytest.raises(ValidationError): + InOrderSamplerConfig(max_lookahead_versions=2, warmup_lookahead_versions=1) + + def test_discriminated_union_parses_the_warmup_window(self): + cfg = TypeAdapter(SamplerConfig).validate_python( + { + "name": "in_order", + "max_lookahead_versions": 1, + "warmup_lookahead_versions": 4, + } + ) + + assert isinstance(cfg, InOrderSamplerConfig) + assert cfg.warmup_lookahead_versions == 4 + + class TestCustomFqnSampler: def test_custom_target_loads_out_of_repo_sampler(self): # A user sampler defined anywhere importable; here, this test module. diff --git a/tests/unit/single_controller/test_single_controller_actor.py b/tests/unit/single_controller/test_single_controller_actor.py index 463b3ce42e3..185b0990b95 100644 --- a/tests/unit/single_controller/test_single_controller_actor.py +++ b/tests/unit/single_controller/test_single_controller_actor.py @@ -336,6 +336,105 @@ def test_logs_setup_timing_metrics(monkeypatch, tmp_path) -> None: ) +def _lookahead_controller( + *, + trainer_version: int, + policy_training_start_step: int, + max_lookahead_versions: int = 1, + warmup_lookahead_versions: int | None = None, + is_ppo: bool = True, +): + """Bare actor carrying only what the lookahead schedule reads.""" + controller_cls = SingleControllerActor.__ray_metadata__.modified_class + ctrl = object.__new__(controller_cls) + ctrl._is_ppo = is_ppo + ctrl._trainer_version = trainer_version + ctrl._algo_cfg = SimpleNamespace( + policy_training_start_step=policy_training_start_step + ) + ctrl._async_cfg = SimpleNamespace( + sampler=SimpleNamespace( + max_lookahead_versions=max_lookahead_versions, + warmup_lookahead_versions=warmup_lookahead_versions, + ) + ) + ctrl._sampler = MagicMock() + return ctrl + + +class TestLookaheadSchedule: + """Port of ppo.py's _async_ppo_generation_lead_steps. + + Generation may run further ahead while the policy is frozen, then the window + has to converge back before warmup-era rollouts stop being trainable. + """ + + @pytest.mark.parametrize( + ("trainer_version", "expected"), + [ + # start=4, steady=1, warmup=5 -> frontier is 4+1=5. + (0, 5), # min(5, 5-0) = 5 + (1, 4), # min(5, 5-1) = 4, already shrinking + (3, 2), + (4, 1), # warmup over: back to steady + (9, 1), + ], + ) + def test_window_widens_then_converges(self, trainer_version, expected): + ctrl = _lookahead_controller( + trainer_version=trainer_version, + policy_training_start_step=4, + max_lookahead_versions=1, + warmup_lookahead_versions=5, + ) + + ctrl._retune_lookahead_versions() + + ctrl._sampler.set_gate_window.assert_called_once_with(expected) + + def test_steady_value_is_a_defensive_floor(self): + """Kept because ppo.py has it, but unreachable through a valid config. + + The floor only binds when warmup < steady, which + InOrderSamplerConfig.validate_warmup_lookahead already rejects -- inside + the warmup branch the frontier term is always greater than steady. This + builds the config by hand to reach it at all. + """ + ctrl = _lookahead_controller( + trainer_version=2, + policy_training_start_step=100, + max_lookahead_versions=3, + warmup_lookahead_versions=2, + ) + + ctrl._retune_lookahead_versions() + + ctrl._sampler.set_gate_window.assert_called_once_with(3) + + def test_unset_warmup_window_pins_the_steady_value(self): + ctrl = _lookahead_controller( + trainer_version=0, + policy_training_start_step=4, + max_lookahead_versions=2, + warmup_lookahead_versions=None, + ) + + ctrl._retune_lookahead_versions() + + ctrl._sampler.set_gate_window.assert_called_once_with(2) + + def test_retune_is_a_noop_off_the_ppo_path(self): + ctrl = _lookahead_controller( + trainer_version=0, + policy_training_start_step=0, + is_ppo=False, + ) + + ctrl._retune_lookahead_versions() + + ctrl._sampler.set_gate_window.assert_not_called() + + @pytest.mark.parametrize( ("recompute_kv_cache", "expected_invalidation_calls"), [(False, 0), (True, 1)], @@ -685,6 +784,9 @@ async def evict(self, *, current_train_weight: int) -> int: del current_train_weight return 0 + def set_gate_window(self, gate_window: int) -> None: + self.gate_window = gate_window + async def select(self, **kwargs): del kwargs return None, 0 @@ -846,6 +948,10 @@ def _train_pump_controller(*, sampler) -> object: ctrl._async_cfg = SimpleNamespace( min_groups_for_streaming_train=1, rollout_failure=SimpleNamespace(min_step_batch_fraction=0.9), + sampler=SimpleNamespace( + max_lookahead_versions=1, + warmup_lookahead_versions=None, + ), ) ctrl._consumed_samples = 0 ctrl._total_valid_tokens = 0 @@ -1330,7 +1436,9 @@ def test_train_pump_skips_the_critic_on_an_empty_chunk(monkeypatch) -> None: assert "get_values_from_meta" in value.calls -def test_train_pump_freezes_the_policy_during_critic_warmup(monkeypatch) -> None: +def test_train_pump_freezes_the_policy_during_critic_warmup( + monkeypatch, capsys +) -> None: """Below policy_training_start_step the critic trains alone: no optimizer step, and no weight transfer to generation either.""" meta = _single_group_meta() @@ -1355,9 +1463,10 @@ def test_train_pump_freezes_the_policy_during_critic_warmup(monkeypatch) -> None assert ctrl._train_steps == 1 assert ctrl._trainer_version == 1 ctrl._rollout_manager.set_weight_version.assert_called_once_with(1) + assert "Critic warmup complete" not in capsys.readouterr().out -def test_train_pump_trains_the_policy_once_warmup_is_over(monkeypatch) -> None: +def test_train_pump_trains_the_policy_once_warmup_is_over(monkeypatch, capsys) -> None: meta = _single_group_meta() ctrl, _ = _ppo_train_pump_controller( sampler=_OneThenEmptySampler(meta), @@ -1377,6 +1486,8 @@ def test_train_pump_trains_the_policy_once_warmup_is_over(monkeypatch) -> None: trainer.begin_train_step.assert_called_once() trainer.finish_train_step.assert_called_once_with() ctrl._sync_weights.assert_awaited_once_with(calibration_data=None) + # Announced exactly once, on the step that crosses the boundary. + assert capsys.readouterr().out.count("Critic warmup complete") == 1 def test_train_pump_steps_both_optimizers_once_per_ppo_epoch(monkeypatch) -> None: From 7e9d9a4aadca7eeb9d14246876864d442464151b Mon Sep 17 00:00:00 2001 From: Yuki Huang Date: Sun, 23 Aug 2026 08:11:14 -0700 Subject: [PATCH 13/31] feat(sc): widen the sampler lookahead during PPO critic warmup, port the remaining warmup behaviours from ppo.py Signed-off-by: Yuki Huang --- nemo_rl/data_plane/driver_mixin.py | 26 +++++++++++ nemo_rl/models/policy/tq_policy.py | 10 ++--- nemo_rl/models/value/tq_value.py | 7 +-- tests/unit/models/value/test_tq_value.py | 55 +++++++++++++++++++++++- 4 files changed, 85 insertions(+), 13 deletions(-) diff --git a/nemo_rl/data_plane/driver_mixin.py b/nemo_rl/data_plane/driver_mixin.py index ce6c2b2d2c1..2dd6df0feed 100644 --- a/nemo_rl/data_plane/driver_mixin.py +++ b/nemo_rl/data_plane/driver_mixin.py @@ -15,6 +15,7 @@ from __future__ import annotations +from dataclasses import replace from typing import Any, Optional from nemo_rl.data_plane.column_io import read_columns, round_up, write_columns @@ -70,6 +71,31 @@ def _stamp_pad_seqlen(self, meta: KVBatchMeta) -> None: max(meta.sequence_lengths), max(pad_mult, seq_round) ) + def _isolated_meta( + self, + meta: KVBatchMeta, + *, + fields: list[str], + task_name: str, + ) -> KVBatchMeta: + """Narrow ``meta`` for one model's dispatch and mint it a fresh pad target. + + The mint is idempotent, so sharing or inheriting the target would let + whichever model dispatches first decide the forward pad for the rest -- + and with ppo_epochs > 1 the caller's meta is already stamped when the + critic dispatches again. + """ + extra_info = dict(meta.extra_info) + extra_info.pop(GLOBAL_FORWARD_PAD_SEQLEN, None) + isolated = replace( + meta, + fields=fields, + task_name=task_name, + extra_info=extra_info, + ) + self._stamp_pad_seqlen(isolated) + return isolated + def read_from_dataplane( self, meta: KVBatchMeta, diff --git a/nemo_rl/models/policy/tq_policy.py b/nemo_rl/models/policy/tq_policy.py index 6fb4becde70..bbc412cf1d0 100644 --- a/nemo_rl/models/policy/tq_policy.py +++ b/nemo_rl/models/policy/tq_policy.py @@ -32,7 +32,6 @@ import warnings from collections import defaultdict from contextlib import nullcontext -from dataclasses import replace from typing import Any, Optional import ray @@ -219,9 +218,8 @@ def _logprob_dispatch( leader-rank ``_write_back_result_field``; the Ray return is always None, so this dispatcher just waits for completion. """ - self._stamp_pad_seqlen(meta) spa, dba = self._packing_args("logprob_mb_tokens") - lp_meta = replace( + lp_meta = self._isolated_meta( meta, fields=fields_with_optional_routed_experts( LP_SEED_FIELDS, @@ -323,13 +321,12 @@ def train_from_meta( batch_size = gbs or self.cfg["train_global_batch_size"] micro_batch_size = mbs or self.cfg["train_micro_batch_size"] - self._stamp_pad_seqlen(meta) spa, dba = self._packing_args("train_mb_tokens") # ``train_fields`` (rollout + logprob deltas + advantages + sample_mask; # default ``DP_TRAIN_FIELDS``) must be in TQ before this call — written # by workers + driver delta-writes. Caller may narrow to drop columns # skipped this step (e.g. ``prev_logprobs`` under force_on_policy_ratio). - train_meta = replace( + train_meta = self._isolated_meta( meta, fields=fields_with_optional_routed_experts( train_fields, enabled=self._router_replay_enabled @@ -451,9 +448,8 @@ def train_microbatches_from_meta( the workers' open-step state and surface once via :meth:`finish_train_step`. """ - self._stamp_pad_seqlen(meta) spa, dba = self._packing_args("train_mb_tokens") - train_meta = replace( + train_meta = self._isolated_meta( meta, fields=fields_with_optional_routed_experts( DP_TRAIN_FIELDS, enabled=self._router_replay_enabled diff --git a/nemo_rl/models/value/tq_value.py b/nemo_rl/models/value/tq_value.py index 317d155373f..062df22f12a 100644 --- a/nemo_rl/models/value/tq_value.py +++ b/nemo_rl/models/value/tq_value.py @@ -17,7 +17,6 @@ import warnings from contextlib import nullcontext -from dataclasses import replace from typing import Any, Optional import ray @@ -110,9 +109,8 @@ def get_values_from_meta( micro_batch_size: Inference micro batch size; None uses the config default. timer: Optional timer for nested get_values measurements. """ - self._stamp_pad_seqlen(meta) spa, dba = self._packing_args("logprob_mb_tokens") - value_meta = replace( + value_meta = self._isolated_meta( meta, fields=list(VALUE_SEED_FIELDS), task_name="value_fwd", @@ -162,9 +160,8 @@ def train_from_meta( batch_size = gbs or self.cfg["train_global_batch_size"] micro_batch_size = mbs or self.cfg["train_micro_batch_size"] - self._stamp_pad_seqlen(meta) spa, dba = self._packing_args("train_mb_tokens") - train_meta = replace( + train_meta = self._isolated_meta( meta, fields=list(DP_VALUE_TRAIN_FIELDS), task_name="value_train", diff --git a/tests/unit/models/value/test_tq_value.py b/tests/unit/models/value/test_tq_value.py index 7080f0a9583..9c6173e4daa 100644 --- a/tests/unit/models/value/test_tq_value.py +++ b/tests/unit/models/value/test_tq_value.py @@ -29,7 +29,11 @@ import torch from nemo_rl.data_plane import KVBatchMeta -from nemo_rl.data_plane.schema import DP_VALUE_TRAIN_FIELDS, VALUE_SEED_FIELDS +from nemo_rl.data_plane.schema import ( + DP_VALUE_TRAIN_FIELDS, + GLOBAL_FORWARD_PAD_SEQLEN, + VALUE_SEED_FIELDS, +) from nemo_rl.data_plane.worker_mixin import TQWorkerMixin from nemo_rl.models.value.tq_value import TQValue @@ -212,3 +216,52 @@ def _result(loss: float) -> dict: out = v.train_from_meta(meta, loss_fn="LF") assert out["all_mb_metrics"]["loss"] == [0.1, 0.2] + + +class TestPadTargetIsolation: + """Each dispatch mints its own pad target, in both directions: the critic + goes first within a step, and with ppo_epochs > 1 the policy has already + stamped by the time the critic runs again.""" + + def test_dispatch_does_not_stamp_the_callers_meta(self): + v, _ = _make_tq_value() + v.use_dynamic_batches = False + v.use_sequence_packing = False + meta = KVBatchMeta( + partition_id="rollout_data", + task_name="train", + sample_ids=["s0", "s1"], + sequence_lengths=[7, 9], + ) + with patch( + "nemo_rl.models.value.tq_value.shard_meta_for_dp", + return_value=([meta, meta], None), + ) as mock_shard: + v.get_values_from_meta(meta) + + assert GLOBAL_FORWARD_PAD_SEQLEN not in meta.extra_info + # ...but the dispatched meta carries one, so DP ranks still agree. + assert GLOBAL_FORWARD_PAD_SEQLEN in mock_shard.call_args.args[0].extra_info + + def test_a_stamped_caller_meta_does_not_decide_the_critic_pad(self): + """ppo_epochs > 1: the policy has stamped the shared meta by epoch 1.""" + v, _ = _make_tq_value() + v.use_dynamic_batches = False + v.use_sequence_packing = False + meta = KVBatchMeta( + partition_id="rollout_data", + task_name="train", + sample_ids=["s0", "s1"], + sequence_lengths=[7, 9], + extra_info={GLOBAL_FORWARD_PAD_SEQLEN: 4096}, + ) + with patch( + "nemo_rl.models.value.tq_value.shard_meta_for_dp", + return_value=([meta, meta], None), + ) as mock_shard: + v.get_values_from_meta(meta) + + dispatched = mock_shard.call_args.args[0] + assert dispatched.extra_info[GLOBAL_FORWARD_PAD_SEQLEN] != 4096 + # The caller's value survives for whoever minted it. + assert meta.extra_info[GLOBAL_FORWARD_PAD_SEQLEN] == 4096 From a63dcc9691436a8b38046d1187c8cf56268ce35d Mon Sep 17 00:00:00 2001 From: Yuki Huang Date: Sun, 23 Aug 2026 08:57:43 -0700 Subject: [PATCH 14/31] test(sc): add an async PPO SingleController functional test and its Megatron exemplar Signed-off-by: Yuki Huang --- ...po_math_1B_megatron_single_controller.yaml | 200 ++++++++++++++++++ .../L1_Functional_Tests_SingleController.sh | 6 + .../functional/ppo_async_single_controller.sh | 123 +++++++++++ 3 files changed, 329 insertions(+) create mode 100644 examples/configs/ppo_math_1B_megatron_single_controller.yaml create mode 100755 tests/functional/ppo_async_single_controller.sh diff --git a/examples/configs/ppo_math_1B_megatron_single_controller.yaml b/examples/configs/ppo_math_1B_megatron_single_controller.yaml new file mode 100644 index 00000000000..ecb6aca5d62 --- /dev/null +++ b/examples/configs/ppo_math_1B_megatron_single_controller.yaml @@ -0,0 +1,200 @@ +# SingleController variant of ppo_math_1B_megatron.yaml. +defaults: ppo_math_1B_megatron.yaml + +ppo: + # Null out the legacy async_ppo block inherited from ppo_math_1B.yaml: SC reads + # async_rl instead, and the entrypoint raises if this block is set. It defaults + # to a non-null AsyncPPOConfig, so omitting it is not the same as nulling it. + async_ppo: null + # SC does not support validation yet. + val_period: 0 + val_at_start: false + # SC implements neither, so leaving them enabled would describe shaping this + # run does not do. + reward_shaping: + enabled: false + reward_scaling: + enabled: false + +async_rl: + sampler: + name: in_order + # How far generation may run ahead of the trainer. + max_lookahead_versions: 1 + # Widened while the policy is frozen by policy_training_start_step, so critic + # warmup is not also a generation stall. null keeps the value above + # throughout. Buffer capacity below must cover the wider of the two. + warmup_lookahead_versions: null + # Recompute generation KV caches after each weight update. + recompute_kv_cache_after_weight_updates: false + # Min ready groups the streaming trainer waits for before dispatching a batch. + min_groups_for_streaming_train: ${ppo.num_prompts_per_step} + # Cap on in-flight generate_and_push calls in the rollout pump. + max_inflight_prompts: ${ppo.num_prompts_per_step} + # Cap on unconsumed rollout groups buffered in the DataPlane (backpressure). + max_buffered_rollouts: 64 + # Enable per-rollout diagnostic prints (prompt content / completion previews). + diagnostics: false + + # ── Resiliency ───────────────────────────────────────────────────────────── + # What happens to a prompt whose rollout fails. Infrastructure failures are + # re-dispatched (the retry re-enters shard selection, so it lands elsewhere); + # deterministic per-prompt failures get a much smaller budget because another + # shard would fail identically. + # + # The budgets here govern BOTH rollout paths -- the retry loop sits above the + # native/NeMo-Gym split. Anything path-specific lives in the sub-blocks below, + # so the structure says which knob applies where. Filling in the block for the + # path a run is not taking is rejected at setup rather than silently ignored. + rollout_failure: + max_infra_attempts_per_prompt: 5 # infra budget; exhausting it drops the prompt + max_data_attempts_per_prompt: 2 # data budget; 1 retry separates transient from deterministic + backoff_base_s: 1.0 + max_backoff_s: 30.0 + # Prompts allowed to exhaust their data budget and be dropped. 0 fails the run + # on the first one, propagating the original error. The two budgets above are + # INDEPENDENT counters, so one prompt can consume up to their sum minus one. + max_skipped_prompts: 0 + # CONSECUTIVE prompts allowed to exhaust their infra budget and be dropped; any + # committed rollout resets the count. 0 fails the run on the first one. Raise it + # on a large fleet where losing the odd shard is expected and a whole run is + # expensive to restart -- a step short a few groups still trains, and the count + # only keeps climbing if nothing at all is coming back. + max_consecutive_dropped_prompts: 0 + # Smallest batch a step may train on, as a fraction of num_prompts_per_step. + # The budgets above are run-scoped and cannot bound how short one step gets, so + # this is what guarantees the gradient is computed from most of a batch. The floor + # is ceil(fraction * num_prompts_per_step), so small batches tolerate no drops. + min_step_batch_fraction: 0.9 + # What happens to the step a dropped prompt was stamped for (in_order sampler only; + # an unstamped sampler strands nothing). "shrink" trains the step on fewer groups, + # bounded by the floor above. "replace" substitutes a fresh prompt so the step keeps + # its configured batch size -- what the v1 stack does -- and still shrinks if the + # attempts or spares below run out. When a later step already has a finished group, + # replace borrows it to close the dropped step immediately and sends the fresh + # prompt to repay that step instead, so the trainer does not idle waiting on a new + # rollout; promoted_prompt_groups reports when that happened. + on_dropped_prompt: shrink # shrink | replace + max_replacement_attempts: 1 # fresh prompts tried per lost group + # Low-water mark for the spare-prompt pool. Refilled by diverting one whole batch + # from the dataloader before it is admitted, so a refill yields a batch of spares. + replacement_reserve_prompts: 1 + + # Deadlines. null disables, which is the default and reproduces the historical + # behaviour of waiting forever. Set them in any run that matters: without one a + # wedged generation engine parks a rollout permanently, holding a + # max_inflight_prompts slot for the rest of the job. + native: # AsyncRolloutImpl only + generation_timeout_s: null # one generate_async turn + env_timeout_s: null # one environment step + nemo_gym: # AsyncNemoGymRolloutImpl only + rollout_timeout_s: null # the whole prompt-group rollout, retries included + # Re-send just the rows that never arrived before retrying the whole group. + # Gym's stream dies on its first failing row, so one bad row loses every later + # one; recovering those individually beats redoing all N. + max_row_attempts: 3 + + # Last-resort stall detection. stall_timeout_s must exceed interval_s and every + # deadline above, so a merely-slow rollout is not reported as a stall. + stall_watchdog: + interval_s: 30.0 + stall_timeout_s: 600.0 + stall_action: warn # warn | abort + gym_subprocess_check: true # polls NeMo-Gym's RunHelper for dead servers + + # Per-shard liveness for the generation fleet. Off by default, in which case shard + # selection keeps its historical health-blind round-robin. On, a shard that fails + # probes or real requests is quarantined instead of keeping its 1/N of traffic, and + # the run aborts with an attributable error once too few shards remain. + generation_fleet_health: + enabled: false + probe_interval_s: 5.0 # must exceed probe_timeout_s + probe_timeout_s: 2.0 + unhealthy_threshold: 3 # ~15s to quarantine at the default interval + healthy_threshold: 2 # successes before a suspect shard is trusted again + selection: least_outstanding # steers away from a slow shard without diagnosing it + on_dead_shard: fail_fast # recovery modes arrive with the communicator rebuild + max_restart_attempts_per_shard: 5 + min_healthy_shards: 1 # below this the run aborts + + # NeMo-Gym-facing router in front of the generation fleet. Gym picks an endpoint by + # round-robin over a list fixed at process start and never fails over, so it is handed + # one NeMo-RL-owned URL instead and the routing decision moves next to generation_fleet_health. + # Only meaningful on the NeMo-Gym path; the native path calls generation over Ray. + generation_router: + enabled: false + port_range_low: 6000 # distinct from Gym (5000-5999) and vLLM (7000-8999) + port_range_high: 6099 + backend_timeout_s: 600.0 # covers a whole generation; Gym sets none at all + connect_timeout_s: 5.0 # a local handshake is ms-or-never + no_healthy_backend_status: 409 # must stay outside Gym's retry set + +checkpointing: + enabled: false + checkpoint_dir: results/ppo-single-controller + metric_name: null + +policy: + tokenizer: + # Opt-in Rust-backed BPE tokenizer (~10x faster encode); requires the + # fastokens-b10 wheel. NRL_USE_FASTOKENS overrides at runtime. + use_fastokens: false + # The only thing on the SC path that takes the policy optimizer off the GPU. + # Without it the critic's forward and train stages run on top of it. + offload_optimizer_for_logprob: true + + dtensor_cfg: + enabled: false + megatron_cfg: + enabled: true + + # Both are read on the SC path but absent from the PPO exemplar chain, which + # predates them. Values copied from grpo_math_1B.yaml. + draft: + enabled: false + model_name: null + loss_weight: 0.1 + num_layers: null + aux_layer_indices: null + + generation: + port_range_low: 11001 + port_range_high: 15000 + mcore_generation_config: + unified_memory_level: 0 + vllm_cfg: + async_engine: true + colocated: + enabled: false + resources: + gpus_per_node: 1 + num_nodes: 1 + +logger: + wandb: + name: ppo-single-controller-dev + swanlab: + name: ppo-single-controller-dev + mlflow: + run_name: ppo-single-controller-dev + +# TransferQueue-mediated data plane for sync GRPO. +# Off by default — the legacy grpo_train trainer never engages this. +# Flip enabled=true and run grpo_train_sync to use TQ-mediated bulk +# transfer between rollout and train. See nemo_rl/data_plane/README.md. +data_plane: + enabled: true + impl: transfer_queue + backend: "simple" # TQ storage backend ('simple' or 'mooncake_cpu') + storage_capacity: 1000000 # max samples retained per partition + num_storage_units: 2 # storage shards + claim_meta_poll_interval_s: 0.5 # blocking-claim poll cadence + global_segment_size: 549755813888 # 512 GiB — used when backend == "mooncake_cpu" + local_buffer_size: 68719476736 # 64 GiB — used when backend == "mooncake_cpu" + # observability: # NotRequired + # enabled: false + +cluster: + gpus_per_node: 2 + master_port_range_low: 25000 + master_port_range_high: 28000 diff --git a/tests/functional/L1_Functional_Tests_SingleController.sh b/tests/functional/L1_Functional_Tests_SingleController.sh index bad4dce95c1..3f1d3ee71db 100755 --- a/tests/functional/L1_Functional_Tests_SingleController.sh +++ b/tests/functional/L1_Functional_Tests_SingleController.sh @@ -81,6 +81,12 @@ run_test env VICTIM_STATE=serving uv run --no-sync bash ./tests/functional/grpo_ # Checkpoint save/restore (upstream #3429). run_test uv run --no-sync bash ./tests/functional/grpo_checkpoint_single_controller.sh +# PPO on SC: the critic forward/train stages, GAE returns reaching the data plane, +# critic-only warmup, ppo_epochs, and the value model's half of the checkpoint. None +# of the GRPO runs above touch any of it -- they build no value model at all -- so a +# regression in the PPO path is invisible to every other test in this file. +run_test fast uv run --no-sync bash ./tests/functional/ppo_async_single_controller.sh + cd ${PROJECT_ROOT}/tests if compgen -G ".coverage*" > /dev/null; then coverage combine .coverage* diff --git a/tests/functional/ppo_async_single_controller.sh b/tests/functional/ppo_async_single_controller.sh new file mode 100755 index 00000000000..5b1d21370e2 --- /dev/null +++ b/tests/functional/ppo_async_single_controller.sh @@ -0,0 +1,123 @@ +#!/bin/bash +# SingleController counterpart of tests/functional/ppo_async_megatron.sh: +# same model, batch shape, critic warmup and checkpoint/restore coverage, run +# through examples/run_grpo_single_controller.py instead of run_ppo.py. +# +# Two settings could not carry over, and each is forced by what SC is: +# - train_global_batch_size is 8, not 4: SC requires +# num_prompts_per_step * num_generations_per_prompt == the global batch, so +# one RL step maps to one optimizer step. +# - No validation: SC has no validation loop yet. +# +# Step config lives under ppo:, not grpo: -- the two blocks are mutually +# exclusive and the SC PPO exemplar nulls grpo out. + +set -euo pipefail + +SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd) +PROJECT_ROOT=$(realpath "${SCRIPT_DIR}/../..") +EXP_NAME=$(basename "$0" .sh) +EXP_DIR="${SCRIPT_DIR}/${EXP_NAME}" +CKPT_DIR="${EXP_DIR}/checkpoints" +export PYTHONPATH="${PROJECT_ROOT}:${PYTHONPATH:-}" + +rm -rf "${EXP_DIR}" +mkdir -p "${EXP_DIR}" + +TRAIN_CMD=( + uv run coverage run -a + --data-file="${PROJECT_ROOT}/tests/.coverage" + --source="${PROJECT_ROOT}/nemo_rl" + "${PROJECT_ROOT}/examples/run_grpo_single_controller.py" + --config "${PROJECT_ROOT}/examples/configs/ppo_math_1B_megatron_single_controller.yaml" + policy.model_name=Qwen/Qwen2.5-0.5B + value.model_name=Qwen/Qwen2.5-0.5B + ppo.num_prompts_per_step=2 + ppo.num_generations_per_prompt=4 + ppo.max_num_epochs=1000 + ppo.seq_logprob_error_threshold=1000 + ppo.policy_training_start_step=1 + ppo.ppo_epochs=2 + policy.train_global_batch_size=8 + policy.logprob_batch_size=4 + policy.train_micro_batch_size=1 + +policy.megatron_cfg.scheduler.override_opt_param_scheduler=true + policy.generation.colocated.enabled=false + policy.generation.colocated.resources.gpus_per_node=1 + policy.generation.colocated.resources.num_nodes=1 + policy.generation.vllm_cfg.async_engine=true + loss_fn.use_importance_sampling_correction=true + value.train_global_batch_size=8 + value.train_micro_batch_size=1 + +value.megatron_cfg.scheduler.override_opt_param_scheduler=true + data.use_multiple_dataloader=false + data_plane.enabled=true + data_plane.impl=transfer_queue + data_plane.backend=simple + async_rl.sampler.name=in_order + async_rl.sampler.max_lookahead_versions=1 + async_rl.sampler.warmup_lookahead_versions=2 + async_rl.min_groups_for_streaming_train=2 + async_rl.max_inflight_prompts=6 + async_rl.max_buffered_rollouts=6 + cluster.gpus_per_node=2 + logger.tensorboard_enabled=true + logger.wandb_enabled=false + logger.monitor_gpus=true + checkpointing.enabled=true + checkpointing.checkpoint_dir="${CKPT_DIR}" + checkpointing.metric_name=null + checkpointing.save_period=1 +) + +cd "${PROJECT_ROOT}" + +"${TRAIN_CMD[@]}" \ + ppo.max_num_steps=2 \ + logger.log_dir="${EXP_DIR}/logs_run1" \ + "$@" \ + 2>&1 | tee "${EXP_DIR}/run1.log" + +grep -q "Value init:" "${EXP_DIR}/run1.log" +grep -q "weight_sync=CollectiveWeightSynchronizer" "${EXP_DIR}/run1.log" +# policy_training_start_step=1, so step 0 trains the critic alone and step 1 is +# where the policy joins. The banner fires exactly on that transition. +test "$(grep -c "Critic warmup complete" "${EXP_DIR}/run1.log")" -eq 1 +test -f "${CKPT_DIR}/step_1/replay_buffer.pt" +test -f "${CKPT_DIR}/step_2/replay_buffer.pt" +test -d "${CKPT_DIR}/step_1/value/weights" + +"${TRAIN_CMD[@]}" \ + ppo.max_num_steps=4 \ + logger.log_dir="${EXP_DIR}/logs_run2" \ + "$@" \ + 2>&1 | tee "${EXP_DIR}/run2.log" + +grep -q "Restoring replay buffer from checkpoint" "${EXP_DIR}/run2.log" +grep -qF "replay group(s) from checkpoint" "${EXP_DIR}/run2.log" +# Warmup is behind us on the resumed run, so the policy trains every step and the +# transition never happens again. +assert_no_warmup=$(grep -c "Critic warmup complete" "${EXP_DIR}/run2.log" || true) +test "${assert_no_warmup}" -eq 0 +test -d "${CKPT_DIR}/step_4/policy/weights" +test -d "${CKPT_DIR}/step_4/value/weights" + +# run1 trains the policy on 1 of its 2 steps (step 0 is critic warmup); run2 +# resumes past the warmup and trains it on both. +for run_spec in "run1 1" "run2 2"; do + read -r run expected_policy_steps <<< "${run_spec}" + metrics="${EXP_DIR}/metrics_${run}.json" + uv run tests/json_dump_tb_logs.py "${EXP_DIR}/logs_${run}" \ + --output_path "${metrics}" + uv run tests/check_metrics.py "${metrics}" \ + 'len(data["train/reward"]) == 2' \ + "len(data[\"train/loss\"]) == ${expected_policy_steps}" \ + 'len(data["train/critic/loss"]) == 2' \ + 'min(data["train/probs_ratio_clamped_min"]) > 0.79' \ + 'max(data["train/probs_ratio_clamped_min"]) < 1.21' \ + 'min(data["train/probs_ratio_clamped_max"]) > 0.79' \ + 'max(data["train/probs_ratio_clamped_max"]) < 1.29' \ + 'max(data["train/token_mult_prob_error"]) < 1.05' \ + 'max(data["train/critic/loss"]) < 6.0' \ + 'min(data["train/critic/loss"]) >= 0' +done From 42f706211ec2a7eb11f449423ec3b8a52dc13124 Mon Sep 17 00:00:00 2001 From: Yuki Huang Date: Sun, 23 Aug 2026 09:56:06 -0700 Subject: [PATCH 15/31] test(sc): add an async PPO SingleController nightly recipe Signed-off-by: Yuki Huang --- ...po_math_1B_megatron_single_controller.yaml | 2 - ...-noncolocated-async-single-controller.yaml | 61 +++++++++++++++++++ ...ch-noncolocated-async-single-controller.sh | 54 ++++++++++++++++ tests/test_suites/nightly.txt | 1 + 4 files changed, 116 insertions(+), 2 deletions(-) create mode 100644 examples/configs/recipes/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.yaml create mode 100755 tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.sh diff --git a/examples/configs/ppo_math_1B_megatron_single_controller.yaml b/examples/configs/ppo_math_1B_megatron_single_controller.yaml index ecb6aca5d62..876e575135e 100644 --- a/examples/configs/ppo_math_1B_megatron_single_controller.yaml +++ b/examples/configs/ppo_math_1B_megatron_single_controller.yaml @@ -148,8 +148,6 @@ policy: megatron_cfg: enabled: true - # Both are read on the SC path but absent from the PPO exemplar chain, which - # predates them. Values copied from grpo_math_1B.yaml. draft: enabled: false model_name: null diff --git a/examples/configs/recipes/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.yaml b/examples/configs/recipes/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.yaml new file mode 100644 index 00000000000..72832a98473 --- /dev/null +++ b/examples/configs/recipes/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.yaml @@ -0,0 +1,61 @@ +# SingleController variant of ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async.yaml. +# Override some settings for testing complete features. +defaults: ./ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async.yaml + +logger: + log_dir: logs/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller + wandb: + name: ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller + +ppo: + # Null out the legacy async_ppo block: SC reads async_rl instead, and the + # entrypoint raises if this block is set. + async_ppo: null + # SC does not support validation yet. + val_period: 0 + val_at_start: false + # override from the inherited recipe. + num_prompts_per_step: 256 + max_num_epochs: 1000 + ppo_epochs: 4 + policy_training_start_step: 10 + +checkpointing: + enabled: false + checkpoint_dir: results/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller + # Without a validation loop there is no val:accuracy to select on, and SC + # rejects any metric_name outside the train: namespace. + metric_name: null + +policy: + # The only thing on the SC path that takes the policy optimizer off the GPU, + # which the critic needs before its forward and train stages. + offload_optimizer_for_logprob: true + +# TransferQueue data plane is mandatory for the SingleController path. The whole +# block is spelled out because the PPO config chain has none; the fields have no +# defaults. Values are grpo_math_1B.yaml's. +data_plane: + enabled: true + impl: transfer_queue + backend: "simple" + storage_capacity: 1000000 + num_storage_units: 2 + claim_meta_poll_interval_s: 0.5 + global_segment_size: 549755813888 + local_buffer_size: 68719476736 + +# SC async-RL runtime knobs, replacing the nulled ppo.async_ppo block. +async_rl: + sampler: + name: in_order + # Matches grpo.async_grpo.max_trajectory_age_steps=1. + max_lookahead_versions: 1 + # override from the inherited recipe. + warmup_lookahead_versions: 4 + # SC PPO assembles each step from a single chunk; the critic has no split + # train API, so this must equal num_prompts_per_step. + min_groups_for_streaming_train: ${ppo.num_prompts_per_step} + max_inflight_prompts: 1280 + # num_prompts_per_step * (max_lookahead_versions + 1). + max_buffered_rollouts: 1280 diff --git a/tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.sh b/tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.sh new file mode 100755 index 00000000000..a7f07c6c2cb --- /dev/null +++ b/tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.sh @@ -0,0 +1,54 @@ +#!/bin/bash +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source $SCRIPT_DIR/common.env + +# ===== BEGIN CONFIG ===== +NUM_NODES=2 +STEPS_PER_RUN=40 +MAX_STEPS=40 +NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) # Round up +NUM_MINUTES=60 +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +# Run the experiment +cd $PROJECT_ROOT +uv run examples/run_grpo_single_controller.py \ + --config $CONFIG_PATH \ + ppo.max_num_steps=$MAX_STEPS \ + logger.log_dir=$LOG_DIR \ + logger.wandb_enabled=True \ + logger.wandb.project=nemo-rl \ + logger.wandb.name=$EXP_NAME \ + logger.monitor_gpus=True \ + logger.tensorboard_enabled=True \ + checkpointing.enabled=True \ + checkpointing.checkpoint_dir=$CKPT_DIR \ + $@ \ + 2>&1 | tee $RUN_LOG + +# The critic came up, and the policy refits across clusters (non-colocated +# generation). The legacy script's "Separate PPO clusters initialized" and +# "Using GAE advantage estimator" are printed by ppo.py only. +grep -q "Value init:" $RUN_LOG +grep -q "weight_sync=CollectiveWeightSynchronizer" $RUN_LOG + +# Convert tensorboard logs to json +uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS + +# Only run metrics if the target step is reached. Thresholds are the legacy +# recipe's, at the step where this run has consumed the same data. SC has no +# validation loop, so its validation/accuracy assertion has no counterpart. +if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | map(tonumber) | max' $JSON_METRICS) -ge $MAX_STEPS ]]; then + uv run tests/check_metrics.py $JSON_METRICS \ + 'median(data["train/token_mult_prob_error"]) < 1.1' \ + 'data["train/token_mult_prob_error"]["160"] < 1.1' \ + 'median(data["train/max_seq_mult_prob_error"]) < 1.2' \ + 'len(data["train/critic/loss"]) == 160' \ + 'min(data["train/critic/loss"]) >= 0' \ + 'data["train/reward"]["160"] > 0.75' + + # Clean up checkpoint directory after successful run to save space. + rm -rf "$CKPT_DIR" +fi diff --git a/tests/test_suites/nightly.txt b/tests/test_suites/nightly.txt index 65e2d14ef37..894600d37d6 100644 --- a/tests/test_suites/nightly.txt +++ b/tests/test_suites/nightly.txt @@ -175,6 +175,7 @@ tests/test_suites/llm/prorlv2-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v2-tq_moo # Single Controller (SC) tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-async-1off-single-controller-streaming2.sh tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron-single-controller-sync.sh +tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.sh ######## # DAPO # From e95ce039496321be56920dfca3f09cfe04f7c7e5 Mon Sep 17 00:00:00 2001 From: Yuki Huang Date: Sun, 23 Aug 2026 09:59:07 -0700 Subject: [PATCH 16/31] pyrefly Signed-off-by: Yuki Huang --- pyrefly.toml | 20 ++++++++------------ 1 file changed, 8 insertions(+), 12 deletions(-) diff --git a/pyrefly.toml b/pyrefly.toml index de3ddb4018a..e3ba33a251b 100644 --- a/pyrefly.toml +++ b/pyrefly.toml @@ -61,10 +61,6 @@ project-includes = [ "nemo_rl/algorithms/opd.py", "nemo_rl/algorithms/reward_functions.py", "nemo_rl/algorithms/single_controller.py", - "nemo_rl/algorithms/single_controller_utils/__init__.py", - "nemo_rl/algorithms/single_controller_utils/config.py", - "nemo_rl/algorithms/single_controller_utils/setup.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", @@ -172,7 +168,6 @@ project-includes = [ "nemo_rl/models/dtensor/parallelize.py", "nemo_rl/models/generation/__init__.py", "nemo_rl/models/generation/constants.py", - "nemo_rl/models/generation/fleet_health.py", "nemo_rl/models/generation/dynamo/__init__.py", "nemo_rl/models/generation/dynamo/arguments.py", "nemo_rl/models/generation/dynamo/config.py", @@ -186,8 +181,9 @@ project-includes = [ "nemo_rl/models/generation/dynamo/validate_dynamo_vllm_args.py", "nemo_rl/models/generation/dynamo/venv.py", "nemo_rl/models/generation/dynamo/worker_pool.py", - "nemo_rl/models/generation/interfaces.py", + "nemo_rl/models/generation/fleet_health.py", "nemo_rl/models/generation/generation_router.py", + "nemo_rl/models/generation/interfaces.py", "nemo_rl/models/generation/megatron/__init__.py", "nemo_rl/models/generation/megatron/config.py", "nemo_rl/models/generation/megatron/utils.py", @@ -219,23 +215,24 @@ project-includes = [ "nemo_rl/models/generation/vllm/worker_utils.py", "nemo_rl/models/huggingface/__init__.py", "nemo_rl/models/megatron/__init__.py", - "nemo_rl/models/megatron/draft/__init__.py", "nemo_rl/models/megatron/memory_saver.py", "nemo_rl/models/policy/__init__.py", "nemo_rl/models/policy/interfaces.py", + "nemo_rl/models/policy/tq_policy.py", "nemo_rl/models/policy/utils.py", "nemo_rl/models/policy/workers/__init__.py", - "nemo_rl/models/policy/workers/megatron_remote_sparse_refit.py", "nemo_rl/models/policy/workers/checkpoint_engine.py", + "nemo_rl/models/policy/workers/megatron_remote_sparse_refit.py", "nemo_rl/models/policy/workers/patches.py", "nemo_rl/models/value/__init__.py", "nemo_rl/models/value/config.py", + "nemo_rl/models/value/tq_value.py", "nemo_rl/models/value/workers/__init__.py", "nemo_rl/utils/__init__.py", + "nemo_rl/utils/checkpoint.py", "nemo_rl/utils/checkpoint_engines/__init__.py", "nemo_rl/utils/checkpoint_engines/base.py", "nemo_rl/utils/checkpoint_engines/nixl.py", - "nemo_rl/utils/checkpoint.py", "nemo_rl/utils/config.py", "nemo_rl/utils/fastokens.py", "nemo_rl/utils/grad_norm.py", @@ -262,20 +259,19 @@ project-includes = [ "nemo_rl/weight_sync/interfaces.py", "nemo_rl/weight_sync/ipc_weight_synchronizer.py", "nemo_rl/weight_sync/megatron_weight_synchronizer.py", - "nemo_rl/weight_sync/vllm_remote_sparse_weight_synchronizer.py", "nemo_rl/weight_sync/nccl_reshard_utils.py", "nemo_rl/weight_sync/nccl_reshard_weight_synchronizer.py", + "nemo_rl/weight_sync/vllm_remote_sparse_weight_synchronizer.py", "nemo_rl/weight_sync/xferdtensor.py", "nemo_rl/weight_sync/xferdtensor_python.py", - "tools/external_gym_vllm/vllm_pool_lb.py", "tools/external_gym_vllm/serve_vllm_on_ray.py", + "tools/external_gym_vllm/vllm_pool_lb.py", "tools/model_diagnostics/1.max_model_len_respected.py", "tools/model_diagnostics/2.long_generation_decode_vs_prefill.py", "tools/model_diagnostics/3.check_and_reinit_hf_model_embeddings_untrained.py", "tools/model_diagnostics/4.vllm_precision_compilation_test.py", "tools/model_diagnostics/5.prefix_caching_nan.py", "tools/model_diagnostics/6.vllm_routed_experts_completeness.py", - "tools/refit_bandwidth_calculator.py", "tools/x_token/__init__.py", "tools/x_token/reapply_exact_map.py", "tools/x_token/sort_and_cut_projection_matrix.py", From 59b2845b20ddeebfbde017ee1b9c3f9f8525dac6 Mon Sep 17 00:00:00 2001 From: Yuki Huang Date: Sun, 23 Aug 2026 20:41:20 -0700 Subject: [PATCH 17/31] test(sc): regroup the SingleController L1 functional test list, raise the nightly GPU-hour cap to 4190 Signed-off-by: Yuki Huang --- .../L1_Functional_Tests_SingleController.sh | 56 +++---------------- tests/unit/test_recipes_and_test_suites.py | 6 +- 2 files changed, 11 insertions(+), 51 deletions(-) diff --git a/tests/functional/L1_Functional_Tests_SingleController.sh b/tests/functional/L1_Functional_Tests_SingleController.sh index 3f1d3ee71db..fee76a3699e 100755 --- a/tests/functional/L1_Functional_Tests_SingleController.sh +++ b/tests/functional/L1_Functional_Tests_SingleController.sh @@ -35,57 +35,17 @@ run_test() { } run_test fast uv run --no-sync bash ./tests/functional/grpo_dp_single_controller.sh +run_test uv run --no-sync bash ./tests/functional/grpo_checkpoint_single_controller.sh +run_test fast uv run --no-sync bash ./tests/functional/ppo_async_single_controller.sh +# nemogym test run_test fast uv run --no-sync bash ./tests/functional/grpo_async_gym_single_controller.sh -# Full mode only (~10 min): SIGKILLs a generation worker and asserts the job fails fast -# and attributably instead of wedging. This is the ONLY end-to-end check of the -# containment behaviour -- without it, a regression that restores the silent wedge is -# caught by nothing, because a wedged job produces no exception and no failing assertion -# anywhere else. -run_test uv run --no-sync bash ./tests/functional/grpo_dp_single_controller_chaos.sh -# Full mode only: the same Gym run, but with NeMo-Gym pointed at the NeMo-RL-owned router. -# Without this the router has no functional coverage at all -- the default Gym run above -# leaves it disabled, so a regression in the proxy would ship silently. -# -# gen_kl_error is the assertion that earns its keep here: it compares vLLM's logprobs -# against the trainer's recomputation, so a proxy that corrupts or truncates a response -# blows it up. A run that merely completes would not prove the payload survived the hop. -run_test uv run --no-sync bash ./tests/functional/grpo_async_gym_single_controller.sh \ +run_test uv run --no-sync bash ./tests/functional/grpo_async_gym_single_controller.sh \ ++async_rl.generation_router.enabled=true \ ++async_rl.generation_fleet_health.enabled=true - -# ...and the property that run CANNOT prove. It is dp_size=1, so _pick_backend has one -# choice, the serving set never shrinks, and the no-healthy-backend path never fires: it -# demonstrates pass-through, not failover. This one runs two generation shards, kills one -# mid-run, and asserts the serving set shrinks so NeMo-Gym stops being handed the corpse. -# -# EXPECT defaults to quarantine deliberately. Surviving the loss needs the communicator -# rebuild that lands later in this stack -- without it the next refit broadcasts to the -# dead rank and hangs in NCCL (job 6258553 sat there for 33 minutes). Asserting survival -# here would assert a property this part does not implement. -# -# Needs >= 3 GPUs (2 generation + 1 trainer) and self-skips below that, so it is inert on -# the 2-GPU L1 runners and only does its job on a larger box. -run_test uv run --no-sync bash ./tests/functional/grpo_sc_gym_router_failover.sh - -# grpo_dp_single_controller_chaos.sh again, this time killing a worker that is mid-rollout -# rather than between calls. Registered because pinning the victim state -- which is what -# makes that test reproducible at all -- would otherwise silently drop a scenario the old, -# non-deterministic selection used to hit by chance. The two fail by different routes: -# killing an idle worker leaves the loss to be *detected*, killing a serving one destroys -# an in-flight RPC that surfaces at once (222s vs 12s when measured). A regression in -# either is invisible to the other. -# -# Cheap to add: the serving path fails in seconds, so this is dominated by startup. -run_test env VICTIM_STATE=serving uv run --no-sync bash ./tests/functional/grpo_dp_single_controller_chaos.sh - -# Checkpoint save/restore (upstream #3429). -run_test uv run --no-sync bash ./tests/functional/grpo_checkpoint_single_controller.sh - -# PPO on SC: the critic forward/train stages, GAE returns reaching the data plane, -# critic-only warmup, ppo_epochs, and the value model's half of the checkpoint. None -# of the GRPO runs above touch any of it -- they build no value model at all -- so a -# regression in the PPO path is invisible to every other test in this file. -run_test fast uv run --no-sync bash ./tests/functional/ppo_async_single_controller.sh +# chaos test +run_test uv run --no-sync bash ./tests/functional/grpo_dp_single_controller_chaos.sh +run_test env VICTIM_STATE=serving uv run --no-sync bash ./tests/functional/grpo_dp_single_controller_chaos.sh +run_test uv run --no-sync bash ./tests/functional/grpo_sc_gym_router_failover.sh cd ${PROJECT_ROOT}/tests if compgen -G ".coverage*" > /dev/null; then diff --git a/tests/unit/test_recipes_and_test_suites.py b/tests/unit/test_recipes_and_test_suites.py index 9bf0d00e130..8739f7b135b 100644 --- a/tests/unit/test_recipes_and_test_suites.py +++ b/tests/unit/test_recipes_and_test_suites.py @@ -256,7 +256,7 @@ def test_all_recipe_yamls_accounted_for_in_test_suites( ) -def test_nightly_compute_stays_below_4139_hours(nightly_test_suite, tracker): +def test_nightly_compute_stays_below_4190_hours(nightly_test_suite, tracker): command = f"DRYRUN=1 HF_HOME=... HF_DATASETS_CACHE=... CONTAINER= ACCOUNT= PARTITION= ./tools/launch {' '.join(nightly_test_suite)}" print(f"Running command: {command}") @@ -288,8 +288,8 @@ def test_nightly_compute_stays_below_4139_hours(nightly_test_suite, tracker): f"Last line of output was not as expected: '{last_line}'" ) total_gpu_hours = float(last_line.split(":")[-1].strip()) - assert total_gpu_hours <= 4139, ( - f"Total GPU hours exceeded 4139: {last_line}. We should revisit the test suites to reduce the total GPU hours." + assert total_gpu_hours <= 4190, ( + f"Total GPU hours exceeded 4190: {last_line}. We should revisit the test suites to reduce the total GPU hours." ) tracker.track("total_nightly_gpu_hours", total_gpu_hours) From b9bd64c77b60600d1cc47047a7568b67bef5990e Mon Sep 17 00:00:00 2001 From: Yuki Huang Date: Mon, 24 Aug 2026 01:48:08 -0700 Subject: [PATCH 18/31] fix(sc): mark returns/values token-aligned, migrate the data_plane configs to backend-specific blocks, correct the async SC nightly metric step indices Signed-off-by: Yuki Huang --- ...po_math_1B_megatron_single_controller.yaml | 24 ++++++++++++------- ...-noncolocated-async-single-controller.yaml | 22 +++++++++++++---- nemo_rl/data_plane/column_io.py | 2 ++ .../functional/ppo_async_single_controller.sh | 3 --- ...ch-noncolocated-async-single-controller.sh | 6 ++--- 5 files changed, 38 insertions(+), 19 deletions(-) diff --git a/examples/configs/ppo_math_1B_megatron_single_controller.yaml b/examples/configs/ppo_math_1B_megatron_single_controller.yaml index 876e575135e..777903e4a79 100644 --- a/examples/configs/ppo_math_1B_megatron_single_controller.yaml +++ b/examples/configs/ppo_math_1B_megatron_single_controller.yaml @@ -176,19 +176,27 @@ logger: mlflow: run_name: ppo-single-controller-dev -# TransferQueue-mediated data plane for sync GRPO. -# Off by default — the legacy grpo_train trainer never engages this. -# Flip enabled=true and run grpo_train_sync to use TQ-mediated bulk -# transfer between rollout and train. See nemo_rl/data_plane/README.md. +# TransferQueue data plane is mandatory for the SingleController path. The whole +# block is spelled out because the PPO config chain has none; the fields have no +# defaults. Values are grpo_math_1B.yaml's except defualt enabled. +# See nemo_rl/data_plane/README.md for more details. data_plane: enabled: true impl: transfer_queue backend: "simple" # TQ storage backend ('simple' or 'mooncake_cpu') - storage_capacity: 1000000 # max samples retained per partition - num_storage_units: 2 # storage shards claim_meta_poll_interval_s: 0.5 # blocking-claim poll cadence - global_segment_size: 549755813888 # 512 GiB — used when backend == "mooncake_cpu" - local_buffer_size: 68719476736 # 64 GiB — used when backend == "mooncake_cpu" + # Backend-specific blocks: only the one named by `backend` is read, and an + # absent block means that backend's defaults (see SimpleStorageConfig / + # MooncakeCpuConfig in nemo_rl/data_plane/interfaces.py). + simple: + storage_capacity: 1000000 # max samples retained per partition + num_storage_units: ${mul:2, ${cluster.num_nodes}} # TQ wants >= 2 per node + mooncake_cpu: + # Per client process — see MooncakeCpuConfig before raising these. + global_segment_size: 68719476736 # 64 GiB/process + local_buffer_size: 4294967296 # 4 GiB/process + reuse_registered_buffers: true # reuse RDMA-registered buffers + staging_buffer_size: 268435456 # 256 MiB/pool slot; bigger transfers bypass the pool # observability: # NotRequired # enabled: false diff --git a/examples/configs/recipes/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.yaml b/examples/configs/recipes/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.yaml index 72832a98473..e5d20a3d545 100644 --- a/examples/configs/recipes/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.yaml +++ b/examples/configs/recipes/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.yaml @@ -32,18 +32,29 @@ policy: # which the critic needs before its forward and train stages. offload_optimizer_for_logprob: true + draft: + enabled: false + model_name: null + loss_weight: 0.1 + num_layers: null + aux_layer_indices: null + # TransferQueue data plane is mandatory for the SingleController path. The whole # block is spelled out because the PPO config chain has none; the fields have no -# defaults. Values are grpo_math_1B.yaml's. +# defaults. Values are grpo_math_1B.yaml's except defualt enabled. data_plane: enabled: true impl: transfer_queue backend: "simple" - storage_capacity: 1000000 - num_storage_units: 2 claim_meta_poll_interval_s: 0.5 - global_segment_size: 549755813888 - local_buffer_size: 68719476736 + simple: + storage_capacity: 1000000 + num_storage_units: ${mul:2, ${cluster.num_nodes}} + mooncake_cpu: + global_segment_size: 68719476736 + local_buffer_size: 4294967296 + reuse_registered_buffers: true + staging_buffer_size: 268435456 # SC async-RL runtime knobs, replacing the nulled ppo.async_ppo block. async_rl: @@ -53,6 +64,7 @@ async_rl: max_lookahead_versions: 1 # override from the inherited recipe. warmup_lookahead_versions: 4 + recompute_kv_cache_after_weight_updates: true # SC PPO assembles each step from a single chunk; the critic has no split # train API, so this must equal num_prompts_per_step. min_groups_for_streaming_train: ${ppo.num_prompts_per_step} diff --git a/nemo_rl/data_plane/column_io.py b/nemo_rl/data_plane/column_io.py index e8196866b50..5242f282263 100644 --- a/nemo_rl/data_plane/column_io.py +++ b/nemo_rl/data_plane/column_io.py @@ -47,6 +47,8 @@ "prev_logprobs", "reference_policy_logprobs", "advantages", + "returns", + "values", "token_mask", "sample_mask", "routed_experts", diff --git a/tests/functional/ppo_async_single_controller.sh b/tests/functional/ppo_async_single_controller.sh index 5b1d21370e2..aa3cb1fd637 100755 --- a/tests/functional/ppo_async_single_controller.sh +++ b/tests/functional/ppo_async_single_controller.sh @@ -51,9 +51,6 @@ TRAIN_CMD=( value.train_micro_batch_size=1 +value.megatron_cfg.scheduler.override_opt_param_scheduler=true data.use_multiple_dataloader=false - data_plane.enabled=true - data_plane.impl=transfer_queue - data_plane.backend=simple async_rl.sampler.name=in_order async_rl.sampler.max_lookahead_versions=1 async_rl.sampler.warmup_lookahead_versions=2 diff --git a/tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.sh b/tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.sh index a7f07c6c2cb..aaace3a86cf 100755 --- a/tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.sh +++ b/tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.sh @@ -43,11 +43,11 @@ uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | map(tonumber) | max' $JSON_METRICS) -ge $MAX_STEPS ]]; then uv run tests/check_metrics.py $JSON_METRICS \ 'median(data["train/token_mult_prob_error"]) < 1.1' \ - 'data["train/token_mult_prob_error"]["160"] < 1.1' \ + 'data["train/token_mult_prob_error"]["40"] < 1.1' \ 'median(data["train/max_seq_mult_prob_error"]) < 1.2' \ - 'len(data["train/critic/loss"]) == 160' \ + 'len(data["train/critic/loss"]) == 40' \ 'min(data["train/critic/loss"]) >= 0' \ - 'data["train/reward"]["160"] > 0.75' + 'data["train/reward"]["40"] > 0.75' # Clean up checkpoint directory after successful run to save space. rm -rf "$CKPT_DIR" From 4025501eba8eca98e93cd1149e62b4a62eee7815 Mon Sep 17 00:00:00 2001 From: Yuki Huang Date: Mon, 24 Aug 2026 07:15:29 -0700 Subject: [PATCH 19/31] fix(sc): reject the shaping knobs the SingleController path does not implement, stop max_num_epochs<=0 from zeroing max_num_steps, correct the PPO recipe's buffer-size comment Signed-off-by: Yuki Huang --- ...po_math_1B_megatron_single_controller.yaml | 2 +- ...-noncolocated-async-single-controller.yaml | 7 ++-- .../single_controller_utils/config.py | 27 ++++++++++++-- .../single_controller_utils/setup.py | 2 +- .../unit/single_controller/test_ppo_setup.py | 35 ++++++++++++++++++- 5 files changed, 66 insertions(+), 7 deletions(-) diff --git a/examples/configs/ppo_math_1B_megatron_single_controller.yaml b/examples/configs/ppo_math_1B_megatron_single_controller.yaml index 777903e4a79..b63a64562db 100644 --- a/examples/configs/ppo_math_1B_megatron_single_controller.yaml +++ b/examples/configs/ppo_math_1B_megatron_single_controller.yaml @@ -178,7 +178,7 @@ logger: # TransferQueue data plane is mandatory for the SingleController path. The whole # block is spelled out because the PPO config chain has none; the fields have no -# defaults. Values are grpo_math_1B.yaml's except defualt enabled. +# defaults. Values are grpo_math_1B.yaml's except default enabled. # See nemo_rl/data_plane/README.md for more details. data_plane: enabled: true diff --git a/examples/configs/recipes/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.yaml b/examples/configs/recipes/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.yaml index e5d20a3d545..bbddb1beccf 100644 --- a/examples/configs/recipes/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.yaml +++ b/examples/configs/recipes/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.yaml @@ -14,6 +14,9 @@ ppo: # SC does not support validation yet. val_period: 0 val_at_start: false + # SC does not implement it, so leaving it enabled would describe filtering + # this run does not do. + overlong_filtering: false # override from the inherited recipe. num_prompts_per_step: 256 max_num_epochs: 1000 @@ -41,7 +44,7 @@ policy: # TransferQueue data plane is mandatory for the SingleController path. The whole # block is spelled out because the PPO config chain has none; the fields have no -# defaults. Values are grpo_math_1B.yaml's except defualt enabled. +# defaults. Values are grpo_math_1B.yaml's except default enabled. data_plane: enabled: true impl: transfer_queue @@ -69,5 +72,5 @@ async_rl: # train API, so this must equal num_prompts_per_step. min_groups_for_streaming_train: ${ppo.num_prompts_per_step} max_inflight_prompts: 1280 - # num_prompts_per_step * (max_lookahead_versions + 1). + # num_prompts_per_step * (warmup_lookahead_versions + 1) -- sized for the peak window. max_buffered_rollouts: 1280 diff --git a/nemo_rl/algorithms/single_controller_utils/config.py b/nemo_rl/algorithms/single_controller_utils/config.py index b1cbf4b2099..3e1f63ade3b 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/nemo_rl/algorithms/single_controller_utils/config.py @@ -708,8 +708,9 @@ def _validate_failure_settings( def _validate_algo_settings(master_config: MasterConfig) -> None: """Reject algorithm blocks the SingleController path cannot honour. - Both directions: a critic the PPO path needs and does not have, and a critic - a GRPO run carries and would never build. + Both directions on the critic: one the PPO path needs and does not have, and + one a GRPO run carries and would never build. Plus the reward shaping and + filtering knobs SC reads on neither path. """ grpo = getattr(master_config, "grpo", None) ppo = getattr(master_config, "ppo", None) @@ -721,6 +722,28 @@ def _validate_algo_settings(master_config: MasterConfig) -> None: ) algo_cfg = algo_config(master_config) + + # SC reads none of these on either path, so an enabled one describes shaping + # this run does not do. Async GRPO rejects three of them the same way. + unsupported = [ + name + for name, enabled in ( + ("overlong_filtering", algo_cfg.overlong_filtering), + ("use_dynamic_sampling", algo_cfg.use_dynamic_sampling), + ("reward_scaling", algo_cfg.reward_scaling.enabled), + ("reward_shaping", algo_cfg.reward_shaping.enabled), + ) + if enabled + ] + if unsupported: + prefix = "ppo" if is_ppo_run(master_config) else "grpo" + names = ", ".join(f"{prefix}.{name}" for name in unsupported) + raise NotImplementedError( + f"{names} not supported on the SingleController path, which " + "implements none of them -- the run would silently skip the " + "shaping. Disable them." + ) + if not is_ppo_run(master_config): # A value block without `ppo` is inert -- nothing builds the critic -- # and a config carrying one is asking for PPO by every reading except diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py index 723c80bad4c..7f88866990e 100644 --- a/nemo_rl/algorithms/single_controller_utils/setup.py +++ b/nemo_rl/algorithms/single_controller_utils/setup.py @@ -421,7 +421,7 @@ def _clamp_max_num_steps( """Clamp grpo.max_num_steps to max_num_epochs * len(dataloader).""" algo_cfg = algo_config(master_config) max_num_epochs = algo_cfg.max_num_epochs - if max_num_epochs is None: + if max_num_epochs is None or max_num_epochs <= 0: return algo_cfg.max_num_steps = min( algo_cfg.max_num_steps, diff --git a/tests/unit/single_controller/test_ppo_setup.py b/tests/unit/single_controller/test_ppo_setup.py index c360062edf4..e816335ded1 100644 --- a/tests/unit/single_controller/test_ppo_setup.py +++ b/tests/unit/single_controller/test_ppo_setup.py @@ -246,7 +246,7 @@ def test_rejects_a_ppo_epoch_count_below_one(self): with pytest.raises(ValueError, match="ppo_epochs must be at least 1"): validate_single_controller_config(mc) - def test_rejects_non_gae_estimator(self): + def test_the_ppo_schema_rejects_a_non_ppo_estimator(self): with pytest.raises(ValueError): PPOConfig(adv_estimator={"name": "grpo"}) @@ -293,6 +293,39 @@ def test_accepts_varying_ckpt_structure_with_warmup(self): self._warmup_ckpt_config(constant_structure=False) ) + @pytest.mark.parametrize( + "enable", + [ + lambda cfg: setattr(cfg, "overlong_filtering", True), + lambda cfg: setattr(cfg, "use_dynamic_sampling", True), + lambda cfg: setattr(cfg.reward_scaling, "enabled", True), + lambda cfg: setattr(cfg.reward_shaping, "enabled", True), + ], + ids=[ + "overlong_filtering", + "use_dynamic_sampling", + "reward_scaling", + "reward_shaping", + ], + ) + def test_rejects_shaping_the_sc_path_does_not_implement(self, enable): + mc = _ppo_master_config() + enable(mc.ppo) + + with pytest.raises( + NotImplementedError, match="not supported on the SingleController" + ): + validate_single_controller_config(mc) + + def test_rejects_shaping_on_a_grpo_run_too(self): + mc = _make_master_config() + mc.grpo.overlong_filtering = True + + with pytest.raises( + NotImplementedError, match=r"grpo\.overlong_filtering" + ): + validate_single_controller_config(mc) + def test_grpo_is_free_to_use_any_sampler(self): mc = _make_master_config() mc.async_rl.sampler = WindowedSamplerConfig() From 0ef34c8d02042aedc01d04dcd286c8ebbf1f65a6 Mon Sep 17 00:00:00 2001 From: Yuki Huang Date: Mon, 24 Aug 2026 07:26:03 -0700 Subject: [PATCH 20/31] docs(sc): document PPO on the SingleController path, correct the stale checkpointing and reward-shaping feature gaps Signed-off-by: Yuki Huang --- docs/guides/single-controller.md | 26 ++++++++++++++++---------- docs/index.md | 4 ++-- 2 files changed, 18 insertions(+), 12 deletions(-) diff --git a/docs/guides/single-controller.md b/docs/guides/single-controller.md index 703859d2c8c..209fdacda26 100644 --- a/docs/guides/single-controller.md +++ b/docs/guides/single-controller.md @@ -1,10 +1,10 @@ -# Train with Single-Controller (Async GRPO) +# Train with Single-Controller (Async GRPO and PPO) :::{warning} The Single-Controller path is a **beta feature** and still under active development. The API and configuration surface are not yet stable and may change without notice. Issues and feedback are welcome — please file them at [github.com/NVIDIA-NeMo/RL/issues](https://github.com/NVIDIA-NeMo/RL/issues). ::: -The Single-Controller (SC) path is an alternative async GRPO runtime that runs rollout generation and policy training as two independent *pumps* coordinated by a single Ray actor (`SingleControllerActor`) sitting over a shared TransferQueue (TQ) data plane. Compared to the legacy async GRPO in [async-grpo.md](./async-grpo.md), SC decouples per-prompt rollouts from the per-step batch boundary: producers push finished rollouts into `TQReplayBuffer` at group granularity, and a pluggable `StalenessSampler` decides which groups the trainer consumes on each step. +The Single-Controller (SC) path is an alternative async GRPO and PPO runtime that runs rollout generation and policy training as two independent *pumps* coordinated by a single Ray actor (`SingleControllerActor`) sitting over a shared TransferQueue (TQ) data plane. Compared to the legacy async GRPO in [async-grpo.md](./async-grpo.md), SC decouples per-prompt rollouts from the per-step batch boundary: producers push finished rollouts into `TQReplayBuffer` at group granularity, and a pluggable `StalenessSampler` decides which groups the trainer consumes on each step. ## Configure the Single-Controller Path @@ -14,7 +14,7 @@ The SC path is launched via a dedicated entrypoint: uv run examples/run_grpo_single_controller.py --config ``` -`run_grpo_single_controller.py` mirrors `run_grpo.py` for config loading — the same YAML files apply — but requires a few settings the legacy path does not. The default exemplar lives at [examples/configs/grpo_math_1B_megatron_single_controller.yaml](../../examples/configs/grpo_math_1B_megatron_single_controller.yaml). +`run_grpo_single_controller.py` mirrors `run_grpo.py` for config loading — the same YAML files apply — but requires a few settings the legacy path does not. The default exemplar lives at [examples/configs/grpo_math_1B_megatron_single_controller.yaml](../../examples/configs/grpo_math_1B_megatron_single_controller.yaml); the PPO one at [examples/configs/ppo_math_1B_megatron_single_controller.yaml](../../examples/configs/ppo_math_1B_megatron_single_controller.yaml). ### Mandatory settings @@ -40,10 +40,11 @@ uv run examples/run_grpo_single_controller.py --config gpus_per_node: 4 # inference GPUs; remainder go to training ``` -3. **One RL step = one optimizer step.** The SC train pump does not support multi-mini-step inside a single RL step (see `validate_single_controller_config` in [nemo_rl/algorithms/single_controller_utils/config.py](../../nemo_rl/algorithms/single_controller_utils/config.py)): +3. **One RL step = one training batch.** The batch a step trains on is the whole step (see `validate_single_controller_config` in [nemo_rl/algorithms/single_controller_utils/config.py](../../nemo_rl/algorithms/single_controller_utils/config.py)). A GRPO step is also one optimizer step; a PPO step is `ppo.ppo_epochs` of them over that same batch. ```python num_prompts_per_step * num_generations_per_prompt == policy.train_global_batch_size + num_prompts_per_step * num_generations_per_prompt == value.train_global_batch_size # PPO ``` 4. **Enable importance sampling correction** whenever the sampler admits off-policy data (any `max_staleness_versions > 0` on the `windowed`/`weight_fifo` samplers, or `max_lookahead_versions > 0` on `in_order`). The correction and its derivation are the same as for legacy async GRPO — see [Why Importance Sampling Correction Is Required for Async](./async-grpo.md#why-importance-sampling-correction-is-required-for-async): @@ -53,6 +54,8 @@ uv run examples/run_grpo_single_controller.py --config use_importance_sampling_correction: true ``` +5. **(PPO) Set `ppo:` instead of `grpo:`** — the two algorithm blocks are mutually exclusive, and SC reads every step setting from whichever one is present. A PPO run also needs `value:`, `value_loss_fn:` and `ppo.adv_estimator.name: gae` (same schemas as legacy PPO), a Megatron critic, and `policy.offload_optimizer_for_logprob: true`, which is what keeps the policy optimizer off the GPU while the critic runs. `ppo.policy_training_start_step: N` gives the usual critic warmup: for the first N steps the policy is neither trained nor refit, while the critic trains every step. + ## Async-RL Knobs and Sampler Modes All SC async-RL runtime knobs live under `async_rl:` in the master config. The most important choice is the `sampler`, which sets the staleness policy shared by the rollout pump (how far it may run ahead) and the train pump (which groups it may consume). @@ -68,7 +71,7 @@ Pick one of four modes with `sampler.name`. Each mode takes its own knobs, liste | `sampler.name` | Rollout gating | Train selection | Typical use | | -------------- | ----------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------ | -| `in_order` | Dispatch may lead the trainer by up to `max_lookahead_versions` batches. Each dispatch is stamped with a `target_step`. | Consume the group whose `target_step == current_train_weight`. | Sync mode (`max_lookahead_versions=0`) and legacy-async exact-batch semantics (`max_lookahead_versions>=1`). | +| `in_order` | Dispatch may lead the trainer by up to `max_lookahead_versions` batches. Each dispatch is stamped with a `target_step`. | Consume the group whose `target_step == current_train_weight`. | Sync mode (`max_lookahead_versions=0`) and legacy-async exact-batch semantics (`max_lookahead_versions>=1`). The only mode supported on a PPO run. | | `weight_fifo` | Same gate as `in_order` (`max_staleness_versions` of lookahead). | Drain the oldest in-window `start_weight` first, waiting for that weight's batch to fill. | Strict weight-version FIFO under a bounded lookahead. | | `windowed` | Ungated — rollout keeps producing until the buffer fills. | Take any ready group with `start_weight` in `[train - max_staleness_versions, train]`, optionally freshest-first. | Over-sampled streaming; aged groups outside the window are evicted (wasted compute). | | `custom` | Determined by the imported class. | Determined by the imported class. | `target: "module:ClassName"` — bring your own `PromptGroupSampler`. | @@ -89,8 +92,9 @@ The shipped exemplars cover three of the four modes: Field definitions: -- `max_buffered_rollouts` — hard cap on unconsumed rollout groups buffered in the data plane. Validated at setup against the gated sampler's required capacity; a value too small deadlocks the rollout pump, so setup raises instead of silently blocking. -- `min_groups_for_streaming_train` — minimum ready groups the trainer waits for before dispatching a batch. Set to `num_prompts_per_step` for sync/legacy semantics; lower for streaming. +- `max_buffered_rollouts` — hard cap on unconsumed rollout groups buffered in the data plane. Validated at setup against the gated sampler's required capacity; a value too small deadlocks the rollout pump, so setup raises instead of silently blocking. Sized from the widest window the run ever uses, so `warmup_lookahead_versions` rather than `max_lookahead_versions` when it is set. +- `min_groups_for_streaming_train` — minimum ready groups the trainer waits for before dispatching a batch. Set to `num_prompts_per_step` for sync/legacy semantics; lower for streaming. (PPO) Must equal `num_prompts_per_step` — the critic has no split train API, so one `train_from_meta` call is one optimizer step, and streaming a step across chunks would step the critic once per chunk. +- `sampler.warmup_lookahead_versions` (PPO) — lookahead used while `ppo.policy_training_start_step` critic warmup is in progress, shrinking back to `max_lookahead_versions` afterwards. The SC equivalent of `ppo.async_ppo.warmup_generation_lead_steps`. ## Implementation Structure @@ -154,12 +158,13 @@ The [legacy async GRPO](./async-grpo.md) (`grpo.async_grpo.enabled: true` under ### Migrating a legacy async config -SC reads its async knobs from `async_rl:` and **requires `grpo.async_grpo: null`** — `run_grpo_single_controller.py` raises if a legacy block is still present, so null it out when porting rather than leaving it in place. +SC reads its async knobs from `async_rl:` and **requires `grpo.async_grpo: null`** (or `ppo.async_ppo: null` on a PPO run) — `run_grpo_single_controller.py` raises if a legacy block is still present, so null it out when porting rather than leaving it in place. -| Legacy `grpo.async_grpo.*` | SC equivalent `async_rl.*` | +| Legacy `grpo.async_grpo.*` / `ppo.async_ppo.*` | SC equivalent `async_rl.*` | | -------------------------- | -------------------------- | | `enabled: true` | Implicit — SC is always async; use `sampler.max_lookahead_versions: 0` for sync semantics, `>= 1` for async | | `max_trajectory_age_steps: N` | `sampler.name: in_order` with `sampler.max_lookahead_versions: N` | +| `warmup_generation_lead_steps` (PPO) | `sampler.warmup_lookahead_versions` — the lookahead to use while critic warmup is in progress | | `recompute_kv_cache_after_weight_updates` | `recompute_kv_cache_after_weight_updates` (same) | | `in_flight_weight_updates` | Always effectively true; `false`-equivalent behavior is not yet supported (drain-gate tracked in [issue #2625](https://github.com/NVIDIA-NeMo/RL/issues/2625)) | | *(no legacy equivalent — matches legacy full-batch train semantics)* | `min_groups_for_streaming_train: ${grpo.num_prompts_per_step}` | @@ -172,6 +177,7 @@ The SC path is still under active development. Feature gaps are tracked in [issu - Train backend: only Megatron is supported and validated; the AutoModel training path has not been tested on SC. - Generation backend: only vLLM is supported and validated; Megatron generation, SGLang, and TRT-LLM have not been tested on SC. -- Checkpointing and validation are not yet supported (setup raises if enabled). +- Validation is not yet supported (setup raises on `val_period > 0`, `val_at_start`, or `val_at_end`). +- Reward shaping and sample filtering — `overlong_filtering`, `reward_shaping`, `reward_scaling`, and `use_dynamic_sampling` are implemented on neither algorithm block, so setup rejects them rather than silently skipping the shaping. - The `windowed` sampler has no `over_sampling_ratio` cap — over-produced groups aged past the window are evicted, wasting rollout compute. - The drain gate in refit is not yet supported. diff --git a/docs/index.md b/docs/index.md index f17796dedab..d32943ff391 100644 --- a/docs/index.md +++ b/docs/index.md @@ -191,11 +191,11 @@ Choose among colocated IPC, NCCL, sparse delta, and NIXL refit transports. Use NIXL checkpoint-engine refit to update non-colocated vLLM generation workers from policy workers. ::: -:::{grid-item-card} {octicon}`workflow` Single-Controller (Async GRPO) +:::{grid-item-card} {octicon}`workflow` Single-Controller (Async GRPO and PPO) :link: guides/single-controller :link-type: doc -Run async GRPO via the SingleController path: TransferQueue data plane, pluggable staleness samplers, and streaming trainer. +Run async GRPO or PPO via the SingleController path: TransferQueue data plane, pluggable staleness samplers, and streaming trainer. ::: :::: From aa0664319302033719e6c3a84e354d635292c504 Mon Sep 17 00:00:00 2001 From: Yuki Huang Date: Mon, 24 Aug 2026 07:29:27 -0700 Subject: [PATCH 21/31] fix(sc): re-enable the DDP forward pre-hook in finish_train_step Signed-off-by: Yuki Huang --- .../models/policy/workers/megatron_policy_worker.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index 2377c3c84e9..2df57f7cf8f 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -1665,6 +1665,16 @@ def _finish_train_step_body(self, state: dict[str, Any]) -> dict[str, Any]: num_zeros_in_grad, mp_group=pg_collection.mp ) + # Mirrors train(): without re-enabling the pre-hook __init__ removed, the + # param all-gather never runs and each forward sees only its own shard. + if self._first_train_step_forward_pre_hook_disabled and update_successful: + self.enable_forward_pre_hook() + get_model_config( + self.model + ).param_sync_func = self._first_train_step_param_sync_func + self._first_train_step_param_sync_func = None + self._first_train_step_forward_pre_hook_disabled = False + if self.cfg["megatron_cfg"]["empty_unused_memory_level"] >= 2: torch.cuda.empty_cache() From 0b59fe4fb75a2ba327c5441aeb33c09e07ee0882 Mon Sep 17 00:00:00 2001 From: Yuki Huang Date: Mon, 24 Aug 2026 07:37:46 -0700 Subject: [PATCH 22/31] chore(sc): restore the L1 functional test rationale, put the SC-utils modules back in pyrefly project-includes, add the value stages to the data-flow diagram Signed-off-by: Yuki Huang --- nemo_rl/algorithms/single_controller.py | 4 ++ pyrefly.toml | 6 +++ .../L1_Functional_Tests_SingleController.sh | 45 ++++++++++++++++--- 3 files changed, 50 insertions(+), 5 deletions(-) diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 0de75a38416..9926124d0ec 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -22,9 +22,13 @@ _rollout_pump → gen.generate_and_push(prompt, dp_client) ← RPC to GenWorker GenWorker → dp_client.put_samples(...) _train_pump → sampler.evict/select against TQReplayBuffer + → _value_stage(meta) (PPO only) → value.get_values_from_meta(...) + Value → dp_client.get/put_samples(...) (via its own client) → _advantage_stage(meta) → dp_client.get_samples(...) → adv_estimator.compute_advantage(...) → dp_client.put_samples(...) + → _value_train(meta) (PPO only) → value.train_from_meta(...) + Value → dp_client.get_samples(...) (via its own client) → 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) diff --git a/pyrefly.toml b/pyrefly.toml index e3ba33a251b..7a1691b0877 100644 --- a/pyrefly.toml +++ b/pyrefly.toml @@ -61,6 +61,10 @@ project-includes = [ "nemo_rl/algorithms/opd.py", "nemo_rl/algorithms/reward_functions.py", "nemo_rl/algorithms/single_controller.py", + "nemo_rl/algorithms/single_controller_utils/__init__.py", + "nemo_rl/algorithms/single_controller_utils/config.py", + "nemo_rl/algorithms/single_controller_utils/setup.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", @@ -215,6 +219,7 @@ project-includes = [ "nemo_rl/models/generation/vllm/worker_utils.py", "nemo_rl/models/huggingface/__init__.py", "nemo_rl/models/megatron/__init__.py", + "nemo_rl/models/megatron/draft/__init__.py", "nemo_rl/models/megatron/memory_saver.py", "nemo_rl/models/policy/__init__.py", "nemo_rl/models/policy/interfaces.py", @@ -272,6 +277,7 @@ project-includes = [ "tools/model_diagnostics/4.vllm_precision_compilation_test.py", "tools/model_diagnostics/5.prefix_caching_nan.py", "tools/model_diagnostics/6.vllm_routed_experts_completeness.py", + "tools/refit_bandwidth_calculator.py", "tools/x_token/__init__.py", "tools/x_token/reapply_exact_map.py", "tools/x_token/sort_and_cut_projection_matrix.py", diff --git a/tests/functional/L1_Functional_Tests_SingleController.sh b/tests/functional/L1_Functional_Tests_SingleController.sh index fee76a3699e..7872aaa64f6 100755 --- a/tests/functional/L1_Functional_Tests_SingleController.sh +++ b/tests/functional/L1_Functional_Tests_SingleController.sh @@ -35,18 +35,53 @@ run_test() { } run_test fast uv run --no-sync bash ./tests/functional/grpo_dp_single_controller.sh -run_test uv run --no-sync bash ./tests/functional/grpo_checkpoint_single_controller.sh run_test fast uv run --no-sync bash ./tests/functional/ppo_async_single_controller.sh -# nemogym test run_test fast uv run --no-sync bash ./tests/functional/grpo_async_gym_single_controller.sh +# Full mode only (~10 min): SIGKILLs a generation worker and asserts the job fails fast +# and attributably instead of wedging. This is the ONLY end-to-end check of the +# containment behaviour -- without it, a regression that restores the silent wedge is +# caught by nothing, because a wedged job produces no exception and no failing assertion +# anywhere else. +run_test uv run --no-sync bash ./tests/functional/grpo_dp_single_controller_chaos.sh +# Full mode only: the same Gym run, but with NeMo-Gym pointed at the NeMo-RL-owned router. +# Without this the router has no functional coverage at all -- the default Gym run above +# leaves it disabled, so a regression in the proxy would ship silently. +# +# gen_kl_error is the assertion that earns its keep here: it compares vLLM's logprobs +# against the trainer's recomputation, so a proxy that corrupts or truncates a response +# blows it up. A run that merely completes would not prove the payload survived the hop. run_test uv run --no-sync bash ./tests/functional/grpo_async_gym_single_controller.sh \ ++async_rl.generation_router.enabled=true \ ++async_rl.generation_fleet_health.enabled=true -# chaos test -run_test uv run --no-sync bash ./tests/functional/grpo_dp_single_controller_chaos.sh -run_test env VICTIM_STATE=serving uv run --no-sync bash ./tests/functional/grpo_dp_single_controller_chaos.sh + +# ...and the property that run CANNOT prove. It is dp_size=1, so _pick_backend has one +# choice, the serving set never shrinks, and the no-healthy-backend path never fires: it +# demonstrates pass-through, not failover. This one runs two generation shards, kills one +# mid-run, and asserts the serving set shrinks so NeMo-Gym stops being handed the corpse. +# +# EXPECT defaults to quarantine deliberately. Surviving the loss needs the communicator +# rebuild that lands later in this stack -- without it the next refit broadcasts to the +# dead rank and hangs in NCCL (job 6258553 sat there for 33 minutes). Asserting survival +# here would assert a property this part does not implement. +# +# Needs >= 3 GPUs (2 generation + 1 trainer) and self-skips below that, so it is inert on +# the 2-GPU L1 runners and only does its job on a larger box. run_test uv run --no-sync bash ./tests/functional/grpo_sc_gym_router_failover.sh +# grpo_dp_single_controller_chaos.sh again, this time killing a worker that is mid-rollout +# rather than between calls. Registered because pinning the victim state -- which is what +# makes that test reproducible at all -- would otherwise silently drop a scenario the old, +# non-deterministic selection used to hit by chance. The two fail by different routes: +# killing an idle worker leaves the loss to be *detected*, killing a serving one destroys +# an in-flight RPC that surfaces at once (222s vs 12s when measured). A regression in +# either is invisible to the other. +# +# Cheap to add: the serving path fails in seconds, so this is dominated by startup. +run_test env VICTIM_STATE=serving uv run --no-sync bash ./tests/functional/grpo_dp_single_controller_chaos.sh + +# Checkpoint save/restore (upstream #3429). +run_test uv run --no-sync bash ./tests/functional/grpo_checkpoint_single_controller.sh + cd ${PROJECT_ROOT}/tests if compgen -G ".coverage*" > /dev/null; then coverage combine .coverage* From c253574c348fe195b5686a7e41f170632afdbadc Mon Sep 17 00:00:00 2001 From: Yuki Huang Date: Mon, 24 Aug 2026 08:15:09 -0700 Subject: [PATCH 23/31] fix(sc): require a Megatron critic and reject drop budgets under PPO, hoist the warmup-lookahead guard onto the GRPO path, unprefix shared validation messages, retune the nightly thresholds Signed-off-by: Yuki Huang --- .../single_controller_utils/config.py | 71 ++++++++++++++----- .../single_controller_utils/setup.py | 4 +- ...ch-noncolocated-async-single-controller.sh | 11 ++- .../unit/single_controller/test_ppo_setup.py | 49 +++++++++++-- 4 files changed, 103 insertions(+), 32 deletions(-) diff --git a/nemo_rl/algorithms/single_controller_utils/config.py b/nemo_rl/algorithms/single_controller_utils/config.py index 3e1f63ade3b..827f98dbe05 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/nemo_rl/algorithms/single_controller_utils/config.py @@ -669,7 +669,7 @@ def _validate_failure_settings( warnings.warn( f"async_rl.rollout_failure.min_step_batch_fraction=" f"{failure_config.min_step_batch_fraction} gives a floor of {floor} of " - f"grpo.num_prompts_per_step={num_prompts_per_step}, so no prompt may be " + f"num_prompts_per_step={num_prompts_per_step}, so no prompt may be " f"dropped -- but max_skipped_prompts=" f"{failure_config.max_skipped_prompts} / max_consecutive_dropped_prompts=" f"{failure_config.max_consecutive_dropped_prompts} permit drops, and the " @@ -696,7 +696,7 @@ def _validate_failure_settings( warnings.warn( f"async_rl.rollout_failure.replacement_reserve_prompts=" f"{failure_config.replacement_reserve_prompts} exceeds " - f"grpo.num_prompts_per_step={num_prompts_per_step}. The pool is refilled " + f"num_prompts_per_step={num_prompts_per_step}. The pool is refilled " "by diverting one whole dataloader batch, so a mark above one batch " "diverts several in a row before any is admitted -- and diverting happens " "before admit(), outside the sampler gate and the buffer-capacity valve, " @@ -736,14 +736,33 @@ def _validate_algo_settings(master_config: MasterConfig) -> None: if enabled ] if unsupported: - prefix = "ppo" if is_ppo_run(master_config) else "grpo" - names = ", ".join(f"{prefix}.{name}" for name in unsupported) + names = ", ".join(unsupported) raise NotImplementedError( f"{names} not supported on the SingleController path, which " "implements none of them -- the run would silently skip the " "shaping. Disable them." ) + async_config = master_config.async_rl + # Capacity is sized from the peak window whatever the algorithm, so an inert + # setting still costs buffer and fails setup naming the wrong cause. + if ( + getattr(async_config.sampler, "warmup_lookahead_versions", None) is not None + and getattr(algo_cfg, "policy_training_start_step", 0) == 0 + ): + if is_ppo_run(master_config): + raise ValueError( + "async_rl.sampler.warmup_lookahead_versions requires " + "ppo.policy_training_start_step > 0; without critic warmup there is " + "no frozen-policy window to widen." + ) + raise ValueError( + "async_rl.sampler.warmup_lookahead_versions is a PPO critic-warmup knob " + "and this run has no `ppo` block, so nothing ever widens the window -- " + "but max_buffered_rollouts is still validated against the wider one. " + "Remove it, or add a `ppo` block with policy_training_start_step > 0." + ) + if not is_ppo_run(master_config): # A value block without `ppo` is inert -- nothing builds the critic -- # and a config carrying one is asking for PPO by every reading except @@ -764,10 +783,20 @@ def _validate_algo_settings(master_config: MasterConfig) -> None: "See examples/configs/ppo_math_1B_megatron_single_controller.yaml." ) + # Only megatron_value_worker mixes in TQWorkerMixin; TQValue fans out + # setup_data_plane unconditionally, so a DTensor critic dies in Ray with the + # model already on GPU. ppo_math_1B.yaml ships dtensor_cfg.enabled=true. + value_megatron_cfg = master_config.value.get("megatron_cfg", {}) # type: ignore + if not value_megatron_cfg.get("enabled"): + raise ValueError( + "PPO on the SingleController path requires a Megatron critic " + "(value.megatron_cfg.enabled=true). The DTensor value worker does not " + "carry TQWorkerMixin, so it has no data-plane setup to call (#2625)." + ) + if algo_cfg.ppo_epochs < 1: raise ValueError("ppo.ppo_epochs must be at least 1") - async_config = master_config.async_rl # Without it the critic steps once per chunk and the policy once per step, # which is two effective learning rates from one config, and no error. if async_config.min_groups_for_streaming_train != algo_cfg.num_prompts_per_step: @@ -782,6 +811,22 @@ def _validate_algo_settings(master_config: MasterConfig) -> None: "not have yet (#2625)." ) + failure_config = async_config.rollout_failure + drop_budget = ( + failure_config.max_skipped_prompts + + failure_config.max_consecutive_dropped_prompts + ) + if drop_budget > 0: + raise ValueError( + "PPO on the SingleController path requires " + "async_rl.rollout_failure.max_skipped_prompts=0 and " + "max_consecutive_dropped_prompts=0, but they sum to " + f"{drop_budget}. A drop shortens the step, and the critic shards that " + "step against the configured value.train_global_batch_size rather than " + "its actual size, so the first short step fails a divisibility assert " + "inside the value workers (#2625)." + ) + policy_megatron_cfg = master_config.policy.get("megatron_cfg", {}) # type: ignore if ( getattr(algo_cfg, "policy_training_start_step", 0) > 0 @@ -799,16 +844,6 @@ def _validate_algo_settings(master_config: MasterConfig) -> None: "ppo.policy_training_start_step=0, or checkpointing.save_optimizer=false." ) - if ( - getattr(async_config.sampler, "warmup_lookahead_versions", None) is not None - and getattr(algo_cfg, "policy_training_start_step", 0) == 0 - ): - raise ValueError( - "async_rl.sampler.warmup_lookahead_versions requires " - "ppo.policy_training_start_step > 0; without critic warmup there is " - "no frozen-policy window to widen." - ) - sampler_name = async_config.sampler.name if sampler_name != "in_order": raise ValueError( @@ -842,7 +877,7 @@ def validate_single_controller_config(master_config: MasterConfig) -> None: if algo_cfg.num_prompts_per_step < async_config.min_groups_for_streaming_train: raise ValueError( - f"grpo.num_prompts_per_step ({algo_cfg.num_prompts_per_step}) " + f"num_prompts_per_step ({algo_cfg.num_prompts_per_step}) " f"must be >= async_rl.min_groups_for_streaming_train " f"({async_config.min_groups_for_streaming_train})" ) @@ -921,9 +956,9 @@ def validate_single_controller_config(master_config: MasterConfig) -> None: raise ValueError( "loss_fn.reference_policy_kl_penalty=" f"{reference_policy_kl_penalty} requires reference_policy_logprobs, " - "but grpo.skip_reference_policy_logprobs_calculation=true skips " + "but skip_reference_policy_logprobs_calculation=true skips " "computing them on the SingleController path. Set " - "grpo.skip_reference_policy_logprobs_calculation=false, or set " + "skip_reference_policy_logprobs_calculation=false, or set " "loss_fn.reference_policy_kl_penalty=0." ) diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py index 7f88866990e..682220ac4ce 100644 --- a/nemo_rl/algorithms/single_controller_utils/setup.py +++ b/nemo_rl/algorithms/single_controller_utils/setup.py @@ -418,7 +418,7 @@ def _generation_max_seq_len(generation_config) -> int: def _clamp_max_num_steps( master_config: MasterConfig, dataloader: StatefulDataLoader ) -> None: - """Clamp grpo.max_num_steps to max_num_epochs * len(dataloader).""" + """Clamp max_num_steps to max_num_epochs * len(dataloader).""" algo_cfg = algo_config(master_config) max_num_epochs = algo_cfg.max_num_epochs if max_num_epochs is None or max_num_epochs <= 0: @@ -610,7 +610,7 @@ def setup_single_controller( if algo_cfg.val_period > 0 or algo_cfg.val_at_start or algo_cfg.val_at_end: raise NotImplementedError( "SingleController doesn't support validation now, will support " - "later. Set grpo.val_period=0, val_at_start=false, val_at_end=false." + "later. Set val_period=0, val_at_start=false, val_at_end=false." ) if dp_config is None or not dp_config.get("enabled", False): raise ValueError( diff --git a/tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.sh b/tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.sh index aaace3a86cf..f1666482055 100755 --- a/tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.sh +++ b/tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.sh @@ -7,7 +7,7 @@ NUM_NODES=2 STEPS_PER_RUN=40 MAX_STEPS=40 NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) # Round up -NUM_MINUTES=60 +NUM_MINUTES=40 # ===== END CONFIG ===== exit_if_max_steps_reached @@ -37,17 +37,16 @@ grep -q "weight_sync=CollectiveWeightSynchronizer" $RUN_LOG # Convert tensorboard logs to json uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS -# Only run metrics if the target step is reached. Thresholds are the legacy -# recipe's, at the step where this run has consumed the same data. SC has no -# validation loop, so its validation/accuracy assertion has no counterpart. +# Only run metrics if the target step is reached if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | map(tonumber) | max' $JSON_METRICS) -ge $MAX_STEPS ]]; then uv run tests/check_metrics.py $JSON_METRICS \ 'median(data["train/token_mult_prob_error"]) < 1.1' \ 'data["train/token_mult_prob_error"]["40"] < 1.1' \ 'median(data["train/max_seq_mult_prob_error"]) < 1.2' \ 'len(data["train/critic/loss"]) == 40' \ - 'min(data["train/critic/loss"]) >= 0' \ - 'data["train/reward"]["40"] > 0.75' + 'max(data["train/critic/loss"]) < 1.5' \ + 'mean(data["train/critic/explained_var"], range_start=-10) > 0.5' \ + 'mean(data["train/reward"], range_start=-10) > 0.75' # Clean up checkpoint directory after successful run to save space. rm -rf "$CKPT_DIR" diff --git a/tests/unit/single_controller/test_ppo_setup.py b/tests/unit/single_controller/test_ppo_setup.py index e816335ded1..97770a8f953 100644 --- a/tests/unit/single_controller/test_ppo_setup.py +++ b/tests/unit/single_controller/test_ppo_setup.py @@ -58,7 +58,7 @@ def _value_config( *, - megatron_enabled: bool = False, + megatron_enabled: bool = True, train_global_batch_size: int = _GLOBAL_BATCH_SIZE, ) -> dict: return { @@ -149,7 +149,7 @@ def _ppo_master_config(**kwargs) -> MasterConfig: ), ) kwargs.setdefault( - "value", _value_config(megatron_enabled=kwargs.get("megatron_enabled", False)) + "value", _value_config(megatron_enabled=kwargs.get("megatron_enabled", True)) ) kwargs.setdefault("value_loss_fn", MseValueLossConfig()) return _make_master_config(**kwargs) @@ -322,10 +322,36 @@ def test_rejects_shaping_on_a_grpo_run_too(self): mc.grpo.overlong_filtering = True with pytest.raises( - NotImplementedError, match=r"grpo\.overlong_filtering" + NotImplementedError, match="overlong_filtering not supported" ): validate_single_controller_config(mc) + def test_rejects_a_dtensor_critic(self): + """Only the Megatron value worker carries TQWorkerMixin.""" + mc = _ppo_master_config(value=_value_config(megatron_enabled=False)) + + with pytest.raises(ValueError, match=r"value\.megatron_cfg\.enabled=true"): + validate_single_controller_config(mc) + + @pytest.mark.parametrize( + "field", ["max_skipped_prompts", "max_consecutive_dropped_prompts"] + ) + def test_rejects_a_drop_budget_under_ppo(self, field): + """A drop shortens the step; the critic shards against the configured size.""" + mc = _ppo_master_config() + setattr(mc.async_rl.rollout_failure, field, 1) + + with pytest.raises(ValueError, match=r"max_skipped_prompts=0"): + validate_single_controller_config(mc) + + def test_rejects_warmup_lookahead_on_a_grpo_run(self): + """GRPO never widens the window, but capacity is still sized for it.""" + mc = _make_master_config() + mc.async_rl.sampler.warmup_lookahead_versions = 2 + + with pytest.raises(ValueError, match="PPO critic-warmup knob"): + validate_single_controller_config(mc) + def test_grpo_is_free_to_use_any_sampler(self): mc = _make_master_config() mc.async_rl.sampler = WindowedSamplerConfig() @@ -462,16 +488,27 @@ def test_policy_steps_off_the_gpu_while_the_critic_loads( ): """Both worker groups sit on the training cluster; leaving the policy resident through the critic's init is what OOMs a tight fit.""" + calls: list[str] = [] policy = patched_ppo_factories["policy"] value = patched_ppo_factories["value"] + policy.offload_to_cpu.side_effect = lambda: calls.append("policy.offload") + policy.prepare_for_training.side_effect = lambda: calls.append("policy.onload") + value.finish_training.side_effect = lambda: calls.append("critic.finish") + patched_ppo_factories["_build_value"].side_effect = lambda *a, **k: ( + calls.append("critic.build"), + (value, 2.0), + )[1] setup_single_controller( _ppo_master_config(), tokenizer=MagicMock(pad_token_id=0) ) - policy.offload_to_cpu.assert_called_once_with() - value.finish_training.assert_called_once_with() - policy.prepare_for_training.assert_called_once_with() + assert calls == [ + "policy.offload", + "critic.build", + "critic.finish", + "policy.onload", + ] def test_grpo_run_builds_no_critic(self, patched_ppo_factories): actor_args, timing = setup_single_controller( From 018cc7e6e128dd030001c617e5384074cf14e702 Mon Sep 17 00:00:00 2001 From: Yuki Huang Date: Mon, 24 Aug 2026 09:00:50 -0700 Subject: [PATCH 24/31] test(sc): remove four duplicated PPO unit tests, parametrize five near-identical groups Removed: - test_ppo_setup.py TestPPOValidation::test_the_ppo_schema_rejects_a_non_ppo_estimator asserted exactly what test_ppo.py::test_ppo_schema_rejects_unsupported_estimator_name already does, and it covers the PPOConfig schema rather than SC setup. - test_single_controller_actor.py::test_train_pump_loads_and_offloads_the_critic_around_each_stage asserted a critic call sequence that is a strict subsequence of the interleaved list in test_train_pump_parks_the_policy_on_cpu_across_the_critic_stages. - test_single_controller_actor.py::test_train_pump_steps_both_optimizers_once_per_ppo_epoch asserted per-epoch counts that follow from the exact call list in test_train_pump_offloads_the_policy_between_ppo_epochs; its two unique assertions, one refit and one version bump per RL step, moved there. - test_sampler_interface.py TestWarmupLookaheadWindow::test_set_gate_window_retunes_admission asserted the private _gate_window; the setter's observable effect is covered by test_capacity_does_not_shrink_when_the_gate_is_retuned and by TestLookaheadSchedule. Parametrized, with the same set of cases as before: - TestIsPPORun's three tests differed only in the config handed to a one-line predicate. - TestTrainClusterSizesForTheCritic's five tests were one body over (colocated, backend, algorithm); the _groups helper goes with them. - TestMegatronTrainIters' two injection tests differed only in ppo_epochs. - test_train_pump_runs_every_critic_epoch_during_warmup folds into test_train_pump_freezes_the_policy_during_critic_warmup as ppo_epochs=2, with the critic-train count asserted against ppo_epochs. - test_init_leaves_the_critic_handles_unset_on_a_grpo_run folds into test_init_picks_up_the_critic_handles as the handles-absent case. - TestPPOWarmupCheckpoint's two optimizer-path tests differed only in is_policy_training_step. The metric pair stays split: one expects a warning, the other an exception. - test_tq_value.py's two get_values_from_meta tests shared a patch stack; the packing-budget check is now one assertion in the surviving test. test_rejects_a_config_with_neither_block and test_rejects_a_config_with_both_blocks move to TestPPOValidation, since both exercise validate_single_controller_config rather than algo_config. Signed-off-by: Yuki Huang --- tests/unit/models/value/test_tq_value.py | 24 +--- .../single_controller/test_checkpointing.py | 19 +-- .../unit/single_controller/test_ppo_setup.py | 127 +++++++----------- .../test_sampler_interface.py | 7 - .../test_single_controller_actor.py | 123 +++++------------ 5 files changed, 95 insertions(+), 205 deletions(-) diff --git a/tests/unit/models/value/test_tq_value.py b/tests/unit/models/value/test_tq_value.py index 9c6173e4daa..ac43417a8cd 100644 --- a/tests/unit/models/value/test_tq_value.py +++ b/tests/unit/models/value/test_tq_value.py @@ -121,7 +121,9 @@ def test_get_values_from_meta_narrows_fields_and_returns_none(self): meta = _meta() with ( patch.object(TQValue, "_stamp_pad_seqlen"), - patch.object(TQValue, "_packing_args", return_value=(None, None)), + patch.object( + TQValue, "_packing_args", return_value=(None, None) + ) as mock_packing, patch( "nemo_rl.models.value.tq_value.shard_meta_for_dp", return_value=([meta, meta], None), @@ -137,24 +139,8 @@ def test_get_values_from_meta_narrows_fields_and_returns_none(self): wg.run_all_workers_sharded_data.call_args.args[0] == "get_values_presharded" ) wg.get_all_worker_results.assert_called_once() - - def test_get_values_from_meta_uses_the_logprob_packing_budget(self): - """The forward pass is inference-shaped, so it must size microbatches - off logprob_mb_tokens rather than the (larger) train budget.""" - v, _ = _make_tq_value() - meta = _meta() - with ( - patch.object(TQValue, "_stamp_pad_seqlen"), - patch.object( - TQValue, "_packing_args", return_value=(None, None) - ) as mock_packing, - patch( - "nemo_rl.models.value.tq_value.shard_meta_for_dp", - return_value=([meta, meta], None), - ), - ): - v.get_values_from_meta(meta) - + # The forward pass is inference-shaped, so it sizes microbatches off + # logprob_mb_tokens rather than the (larger) train budget. assert mock_packing.call_args.args[0] == "logprob_mb_tokens" def test_train_from_meta_requests_the_value_train_columns(self): diff --git a/tests/unit/single_controller/test_checkpointing.py b/tests/unit/single_controller/test_checkpointing.py index f0a4472bb21..ae76355e83e 100644 --- a/tests/unit/single_controller/test_checkpointing.py +++ b/tests/unit/single_controller/test_checkpointing.py @@ -889,18 +889,19 @@ def actor(self, tmp_path, monkeypatch): ) return _ppo_save_actor(tmp_path, []) - def test_warmup_step_writes_no_policy_optimizer(self, actor): - asyncio.run(actor._save_checkpoint({}, is_policy_training_step=False)) + @pytest.mark.parametrize("is_policy_training_step", [False, True]) + def test_the_policy_optimizer_is_written_only_once_it_has_stepped( + self, actor, is_policy_training_step + ): + asyncio.run( + actor._save_checkpoint({}, is_policy_training_step=is_policy_training_step) + ) - assert actor._trainer.save_kwargs["optimizer_path"] is None - # The critic trains from step 0, so its optimizer is still written. + written = actor._trainer.save_kwargs["optimizer_path"] is not None + assert written is is_policy_training_step + # The critic trains from step 0, so its optimizer is always written. assert actor._value.save_kwargs["optimizer_path"] is not None - def test_training_step_writes_the_policy_optimizer(self, actor): - asyncio.run(actor._save_checkpoint({}, is_policy_training_step=True)) - - assert actor._trainer.save_kwargs["optimizer_path"] is not None - def test_warmup_step_skips_the_top_k_metric(self, actor): """No policy metrics exist yet, so the checkpoint just is not a candidate.""" actor._master_config.checkpointing["metric_name"] = "train:loss" diff --git a/tests/unit/single_controller/test_ppo_setup.py b/tests/unit/single_controller/test_ppo_setup.py index 97770a8f953..23a0635eb95 100644 --- a/tests/unit/single_controller/test_ppo_setup.py +++ b/tests/unit/single_controller/test_ppo_setup.py @@ -156,15 +156,18 @@ def _ppo_master_config(**kwargs) -> MasterConfig: class TestIsPPORun: - def test_absent_ppo_block_is_grpo(self): - assert is_ppo_run(_make_master_config()) is False - - def test_present_ppo_block_is_ppo(self): - assert is_ppo_run(_ppo_master_config()) is True - - def test_missing_attribute_is_grpo(self): - """model_construct can omit defaulted fields entirely.""" - assert is_ppo_run(MasterConfig.model_construct()) is False + @pytest.mark.parametrize( + ("make_config", "expected"), + [ + (_make_master_config, False), + (_ppo_master_config, True), + # model_construct can omit defaulted fields entirely. + (MasterConfig.model_construct, False), + ], + ids=["grpo_block", "ppo_block", "no_block_at_all"], + ) + def test_the_ppo_block_is_what_selects_the_path(self, make_config, expected): + assert is_ppo_run(make_config()) is expected class TestAlgoConfigSelection: @@ -180,6 +183,11 @@ def test_returns_the_grpo_block_otherwise(self): assert algo_config(mc) is mc.grpo + +class TestPPOValidation: + def test_accepts_a_well_formed_ppo_config(self): + validate_single_controller_config(_ppo_master_config()) + def test_rejects_a_config_with_neither_block(self): """Also caught at setup; algo_config only asserts the invariant.""" mc = _make_master_config() @@ -196,11 +204,6 @@ def test_rejects_a_config_with_both_blocks(self): with pytest.raises(ValueError, match="Only one algorithm block"): validate_single_controller_config(mc) - -class TestPPOValidation: - def test_accepts_a_well_formed_ppo_config(self): - validate_single_controller_config(_ppo_master_config()) - @pytest.mark.parametrize("missing", ["value", "value_loss_fn"]) def test_rejects_ppo_without_its_critic_blocks(self, missing): mc = _ppo_master_config(**{missing: None}) @@ -246,10 +249,6 @@ def test_rejects_a_ppo_epoch_count_below_one(self): with pytest.raises(ValueError, match="ppo_epochs must be at least 1"): validate_single_controller_config(mc) - def test_the_ppo_schema_rejects_a_non_ppo_estimator(self): - with pytest.raises(ValueError): - PPOConfig(adv_estimator={"name": "grpo"}) - @pytest.mark.parametrize( "sampler_config", [WindowedSamplerConfig(), ReadyFirstSamplerConfig(), WeightFifoSamplerConfig()], @@ -388,32 +387,20 @@ def test_grpo_delegates_to_the_group_relative_factory(self): class TestMegatronTrainIters: - def test_injects_into_both_policy_and_value(self): - mc = _ppo_master_config( - megatron_enabled=True, - ppo=PPOConfig.model_construct( - max_num_steps=7, ppo_epochs=1, **_STEP_CONFIG - ), - ) - - sc_setup_mod._maybe_inject_megatron_train_iters(mc) - - assert mc.policy["megatron_cfg"]["train_iters"] == 7 - assert mc.value["megatron_cfg"]["train_iters"] == 7 - - def test_scales_the_tick_budget_by_ppo_epochs(self): + @pytest.mark.parametrize(("ppo_epochs", "expected"), [(1, 7), (3, 21)]) + def test_injects_into_both_policy_and_value(self, ppo_epochs, expected): """Each epoch steps both optimizers, so each is a scheduler tick.""" mc = _ppo_master_config( megatron_enabled=True, ppo=PPOConfig.model_construct( - max_num_steps=7, ppo_epochs=3, **_STEP_CONFIG + max_num_steps=7, ppo_epochs=ppo_epochs, **_STEP_CONFIG ), ) sc_setup_mod._maybe_inject_megatron_train_iters(mc) - assert mc.policy["megatron_cfg"]["train_iters"] == 21 - assert mc.value["megatron_cfg"]["train_iters"] == 21 + assert mc.policy["megatron_cfg"]["train_iters"] == expected + assert mc.value["megatron_cfg"]["train_iters"] == expected def test_skips_a_critic_on_a_non_megatron_backend(self): mc = _ppo_master_config(megatron_enabled=False, max_num_steps=7) @@ -549,50 +536,34 @@ def fake_cluster(self): cls.side_effect = lambda **kwargs: MagicMock(kwargs=kwargs) yield cls - def _groups(self, mc): - train, inference = sc_setup_mod._build_clusters(mc) - return train.kwargs["max_colocated_worker_groups"], inference.kwargs[ - "max_colocated_worker_groups" - ] - - def test_noncolocated_ppo_leaves_a_slot_for_the_critic(self, fake_cluster): - mc = _cluster_config(_ppo_master_config(), colocated=False, backend="vllm") - - train_groups, inference_groups = self._groups(mc) - - assert train_groups == 2 - # The critic never lands on the inference cluster. - assert inference_groups == 1 - - def test_noncolocated_grpo_is_unchanged(self, fake_cluster): - mc = _cluster_config(_make_master_config(), colocated=False, backend="vllm") - - train_groups, inference_groups = self._groups(mc) - - assert train_groups == 1 - assert inference_groups == 1 - - def test_colocated_ppo_adds_the_critic_beside_policy_and_generation( - self, fake_cluster + @pytest.mark.parametrize( + ("make_config", "colocated", "backend", "expected_train_groups"), + [ + (_ppo_master_config, False, "vllm", 2), + (_make_master_config, False, "vllm", 1), + (_ppo_master_config, True, "vllm", 3), + (_make_master_config, True, "vllm", 2), + # The megatron backend generates from the policy's own workers. + (_ppo_master_config, True, "megatron", 2), + ], + ids=[ + "noncolocated_ppo", + "noncolocated_grpo", + "colocated_ppo", + "colocated_grpo", + "colocated_ppo_megatron_generation", + ], + ) + def test_worker_group_slots( + self, fake_cluster, make_config, colocated, backend, expected_train_groups ): - mc = _cluster_config(_ppo_master_config(), colocated=True, backend="vllm") + mc = _cluster_config(make_config(), colocated=colocated, backend=backend) train, inference = sc_setup_mod._build_clusters(mc) - assert train is inference - assert train.kwargs["max_colocated_worker_groups"] == 3 - - def test_colocated_grpo_is_unchanged(self, fake_cluster): - mc = _cluster_config(_make_master_config(), colocated=True, backend="vllm") - - train, _ = sc_setup_mod._build_clusters(mc) - - assert train.kwargs["max_colocated_worker_groups"] == 2 - - def test_colocated_megatron_generation_needs_no_extra_slot(self, fake_cluster): - """The megatron backend generates from the policy's own workers.""" - mc = _cluster_config(_ppo_master_config(), colocated=True, backend="megatron") - - train, _ = sc_setup_mod._build_clusters(mc) - - assert train.kwargs["max_colocated_worker_groups"] == 2 + assert train.kwargs["max_colocated_worker_groups"] == expected_train_groups + if colocated: + assert train is inference + else: + # The critic never lands on the inference cluster. + assert inference.kwargs["max_colocated_worker_groups"] == 1 diff --git a/tests/unit/single_controller/test_sampler_interface.py b/tests/unit/single_controller/test_sampler_interface.py index 973983a5ca7..f5bd1b11f3e 100644 --- a/tests/unit/single_controller/test_sampler_interface.py +++ b/tests/unit/single_controller/test_sampler_interface.py @@ -262,13 +262,6 @@ def test_capacity_does_not_shrink_when_the_gate_is_retuned(self): assert sampler.required_buffer_capacity(groups_per_step=4) == 16 - def test_set_gate_window_retunes_admission(self): - sampler = InOrderSampler(FakeBuffer(), max_lookahead_versions=1) - - sampler.set_gate_window(3) - - assert sampler._gate_window == 3 - def test_set_gate_window_rejects_a_negative_window(self): sampler = InOrderSampler(FakeBuffer(), max_lookahead_versions=1) diff --git a/tests/unit/single_controller/test_single_controller_actor.py b/tests/unit/single_controller/test_single_controller_actor.py index 185b0990b95..54e001d8b53 100644 --- a/tests/unit/single_controller/test_single_controller_actor.py +++ b/tests/unit/single_controller/test_single_controller_actor.py @@ -254,30 +254,29 @@ def test_reference_logprobs_required_only_when_kl_enabled( assert controller._reference_logprobs_required is expected_required -def test_init_picks_up_the_critic_handles(monkeypatch, tmp_path) -> None: - """The PPO path hands the critic and its loss in through actor_args.""" +@pytest.mark.parametrize("with_critic", [True, False], ids=["ppo", "grpo"]) +def test_init_picks_up_the_critic_handles(monkeypatch, tmp_path, with_critic) -> None: + """The PPO path hands the critic and its loss in through actor_args. + + A GRPO run leaves both unset -- actor_args defaults them to None, and older + args may omit them entirely. + """ monkeypatch.setattr(single_controller, "Logger", lambda _: MagicMock()) - value, value_loss_fn = MagicMock(name="value"), MagicMock(name="value_loss_fn") + handles = ( + { + "value_handle": MagicMock(name="value"), + "value_loss_fn": MagicMock(name="value_loss_fn"), + } + if with_critic + else {} + ) ctrl = _init_controller( - _grpo_master_config(tmp_path), - _actor_args_for_init(value_handle=value, value_loss_fn=value_loss_fn), + _grpo_master_config(tmp_path), _actor_args_for_init(**handles) ) - assert ctrl._value is value - assert ctrl._value_loss_fn is value_loss_fn - - -def test_init_leaves_the_critic_handles_unset_on_a_grpo_run( - monkeypatch, tmp_path -) -> None: - """actor_args defaults them to None, and older args may omit them entirely.""" - monkeypatch.setattr(single_controller, "Logger", lambda _: MagicMock()) - - ctrl = _init_controller(_grpo_master_config(tmp_path), _actor_args_for_init()) - - assert ctrl._value is None - assert ctrl._value_loss_fn is None + assert ctrl._value is handles.get("value_handle") + assert ctrl._value_loss_fn is handles.get("value_loss_fn") def test_logs_setup_timing_metrics(monkeypatch, tmp_path) -> None: @@ -1345,32 +1344,6 @@ def _single_group_meta() -> KVBatchMeta: ) -def test_train_pump_loads_and_offloads_the_critic_around_each_stage( - monkeypatch, -) -> None: - """The critic holds the training GPUs only inside its forward and its train.""" - meta = _single_group_meta() - ctrl, value = _ppo_train_pump_controller(sampler=_OneThenEmptySampler(meta)) - ctrl._policy_logprobs_required = True - trainer = _LpRecordingTrainer() - ctrl._trainer = trainer - ctrl._advantage_stage = AsyncMock(return_value=(meta, True)) - monkeypatch.setattr(single_controller.ray, "cluster_resources", lambda: {}) - - asyncio.run(asyncio.wait_for(ctrl._train_pump(), timeout=1.0)) - - assert value.calls == [ - "prepare_for_inference", - "get_values_from_meta", - "finish_inference", - # Critic before actor, mirroring the legacy PPO epoch loop. - "prepare_for_training", - "train_from_meta", - "finish_training", - ] - assert ctrl._train_steps == 1 - - def test_train_pump_parks_the_policy_on_cpu_across_the_critic_stages( monkeypatch, ) -> None: @@ -1436,15 +1409,18 @@ def test_train_pump_skips_the_critic_on_an_empty_chunk(monkeypatch) -> None: assert "get_values_from_meta" in value.calls +@pytest.mark.parametrize("ppo_epochs", [1, 2]) def test_train_pump_freezes_the_policy_during_critic_warmup( - monkeypatch, capsys + monkeypatch, capsys, ppo_epochs ) -> None: """Below policy_training_start_step the critic trains alone: no optimizer - step, and no weight transfer to generation either.""" + step, and no weight transfer to generation either. The frozen policy does + not shorten the critic's own epoch loop.""" meta = _single_group_meta() ctrl, value = _ppo_train_pump_controller( sampler=_OneThenEmptySampler(meta), policy_training_start_step=1, + ppo_epochs=ppo_epochs, ) trainer = MagicMock(spec=_NoOpTrainer) ctrl._trainer = trainer @@ -1453,7 +1429,7 @@ def test_train_pump_freezes_the_policy_during_critic_warmup( asyncio.run(asyncio.wait_for(ctrl._train_pump(), timeout=1.0)) - assert "train_from_meta" in value.calls + assert value.calls.count("train_from_meta") == ppo_epochs trainer.prepare_for_training.assert_not_called() trainer.begin_train_step.assert_not_called() trainer.finish_train_step.assert_not_called() @@ -1490,51 +1466,11 @@ def test_train_pump_trains_the_policy_once_warmup_is_over(monkeypatch, capsys) - assert capsys.readouterr().out.count("Critic warmup complete") == 1 -def test_train_pump_steps_both_optimizers_once_per_ppo_epoch(monkeypatch) -> None: - """ppo_epochs repeats the whole train stage over the step's own batch.""" - meta = _single_group_meta() - ctrl, value = _ppo_train_pump_controller( - sampler=_OneThenEmptySampler(meta), ppo_epochs=2 - ) - trainer = MagicMock(spec=_NoOpTrainer) - trainer.finish_train_step.return_value = {} - ctrl._trainer = trainer - ctrl._advantage_stage = AsyncMock(return_value=(meta, True)) - monkeypatch.setattr(single_controller.ray, "cluster_resources", lambda: {}) - - asyncio.run(asyncio.wait_for(ctrl._train_pump(), timeout=1.0)) - - assert value.calls.count("train_from_meta") == 2 - assert trainer.begin_train_step.call_count == 2 - assert trainer.finish_train_step.call_count == 2 - # Still one RL step, so one refit and one version bump. - ctrl._sync_weights.assert_awaited_once_with(calibration_data=None) - assert ctrl._trainer_version == 1 - - -def test_train_pump_runs_every_critic_epoch_during_warmup(monkeypatch) -> None: - """The frozen policy does not shorten the critic's own epoch loop.""" - meta = _single_group_meta() - ctrl, value = _ppo_train_pump_controller( - sampler=_OneThenEmptySampler(meta), - policy_training_start_step=1, - ppo_epochs=2, - ) - trainer = MagicMock(spec=_NoOpTrainer) - ctrl._trainer = trainer - ctrl._advantage_stage = AsyncMock(return_value=(meta, True)) - monkeypatch.setattr(single_controller.ray, "cluster_resources", lambda: {}) - - asyncio.run(asyncio.wait_for(ctrl._train_pump(), timeout=1.0)) - - assert value.calls.count("train_from_meta") == 2 - trainer.prepare_for_training.assert_not_called() - trainer.finish_train_step.assert_not_called() - - def test_train_pump_offloads_the_policy_between_ppo_epochs(monkeypatch) -> None: - """The two models share the training GPUs, so every critic train runs with - the policy on CPU -- including the ones after the first epoch.""" + """ppo_epochs repeats the whole train stage over the step's own batch. + + The two models share the training GPUs, so every critic train runs with the + policy on CPU -- including the ones after the first epoch.""" meta = _single_group_meta() calls: list[str] = [] ctrl, _ = _ppo_train_pump_controller( @@ -1570,6 +1506,9 @@ def test_train_pump_offloads_the_policy_between_ppo_epochs(monkeypatch) -> None: # No offload after the last epoch: the refit needs the policy resident. "policy.finish_train_step", ] + # Still one RL step, so one refit and one version bump. + ctrl._sync_weights.assert_awaited_once_with(calibration_data=None) + assert ctrl._trainer_version == 1 def test_advantage_stage_writes_gae_returns_alongside_advantages() -> None: From a091a6f50e14b36250523e30e45cdd1e359e058d Mon Sep 17 00:00:00 2001 From: Yuki Huang Date: Mon, 24 Aug 2026 20:15:41 -0700 Subject: [PATCH 25/31] test(sc): set the first-train-step pre-hook state on the fabricated worker, cover the finish-path re-enable Signed-off-by: Yuki Huang --- .../policy/test_megatron_split_state.py | 31 +++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/tests/unit/models/policy/test_megatron_split_state.py b/tests/unit/models/policy/test_megatron_split_state.py index 0df3f18dd8d..61eaf84a855 100644 --- a/tests/unit/models/policy/test_megatron_split_state.py +++ b/tests/unit/models/policy/test_megatron_split_state.py @@ -148,6 +148,10 @@ def _make_worker(loss_type): # The step summary in finish_train_step reads it eagerly to decide whether # this rank prints. w.rank = 0 + # Also set in __init__: the finish path reads them to put the DDP forward + # pre-hook back after the first optimizer step. + w._first_train_step_forward_pre_hook_disabled = False + w._first_train_step_param_sync_func = None # Pure telemetry, and it resets the CUDA peak counters — keep it out of the # way so these tests stay hermetic on GPU shards. w._log_gpu_mem = MagicMock() @@ -535,6 +539,33 @@ def test_restores_grad_sync_func(self, mock_module_symbols): w.finish_train_step() assert w.model.config.grad_sync_func == "ORIGINAL_GRAD_SYNC_FUNC" + @pytest.mark.parametrize("update_successful", [True, False]) + def test_reenables_forward_pre_hook_after_a_successful_step( + self, mock_module_symbols, update_successful + ): + """__init__ disables the pre-hook for the first step; finish has to put + it back once the optimizer stepped, or every later forward sees only its + own param shard.""" + from nemo_rl.algorithms.loss.interfaces import LossType + + w = self._setup_open_step(mock_module_symbols, LossType.TOKEN_LEVEL) + w._first_train_step_forward_pre_hook_disabled = True + w._first_train_step_param_sync_func = "ORIGINAL_PARAM_SYNC_FUNC" + w.enable_forward_pre_hook = MagicMock() + w.optimizer.step.return_value = (update_successful, 0.5, 0) + + with patch(f"{WORKER_MOD}.get_model_config", return_value=w.model.config): + w.finish_train_step() + + if update_successful: + w.enable_forward_pre_hook.assert_called_once_with() + assert w.model.config.param_sync_func == "ORIGINAL_PARAM_SYNC_FUNC" + assert w._first_train_step_forward_pre_hook_disabled is False + assert w._first_train_step_param_sync_func is None + else: + w.enable_forward_pre_hook.assert_not_called() + assert w._first_train_step_forward_pre_hook_disabled is True + def test_clears_train_step_state(self, mock_module_symbols): from nemo_rl.algorithms.loss.interfaces import LossType From 842adb94c7f8517d7247f93b6a5a8ee3f79b9809 Mon Sep 17 00:00:00 2001 From: Yuki Huang Date: Tue, 25 Aug 2026 06:20:59 -0700 Subject: [PATCH 26/31] docs(sc): document the PPO value stages and critic warmup, note the PPO-only drop-budget constraint Signed-off-by: Yuki Huang --- docs/guides/single-controller.md | 7 ++++--- .../configs/ppo_math_1B_megatron_single_controller.yaml | 8 ++++---- nemo_rl/algorithms/async_utils/staleness_sampler.py | 6 ++++++ nemo_rl/algorithms/single_controller.py | 6 ++++++ 4 files changed, 20 insertions(+), 7 deletions(-) diff --git a/docs/guides/single-controller.md b/docs/guides/single-controller.md index 209fdacda26..9a6bacc9688 100644 --- a/docs/guides/single-controller.md +++ b/docs/guides/single-controller.md @@ -132,14 +132,14 @@ The SC path splits the async-GRPO loop across a rollout pump and a train pump th #### 5. `_rollout_pump` and `_train_pump` - `_rollout_pump`: pulls prompts from the dataloader, calls `sampler.admit`, dispatches `RolloutManager.generate_and_push`, and honours `max_inflight_prompts` as a backpressure cap. -- `_train_pump`: `sampler.evict → sampler.select → _advantage_stage → TQPolicy split API (begin_train_step / train_microbatches_from_meta / finish_train_step) → dp_client.clear_samples`. +- `_train_pump`: `sampler.evict → sampler.select → _value_stage (PPO only) → _advantage_stage → _value_train (PPO only) → TQPolicy split API (begin_train_step / train_microbatches_from_meta / finish_train_step) → dp_client.clear_samples`. ### Coordination Flow 1. **Driver setup**: `setup_single_controller` builds the worker groups, virtual cluster, dp client, dataloader, `TQReplayBuffer`, `RolloutManager`, and weight synchronizer, and packs them into a `SingleControllerActorArgs` that the entrypoint cloudpickles into the actor. 2. **Actor startup**: `SingleControllerActor` launches `_rollout_pump` and `_train_pump` concurrently as asyncio tasks; both share the same `TQReplayBuffer` and `StalenessSampler`. 3. **Rollout pump loop**: `sampler.admit` gates dispatch against the current trainer version (returning a `target_step` for `in_order`); the pump then reserves a buffer slot, drives `RolloutManager.generate_and_push`, and commits with the observed `start_weight` / `end_weight`. -4. **Train pump loop**: `sampler.evict` drops out-of-window groups, `sampler.select` picks the next batch, `_advantage_stage` computes advantages, and the TQPolicy split API runs one optimizer step per RL step. +4. **Train pump loop**: `sampler.evict` drops out-of-window groups, `sampler.select` picks the next batch, `_value_stage` and `_value_train` run the critic forward and its optimizer step on a PPO run, `_advantage_stage` computes advantages, and the TQPolicy split API runs one optimizer step per RL step on GRPO, or `ppo.ppo_epochs` of them on PPO. 5. **Weight sync**: after each optimizer step the pump bumps the trainer version, clears rollout permission, calls the weight synchronizer, and re-opens the rollout pump for the next version. ## Relation to Legacy Async GRPO @@ -167,7 +167,7 @@ SC reads its async knobs from `async_rl:` and **requires `grpo.async_grpo: null` | `warmup_generation_lead_steps` (PPO) | `sampler.warmup_lookahead_versions` — the lookahead to use while critic warmup is in progress | | `recompute_kv_cache_after_weight_updates` | `recompute_kv_cache_after_weight_updates` (same) | | `in_flight_weight_updates` | Always effectively true; `false`-equivalent behavior is not yet supported (drain-gate tracked in [issue #2625](https://github.com/NVIDIA-NeMo/RL/issues/2625)) | -| *(no legacy equivalent — matches legacy full-batch train semantics)* | `min_groups_for_streaming_train: ${grpo.num_prompts_per_step}` | +| *(no legacy equivalent — matches legacy full-batch train semantics)* | `min_groups_for_streaming_train: ${grpo.num_prompts_per_step}`, or `${ppo.num_prompts_per_step}` on a PPO run | | *(no legacy equivalent — matches legacy `max_trajectory_age + 1` batches in flight)* | `max_inflight_prompts: num_prompts_per_step × (max_lookahead_versions + 1)` | | *(no legacy equivalent — legacy sizes its buffer to `num_prompts_per_step × max_trajectory_age_steps × 2`)* | `max_buffered_rollouts: num_prompts_per_step × (max_lookahead_versions + 1)` (tight; see the [Config → behavior map](#config--behavior-map) for per-sampler values) | @@ -178,6 +178,7 @@ The SC path is still under active development. Feature gaps are tracked in [issu - Train backend: only Megatron is supported and validated; the AutoModel training path has not been tested on SC. - Generation backend: only vLLM is supported and validated; Megatron generation, SGLang, and TRT-LLM have not been tested on SC. - Validation is not yet supported (setup raises on `val_period > 0`, `val_at_start`, or `val_at_end`). +- (PPO) Rollout drop budgets — `async_rl.rollout_failure.max_skipped_prompts` and `max_consecutive_dropped_prompts` must both be `0`. A drop shortens the step, and the critic shards it against the configured `value.train_global_batch_size` rather than its actual size, so setup rejects a non-zero budget. The resiliency layer stays available on GRPO. - Reward shaping and sample filtering — `overlong_filtering`, `reward_shaping`, `reward_scaling`, and `use_dynamic_sampling` are implemented on neither algorithm block, so setup rejects them rather than silently skipping the shaping. - The `windowed` sampler has no `over_sampling_ratio` cap — over-produced groups aged past the window are evicted, wasting rollout compute. - The drain gate in refit is not yet supported. diff --git a/examples/configs/ppo_math_1B_megatron_single_controller.yaml b/examples/configs/ppo_math_1B_megatron_single_controller.yaml index b63a64562db..f94bcc7a39e 100644 --- a/examples/configs/ppo_math_1B_megatron_single_controller.yaml +++ b/examples/configs/ppo_math_1B_megatron_single_controller.yaml @@ -56,10 +56,10 @@ async_rl: # INDEPENDENT counters, so one prompt can consume up to their sum minus one. max_skipped_prompts: 0 # CONSECUTIVE prompts allowed to exhaust their infra budget and be dropped; any - # committed rollout resets the count. 0 fails the run on the first one. Raise it - # on a large fleet where losing the odd shard is expected and a whole run is - # expensive to restart -- a step short a few groups still trains, and the count - # only keeps climbing if nothing at all is coming back. + # committed rollout resets the count. Must stay 0 on a PPO run: a drop shortens the + # step, and the critic shards it against the configured value.train_global_batch_size + # rather than its actual size, so _validate_algo_settings rejects a non-zero budget. + # (On GRPO you may raise it on a large fleet where losing the odd shard is expected.) max_consecutive_dropped_prompts: 0 # Smallest batch a step may train on, as a fraction of num_prompts_per_step. # The budgets above are run-scoped and cannot bound how short one step gets, so diff --git a/nemo_rl/algorithms/async_utils/staleness_sampler.py b/nemo_rl/algorithms/async_utils/staleness_sampler.py index 4945d54b2f4..58fe5d43ebe 100644 --- a/nemo_rl/algorithms/async_utils/staleness_sampler.py +++ b/nemo_rl/algorithms/async_utils/staleness_sampler.py @@ -443,6 +443,12 @@ class InOrderSampler(_GatedSampler): (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. + + warmup_lookahead_versions widens the gate while the PPO policy is frozen, so more + batches stay in flight during critic warmup. The driver retunes the window every + step and shrinks it back to max_lookahead_versions once the policy starts training, + so the widened lookahead does not turn into permanent extra staleness. Buffer + capacity is sized for the peak of the two. """ def __init__( diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 9926124d0ec..b079d3c8096 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -894,6 +894,12 @@ async def _train_pump(self) -> None: 5. Train the policy model (GRPO) -- finish_train_step all_reduces the accumulated gradients, rescales, and runs optimizer.step. 6. Refit the model. Sync the new policy weights to generation. + + PPO critic warmup (ppo.policy_training_start_step > 0) changes which of those + run. For the first N steps 3a still trains the critic every step, but 3b and 6 + are skipped, so the policy is neither trained nor refit. The trainer version + still advances, and the sampler's lookahead is widened while the policy is + frozen. """ policy_training_start_step = ( self._algo_cfg.policy_training_start_step if self._is_ppo else 0 From 91c7cc9e5ff69cb1623209e76df6fc5bfd356da4 Mon Sep 17 00:00:00 2001 From: Yuki Huang Date: Tue, 25 Aug 2026 06:23:39 -0700 Subject: [PATCH 27/31] test(sc): cover the PPO warmup-lookahead, negative-KL, and gate-window guards, make the warmup top-k assertion observable, tighten the nightly GPU-hour cap to 4149 Signed-off-by: Yuki Huang --- .../single_controller/test_checkpointing.py | 3 +++ .../unit/single_controller/test_ppo_setup.py | 16 +++++++++++ .../test_sampler_interface.py | 27 +++++++++++++++++++ tests/unit/test_recipes_and_test_suites.py | 6 ++--- 4 files changed, 49 insertions(+), 3 deletions(-) diff --git a/tests/unit/single_controller/test_checkpointing.py b/tests/unit/single_controller/test_checkpointing.py index ae76355e83e..804ab335753 100644 --- a/tests/unit/single_controller/test_checkpointing.py +++ b/tests/unit/single_controller/test_checkpointing.py @@ -905,6 +905,9 @@ def test_the_policy_optimizer_is_written_only_once_it_has_stepped( def test_warmup_step_skips_the_top_k_metric(self, actor): """No policy metrics exist yet, so the checkpoint just is not a candidate.""" actor._master_config.checkpointing["metric_name"] = "train:loss" + # Seed it so the delattr in the warmup branch is observable; the bare + # namespace never had the attribute, so the assertion would be vacuous. + setattr(actor._save_state, "train:loss", 1.23) with pytest.warns(UserWarning, match="not available during PPO critic warmup"): asyncio.run(actor._save_checkpoint({}, is_policy_training_step=False)) diff --git a/tests/unit/single_controller/test_ppo_setup.py b/tests/unit/single_controller/test_ppo_setup.py index 23a0635eb95..13e82b6b5b3 100644 --- a/tests/unit/single_controller/test_ppo_setup.py +++ b/tests/unit/single_controller/test_ppo_setup.py @@ -351,6 +351,22 @@ def test_rejects_warmup_lookahead_on_a_grpo_run(self): with pytest.raises(ValueError, match="PPO critic-warmup knob"): validate_single_controller_config(mc) + def test_rejects_warmup_lookahead_without_critic_warmup_on_a_ppo_run(self): + """With policy_training_start_step=0 there is no window to widen.""" + mc = _ppo_master_config() + mc.async_rl.sampler.warmup_lookahead_versions = 2 + + with pytest.raises(ValueError, match="no frozen-policy window to widen"): + validate_single_controller_config(mc) + + def test_rejects_a_negative_reference_policy_kl_penalty(self): + """ClippedPGLossConfig has no ge= constraint, so the guard is reachable.""" + mc = _ppo_master_config() + mc.loss_fn.reference_policy_kl_penalty = -0.1 + + with pytest.raises(ValueError, match="must not be negative"): + validate_single_controller_config(mc) + def test_grpo_is_free_to_use_any_sampler(self): mc = _make_master_config() mc.async_rl.sampler = WindowedSamplerConfig() diff --git a/tests/unit/single_controller/test_sampler_interface.py b/tests/unit/single_controller/test_sampler_interface.py index f5bd1b11f3e..47d588f7174 100644 --- a/tests/unit/single_controller/test_sampler_interface.py +++ b/tests/unit/single_controller/test_sampler_interface.py @@ -262,6 +262,33 @@ def test_capacity_does_not_shrink_when_the_gate_is_retuned(self): assert sampler.required_buffer_capacity(groups_per_step=4) == 16 + def test_retuning_the_gate_reopens_admission(self): + """The live window is what admit gates on, so widening it admits more.""" + s = InOrderSampler( + FakeBuffer(), max_lookahead_versions=1, warmup_lookahead_versions=3 + ) + + # dispatch_index starts at -1; window 1 admits the live batch and one + # lookahead batch against a trainer parked at 0, then blocks. + assert _run(s.admit(trainer_version_fn=lambda: 0)) == 0 + assert _run(s.admit(trainer_version_fn=lambda: 0)) == 1 + with pytest.raises(asyncio.TimeoutError): + _run(asyncio.wait_for(s.admit(trainer_version_fn=lambda: 0), timeout=0.05)) + + s.set_gate_window(3) + + assert _run(s.admit(trainer_version_fn=lambda: 0)) == 2 + assert _run(s.admit(trainer_version_fn=lambda: 0)) == 3 + with pytest.raises(asyncio.TimeoutError): + _run(asyncio.wait_for(s.admit(trainer_version_fn=lambda: 0), timeout=0.05)) + + # ...and shrinking it back closes the gate again: at dispatch_index 3 a + # trainer on version 2 is inside the warmup window but outside the steady one. + s.set_gate_window(1) + + with pytest.raises(asyncio.TimeoutError): + _run(asyncio.wait_for(s.admit(trainer_version_fn=lambda: 2), timeout=0.05)) + def test_set_gate_window_rejects_a_negative_window(self): sampler = InOrderSampler(FakeBuffer(), max_lookahead_versions=1) diff --git a/tests/unit/test_recipes_and_test_suites.py b/tests/unit/test_recipes_and_test_suites.py index 8739f7b135b..35f14c0c22b 100644 --- a/tests/unit/test_recipes_and_test_suites.py +++ b/tests/unit/test_recipes_and_test_suites.py @@ -256,7 +256,7 @@ def test_all_recipe_yamls_accounted_for_in_test_suites( ) -def test_nightly_compute_stays_below_4190_hours(nightly_test_suite, tracker): +def test_nightly_compute_stays_below_4149_hours(nightly_test_suite, tracker): command = f"DRYRUN=1 HF_HOME=... HF_DATASETS_CACHE=... CONTAINER= ACCOUNT= PARTITION= ./tools/launch {' '.join(nightly_test_suite)}" print(f"Running command: {command}") @@ -288,8 +288,8 @@ def test_nightly_compute_stays_below_4190_hours(nightly_test_suite, tracker): f"Last line of output was not as expected: '{last_line}'" ) total_gpu_hours = float(last_line.split(":")[-1].strip()) - assert total_gpu_hours <= 4190, ( - f"Total GPU hours exceeded 4190: {last_line}. We should revisit the test suites to reduce the total GPU hours." + assert total_gpu_hours <= 4149, ( + f"Total GPU hours exceeded 4149: {last_line}. We should revisit the test suites to reduce the total GPU hours." ) tracker.track("total_nightly_gpu_hours", total_gpu_hours) From ba8357a0e81a01403ec131e61396f1f1dbdde270 Mon Sep 17 00:00:00 2001 From: Yuki Huang Date: Tue, 25 Aug 2026 07:33:15 -0700 Subject: [PATCH 28/31] fix(sc): reject colocated generation on the SingleController path Signed-off-by: Yuki Huang --- .../single_controller_utils/config.py | 8 ++++ .../single_controller/test_checkpointing.py | 3 +- .../unit/single_controller/test_ppo_setup.py | 19 +++++++- .../test_resiliency_config.py | 10 ++++- tests/unit/single_controller/test_setup.py | 44 ++++++------------- .../single_controller/test_train_pump_e2e.py | 5 ++- 6 files changed, 52 insertions(+), 37 deletions(-) diff --git a/nemo_rl/algorithms/single_controller_utils/config.py b/nemo_rl/algorithms/single_controller_utils/config.py index 827f98dbe05..5e3facd0004 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/nemo_rl/algorithms/single_controller_utils/config.py @@ -743,6 +743,14 @@ def _validate_algo_settings(master_config: MasterConfig) -> None: "shaping. Disable them." ) + if master_config.policy["generation"]["colocated"]["enabled"]: + raise ValueError( + "The SingleController path requires " + "policy.generation.colocated.enabled=false: SC drives rollout via " + "RolloutManager.generate_and_push, which is only supported on the " + "disaggregated async engine." + ) + async_config = master_config.async_rl # Capacity is sized from the peak window whatever the algorithm, so an inert # setting still costs buffer and fails setup naming the wrong cause. diff --git a/tests/unit/single_controller/test_checkpointing.py b/tests/unit/single_controller/test_checkpointing.py index 804ab335753..481df1ee695 100644 --- a/tests/unit/single_controller/test_checkpointing.py +++ b/tests/unit/single_controller/test_checkpointing.py @@ -333,6 +333,7 @@ def _actor_master_config( policy={ # One optimizer.step per RL step: prompts * generations == gbs. "train_global_batch_size": num_prompts_per_step * 2, + "generation": {"colocated": {"enabled": False}}, }, loss_fn=ClippedPGLossConfig(), env={}, @@ -1036,7 +1037,7 @@ def _setup_master_config(checkpoint_dir: str) -> MasterConfig: "megatron_cfg": {"enabled": False}, "generation": { "backend": "vllm", - "colocated": {"enabled": True, "resources": {}}, + "colocated": {"enabled": False, "resources": {}}, }, }, loss_fn=ClippedPGLossConfig(), diff --git a/tests/unit/single_controller/test_ppo_setup.py b/tests/unit/single_controller/test_ppo_setup.py index 13e82b6b5b3..0eb2bd4d400 100644 --- a/tests/unit/single_controller/test_ppo_setup.py +++ b/tests/unit/single_controller/test_ppo_setup.py @@ -117,7 +117,7 @@ def _make_master_config( "megatron_cfg": {"enabled": megatron_enabled}, "generation": { "backend": "vllm", - "colocated": {"enabled": True, "resources": {}}, + "colocated": {"enabled": False, "resources": {}}, }, }, checkpointing={ @@ -367,6 +367,17 @@ def test_rejects_a_negative_reference_policy_kl_penalty(self): with pytest.raises(ValueError, match="must not be negative"): validate_single_controller_config(mc) + @pytest.mark.parametrize( + "make_config", [_ppo_master_config, _make_master_config], ids=["ppo", "grpo"] + ) + def test_rejects_colocated_generation(self, make_config): + """SC never sleeps generation, so the engine would hold GPUs all step.""" + mc = make_config() + mc.policy["generation"]["colocated"]["enabled"] = True + + with pytest.raises(ValueError, match="colocated.enabled=false"): + validate_single_controller_config(mc) + def test_grpo_is_free_to_use_any_sampler(self): mc = _make_master_config() mc.async_rl.sampler = WindowedSamplerConfig() @@ -544,7 +555,11 @@ def _cluster_config(mc: MasterConfig, *, colocated: bool, backend: str) -> Maste class TestTrainClusterSizesForTheCritic: - """The critic shares the training GPUs, so it needs its own worker-group slot.""" + """The critic shares the training GPUs, so it needs its own worker-group slot. + + The colocated cases call _build_clusters directly, bypassing the validator that + rejects colocated.enabled=true; they pin its arithmetic, not a usable mode. + """ @pytest.fixture def fake_cluster(self): diff --git a/tests/unit/single_controller/test_resiliency_config.py b/tests/unit/single_controller/test_resiliency_config.py index 3242ac92d98..e3408612ffa 100644 --- a/tests/unit/single_controller/test_resiliency_config.py +++ b/tests/unit/single_controller/test_resiliency_config.py @@ -64,7 +64,10 @@ def _master_config(*, num_prompts_per_step: int = 8, **async_kwargs) -> MasterCo num_generations_per_prompt=4, skip_reference_policy_logprobs_calculation=False, ), - policy={"train_global_batch_size": num_prompts_per_step * 4}, + policy={ + "train_global_batch_size": num_prompts_per_step * 4, + "generation": {"colocated": {"enabled": False}}, + }, # The last two are read only on the ready_first branch, which rejects a run # without them before it reaches anything under test here. loss_fn=SimpleNamespace( @@ -369,7 +372,10 @@ def _master_config(*, use_nemo_gym: bool, rollout_failure: dict) -> MasterConfig num_generations_per_prompt=4, skip_reference_policy_logprobs_calculation=False, ), - policy={"train_global_batch_size": 8}, + policy={ + "train_global_batch_size": 8, + "generation": {"colocated": {"enabled": False}}, + }, loss_fn=SimpleNamespace(reference_policy_kl_penalty=0), env={"should_use_nemo_gym": use_nemo_gym}, # Read by the metric_name check upstream #3429 added to this same diff --git a/tests/unit/single_controller/test_setup.py b/tests/unit/single_controller/test_setup.py index efec93d0a93..1a026b5b955 100644 --- a/tests/unit/single_controller/test_setup.py +++ b/tests/unit/single_controller/test_setup.py @@ -40,7 +40,7 @@ def _make_master_config( *, dp_enabled: bool = True, use_multiple_dataloader: bool = False, - colocated: bool = True, + colocated: bool = False, backend: str = "vllm", megatron_enabled: bool = False, env: dict | None = None, @@ -306,7 +306,7 @@ def test_ready_first_sampler_rejects_incompatible_loss_config( patched_factories["_build_clusters"].assert_not_called() def test_returns_actor_args(self, patched_factories): - mc = _make_master_config(colocated=True) + mc = _make_master_config() tokenizer = MagicMock(pad_token_id=0) actor_args, _ = setup_single_controller(mc, tokenizer) @@ -387,7 +387,7 @@ def test_rollout_manager_gets_no_effort_config_when_unset( assert call_kwargs["effort_config"] is None def test_router_replay_requires_routes_in_tq_buffer(self, patched_factories): - mc = _make_master_config(colocated=True) + mc = _make_master_config() mc.policy["router_replay"] = {"enabled": True} actor_args, _ = setup_single_controller(mc, MagicMock(pad_token_id=0)) @@ -501,7 +501,7 @@ def test_megatron_train_iters_not_set_when_disabled(self, patched_factories): def test_nemo_gym_wires_env_handle(self, patched_factories): """When should_use_nemo_gym is True the nemo-gym actor is spun up and stored.""" - mc = _make_master_config(colocated=True, backend="vllm") + mc = _make_master_config(backend="vllm") mc.policy["generation"]["model_name"] = "test-model" mc.policy["generation"]["stop_strings"] = None mc.policy["generation"]["stop_token_ids"] = None @@ -535,9 +535,9 @@ def test_nemo_gym_wires_env_handle(self, patched_factories): ) assert actor_args.env_handles["nemo_gym"] is fake_gym_actor - def test_setup_timing_populated_for_colocated_vllm(self, patched_factories): - """Colocated vLLM records gen+policy+collective+total+worker fields.""" - mc = _make_master_config(colocated=True, backend="vllm") + def test_setup_timing_populated_for_noncolocated_vllm(self, patched_factories): + """Non-colocated vLLM records every per-phase field.""" + mc = _make_master_config(colocated=False, backend="vllm") _, metrics = setup_single_controller(mc, MagicMock(pad_token_id=0)) @@ -552,23 +552,6 @@ def test_setup_timing_populated_for_colocated_vllm(self, patched_factories): value = getattr(metrics, field) assert value is not None, f"missing {field} on {metrics}" assert value >= 0 - # parallel_wall_time_s / parallel_init_enabled are grpo.py-only in the - # shared SetupTimingMetrics — SC does not emit them. - assert metrics.parallel_wall_time_s is None - assert metrics.parallel_init_enabled is None - # Reserve/load split is populated on the gym-on path only. - assert metrics.generation_init_reserve_time_s is None - assert metrics.generation_init_load_time_s is None - - def test_setup_timing_populated_for_noncolocated_vllm(self, patched_factories): - """Non-colocated vLLM records the same per-phase fields as colocated.""" - mc = _make_master_config(colocated=False, backend="vllm") - - _, metrics = setup_single_controller(mc, MagicMock(pad_token_id=0)) - - assert metrics.generation_init_time_s is not None - assert metrics.policy_init_time_s is not None - assert metrics.worker_setup_time_s is not None # parallel_wall_time_s / parallel_init_enabled are grpo.py-only. assert metrics.parallel_wall_time_s is None assert metrics.parallel_init_enabled is None @@ -578,7 +561,7 @@ def test_setup_timing_populated_for_noncolocated_vllm(self, patched_factories): def test_setup_timing_backend_agnostic_for_sglang(self, patched_factories): """SC uses the backend-agnostic generation_init_time_s regardless of backend.""" - mc = _make_master_config(colocated=True, backend="sglang") + mc = _make_master_config(backend="sglang") _, metrics = setup_single_controller(mc, MagicMock(pad_token_id=0)) @@ -586,7 +569,7 @@ def test_setup_timing_backend_agnostic_for_sglang(self, patched_factories): def test_nemo_gym_uses_deferred_vllm_load(self, patched_factories): """NeMo-Gym path reserves vLLM ports up-front and finishes the load afterwards.""" - mc = _make_master_config(colocated=True, backend="vllm") + mc = _make_master_config(backend="vllm") mc.policy["generation"]["model_name"] = "test-model" mc.policy["generation"]["stop_strings"] = None mc.policy["generation"]["stop_token_ids"] = None @@ -612,7 +595,7 @@ def test_nemo_gym_uses_deferred_vllm_load(self, patched_factories): def test_nemo_gym_records_timing_metrics(self, patched_factories): """NeMo-Gym path records per-phase timings (vllm/policy/gym/worker).""" - mc = _make_master_config(colocated=True, backend="vllm") + mc = _make_master_config(backend="vllm") mc.policy["generation"]["model_name"] = "test-model" mc.policy["generation"]["stop_strings"] = None mc.policy["generation"]["stop_token_ids"] = None @@ -665,9 +648,8 @@ def test_nemo_gym_noncolocated_finishes_deferred_load(self, patched_factories): assert metrics.generation_init_time_s is not None assert metrics.policy_init_time_s is not None - @pytest.mark.parametrize("colocated", [True, False]) def test_nemo_gym_generation_init_time_includes_reserve_time( - self, patched_factories, colocated + self, patched_factories ): """generation_init_time_s folds in the deferred-VllmGeneration reserve time. @@ -677,7 +659,7 @@ def test_nemo_gym_generation_init_time_includes_reserve_time( gym-on runs undercount generation setup by the worker-group span. The reserve/load split is also exposed for overlap analysis. """ - mc = _make_master_config(colocated=colocated, backend="vllm") + mc = _make_master_config(colocated=False, backend="vllm") mc.policy["generation"]["model_name"] = "test-model" mc.policy["generation"]["stop_strings"] = None mc.policy["generation"]["stop_token_ids"] = None @@ -708,7 +690,7 @@ def test_nemo_gym_generation_init_time_includes_reserve_time( @pytest.mark.parametrize("backend", ["sglang", "megatron"]) def test_nemo_gym_rejects_non_vllm_backend(self, patched_factories, backend): """SC nemo-gym wiring only supports vLLM; every other backend must raise.""" - mc = _make_master_config(colocated=True, backend=backend) + mc = _make_master_config(backend=backend) patched_factories["setup_response_data"].return_value = ( list(range(8)), None, diff --git a/tests/unit/single_controller/test_train_pump_e2e.py b/tests/unit/single_controller/test_train_pump_e2e.py index e68a0546643..403c91f63d0 100644 --- a/tests/unit/single_controller/test_train_pump_e2e.py +++ b/tests/unit/single_controller/test_train_pump_e2e.py @@ -302,7 +302,10 @@ def test_train_pump_drives_mcore_training_step( ) master_config = MasterConfig.model_construct( - policy={"train_global_batch_size": train_gbs}, + policy={ + "train_global_batch_size": train_gbs, + "generation": {"colocated": {"enabled": False}}, + }, # _sync_weights gates stale-abort on should_use_nemo_gym(env); empty # env -> native path (nemo_gym disabled). env={}, From 7d1f222efe50b911e25b7647ded4351a8e84346f Mon Sep 17 00:00:00 2001 From: Yuki Huang Date: Tue, 25 Aug 2026 08:23:53 -0700 Subject: [PATCH 29/31] fix(sc): move the exactly-one-algorithm-block check onto MasterConfig, guard ppo.async_ppo=null in run_ppo.py, make the advantage-estimator logprob args keyword-only, document the colocated branch as unreachable Signed-off-by: Yuki Huang --- docs/guides/single-controller.md | 2 +- examples/run_ppo.py | 6 ++ nemo_rl/algorithms/advantage_estimator.py | 2 + .../single_controller_utils/config.py | 29 +++++---- .../single_controller_utils/setup.py | 6 +- .../unit/single_controller/test_ppo_setup.py | 61 ++++++++++++++----- 6 files changed, 75 insertions(+), 31 deletions(-) diff --git a/docs/guides/single-controller.md b/docs/guides/single-controller.md index 9a6bacc9688..ace04ec811b 100644 --- a/docs/guides/single-controller.md +++ b/docs/guides/single-controller.md @@ -25,7 +25,7 @@ uv run examples/run_grpo_single_controller.py --config enabled: true ``` -2. **Enable vLLM async engine** and **disable colocated inference** (SC drives rollout via `RolloutManager.generate_and_push`, which is only supported on the disaggregated async engine): +2. **Enable vLLM async engine** and **disable colocated inference** (SC drives rollout via `RolloutManager.generate_and_push`, which is only supported on the disaggregated async engine; setup rejects `colocated.enabled: true`): ```yaml policy: diff --git a/examples/run_ppo.py b/examples/run_ppo.py index ae7e81ffa67..b7e8d355726 100644 --- a/examples/run_ppo.py +++ b/examples/run_ppo.py @@ -110,6 +110,12 @@ def main() -> None: config = MasterConfig(**config) print("Applied CLI overrides") + if config.ppo.async_ppo is None: + raise ValueError( + "ppo.async_ppo: null is only supported by run_grpo_single_controller.py; " + "the legacy PPO entrypoint requires the block." + ) + # Print config print("Final config:") pprint.pprint(config) diff --git a/nemo_rl/algorithms/advantage_estimator.py b/nemo_rl/algorithms/advantage_estimator.py index 46e855eee85..4ad4a27689a 100644 --- a/nemo_rl/algorithms/advantage_estimator.py +++ b/nemo_rl/algorithms/advantage_estimator.py @@ -227,6 +227,7 @@ def compute_advantage( prompt_ids, rewards, mask, + *, logprobs_policy=None, logprobs_reference=None, **kwargs, @@ -453,6 +454,7 @@ def compute_advantage( rewards, mask, values, + *, logprobs_policy=None, logprobs_reference=None, **kwargs, diff --git a/nemo_rl/algorithms/single_controller_utils/config.py b/nemo_rl/algorithms/single_controller_utils/config.py index 5e3facd0004..d336f53d787 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/nemo_rl/algorithms/single_controller_utils/config.py @@ -553,6 +553,20 @@ class MasterConfig(BaseModel, extra="allow"): data_plane: DataPlaneConfig async_rl: AsyncRLConfig + @model_validator(mode="after") + def validate_algorithm_block(self) -> "MasterConfig": + # Both are Optional so a PPO run can omit `grpo`; without this the + # entrypoint dereferences the absent block before validation runs. + if self.grpo is not None and self.ppo is not None: + raise ValueError( + "Only one algorithm block can be set, either `grpo` or `ppo`." + ) + if self.grpo is None and self.ppo is None: + raise ValueError( + "At least one algorithm block must be set, either `grpo` or `ppo`." + ) + return self + def is_ppo_run(master_config: MasterConfig) -> bool: """Whether this SingleController run trains a PPO critic alongside the policy. @@ -568,8 +582,8 @@ def is_ppo_run(master_config: MasterConfig) -> bool: def algo_config(master_config: MasterConfig) -> GRPOConfig | PPOConfig: """The active algorithm block: ``ppo`` on a PPO run, else ``grpo``. - Exactly one of the two is set; _validate_algo_settings checks that, and it - always runs at setup. + Exactly one of the two is set; MasterConfig.validate_algorithm_block checks + that at construction. """ if is_ppo_run(master_config): return master_config.ppo # type: ignore @@ -712,15 +726,6 @@ def _validate_algo_settings(master_config: MasterConfig) -> None: one a GRPO run carries and would never build. Plus the reward shaping and filtering knobs SC reads on neither path. """ - grpo = getattr(master_config, "grpo", None) - ppo = getattr(master_config, "ppo", None) - if grpo is not None and ppo is not None: - raise ValueError("Only one algorithm block can be set, either `grpo` or `ppo`.") - if grpo is None and ppo is None: - raise ValueError( - "At least one algorithm block must be set, either `grpo` or `ppo`." - ) - algo_cfg = algo_config(master_config) # SC reads none of these on either path, so an enabled one describes shaping @@ -876,8 +881,6 @@ def _validate_algo_settings(master_config: MasterConfig) -> None: def validate_single_controller_config(master_config: MasterConfig) -> None: """Validate cross-section SingleController constraints before setup.""" - # First: everything below reads the active algorithm block, which only exists - # once this has confirmed exactly one is set. _validate_algo_settings(master_config) async_config = master_config.async_rl diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py index 682220ac4ce..e56550d044b 100644 --- a/nemo_rl/algorithms/single_controller_utils/setup.py +++ b/nemo_rl/algorithms/single_controller_utils/setup.py @@ -138,7 +138,11 @@ class SingleControllerActorArgs: def _build_clusters( master_config: MasterConfig, ) -> tuple[RayVirtualCluster, RayVirtualCluster]: - """Allocate train + inference clusters; one shared cluster when colocated.""" + """Allocate train + inference clusters; one shared cluster when colocated. + + The colocated branch is unreachable on a real run -- validation rejects + colocated.enabled=true -- and is kept for when SC can support that mode. + """ cluster_config = master_config.cluster generation_config = master_config.policy["generation"] colocated = generation_config["colocated"]["enabled"] diff --git a/tests/unit/single_controller/test_ppo_setup.py b/tests/unit/single_controller/test_ppo_setup.py index 0eb2bd4d400..d63de950e1f 100644 --- a/tests/unit/single_controller/test_ppo_setup.py +++ b/tests/unit/single_controller/test_ppo_setup.py @@ -22,9 +22,12 @@ from __future__ import annotations +from pathlib import Path from unittest.mock import MagicMock, patch import pytest +from omegaconf import OmegaConf +from pydantic import ValidationError import nemo_rl.algorithms.single_controller_utils.setup as sc_setup_mod from nemo_rl.algorithms.advantage_estimator import ( @@ -155,6 +158,48 @@ def _ppo_master_config(**kwargs) -> MasterConfig: return _make_master_config(**kwargs) +class TestAlgorithmBlockValidator: + """Exactly-one-block, enforced at construction so no reader sees a bad pair.""" + + @staticmethod + def _resolved(name: str) -> dict: + from nemo_rl.utils.config import load_config, register_omegaconf_resolvers + + register_omegaconf_resolvers() + repo_root = Path(__file__).parents[3] + raw = load_config(repo_root / "examples/configs" / name) + resolved = OmegaConf.to_container(raw, resolve=True) + assert isinstance(resolved, dict) + return resolved + + @pytest.mark.parametrize( + "name", + [ + "grpo_math_1B_megatron_single_controller.yaml", + "ppo_math_1B_megatron_single_controller.yaml", + ], + ids=["grpo", "ppo"], + ) + def test_the_exemplars_still_validate(self, name): + MasterConfig(**self._resolved(name)) + + def test_rejects_a_config_with_no_algorithm_block(self): + resolved = self._resolved("grpo_math_1B_megatron_single_controller.yaml") + del resolved["grpo"] + + with pytest.raises(ValidationError, match="At least one algorithm block"): + MasterConfig(**resolved) + + def test_rejects_a_config_with_both_blocks(self): + resolved = self._resolved("ppo_math_1B_megatron_single_controller.yaml") + resolved["grpo"] = self._resolved( + "grpo_math_1B_megatron_single_controller.yaml" + )["grpo"] + + with pytest.raises(ValidationError, match="Only one algorithm block"): + MasterConfig(**resolved) + + class TestIsPPORun: @pytest.mark.parametrize( ("make_config", "expected"), @@ -188,22 +233,6 @@ class TestPPOValidation: def test_accepts_a_well_formed_ppo_config(self): validate_single_controller_config(_ppo_master_config()) - def test_rejects_a_config_with_neither_block(self): - """Also caught at setup; algo_config only asserts the invariant.""" - mc = _make_master_config() - mc.grpo = None - - with pytest.raises(ValueError, match="At least one algorithm block"): - validate_single_controller_config(mc) - - def test_rejects_a_config_with_both_blocks(self): - """Caught at setup rather than in algo_config, which runs on every access.""" - mc = _ppo_master_config() - mc.grpo = GRPOConfig.model_construct(**_STEP_CONFIG) - - with pytest.raises(ValueError, match="Only one algorithm block"): - validate_single_controller_config(mc) - @pytest.mark.parametrize("missing", ["value", "value_loss_fn"]) def test_rejects_ppo_without_its_critic_blocks(self, missing): mc = _ppo_master_config(**{missing: None}) From 6d303d38e29172e1579da1fdff51540bbd11e40d Mon Sep 17 00:00:00 2001 From: Yuki Huang Date: Tue, 25 Aug 2026 09:13:38 -0700 Subject: [PATCH 30/31] fix(sc): park the policy before the critic stages when neither logprob is required Signed-off-by: Yuki Huang --- nemo_rl/algorithms/single_controller.py | 5 ++ .../test_single_controller_actor.py | 76 ++++++++++++++++++- 2 files changed, 78 insertions(+), 3 deletions(-) diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index b079d3c8096..8edd5028d86 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -1023,6 +1023,11 @@ async def _train_pump(self) -> None: self._trainer.get_reference_policy_logprobs_from_meta, train_meta, ) + elif self._is_ppo: + # prepare_for_lp_inference is skipped here, and it is the only + # other call that parks the policy optimizer before the critic. + with self._timer.time("value_inference_prep"): + await asyncio.to_thread(self._trainer.offload_to_cpu) # Value model forward if self._is_ppo: diff --git a/tests/unit/single_controller/test_single_controller_actor.py b/tests/unit/single_controller/test_single_controller_actor.py index 54e001d8b53..9dc52b7f31a 100644 --- a/tests/unit/single_controller/test_single_controller_actor.py +++ b/tests/unit/single_controller/test_single_controller_actor.py @@ -906,6 +906,9 @@ def finish_inference(self) -> None: def prepare_for_training(self) -> None: self.calls.append("policy.prepare_for_training") + def offload_to_cpu(self) -> None: + self.calls.append("policy.offload_to_cpu") + class _EpochRecordingTrainer(_OrderRecordingTrainer): """Also records the optimizer-step lifecycle, which the epoch loop repeats.""" @@ -922,9 +925,6 @@ def finish_train_step(self) -> dict: self.calls.append("policy.finish_train_step") return {} - def offload_to_cpu(self) -> None: - self.calls.append("policy.offload_to_cpu") - class _NoOpDataPlane: def clear_samples(self, **kwargs) -> None: @@ -1265,6 +1265,31 @@ def test_train_pump_keeps_train_buffers_once_the_step_is_open(monkeypatch) -> No assert trainer.keep_train_buffers_calls == [False, True] +def test_train_pump_does_not_offload_the_policy_on_a_grpo_run(monkeypatch) -> None: + """The pre-critic offload is PPO-only: GRPO has no critic to make room for.""" + meta = KVBatchMeta( + partition_id="rollout_data", + task_name="train", + sample_ids=["sample-0"], + fields=[], + sequence_lengths=[1], + tags=[{"weight_version": 0}], + ) + calls: list[str] = [] + ctrl = _train_pump_controller(sampler=_ChunkedSampler(meta, chunks=2)) + ctrl._policy_logprobs_required = False + ctrl._reference_logprobs_required = False + ctrl._trainer = _OrderRecordingTrainer(calls) + ctrl._sync_weights = AsyncMock(return_value=1) + ctrl._logger = MagicMock() + monkeypatch.setattr(single_controller.ray, "cluster_resources", lambda: {}) + + asyncio.run(asyncio.wait_for(ctrl._train_pump(), timeout=1.0)) + + assert ctrl._train_steps == 1 + assert "policy.offload_to_cpu" not in calls + + # ── PPO ──────────────────────────────────────────────────────────────────── @@ -1380,6 +1405,49 @@ def test_train_pump_parks_the_policy_on_cpu_across_the_critic_stages( ] +def test_train_pump_parks_the_policy_when_neither_logprob_is_needed( + monkeypatch, +) -> None: + """No logprob means no prepare_for_lp_inference, so nothing else would park + the policy optimizer before the critic runs.""" + meta = _single_group_meta() + calls: list[str] = [] + ctrl, _ = _ppo_train_pump_controller( + sampler=_OneThenEmptySampler(meta), + value=_NoOpValue(calls=calls, prefix="critic."), + ) + ctrl._policy_logprobs_required = False + ctrl._reference_logprobs_required = False + ctrl._trainer = _OrderRecordingTrainer(calls) + ctrl._advantage_stage = AsyncMock(return_value=(meta, True)) + monkeypatch.setattr(single_controller.ray, "cluster_resources", lambda: {}) + + asyncio.run(asyncio.wait_for(ctrl._train_pump(), timeout=1.0)) + + assert "policy.prepare_for_lp_inference" not in calls + assert calls.index("policy.offload_to_cpu") < calls.index( + "critic.prepare_for_inference" + ) + + +def test_train_pump_does_not_double_offload_when_logprobs_run(monkeypatch) -> None: + """The logprob path already parks the optimizer, so the elif must not fire.""" + meta = _single_group_meta() + calls: list[str] = [] + ctrl, _ = _ppo_train_pump_controller( + sampler=_OneThenEmptySampler(meta), + value=_NoOpValue(calls=calls, prefix="critic."), + ) + ctrl._policy_logprobs_required = True + ctrl._trainer = _OrderRecordingTrainer(calls) + ctrl._advantage_stage = AsyncMock(return_value=(meta, True)) + monkeypatch.setattr(single_controller.ray, "cluster_resources", lambda: {}) + + asyncio.run(asyncio.wait_for(ctrl._train_pump(), timeout=1.0)) + + assert "policy.offload_to_cpu" not in calls + + def test_train_pump_logs_critic_metrics(monkeypatch) -> None: meta = _single_group_meta() ctrl, _ = _ppo_train_pump_controller(sampler=_OneThenEmptySampler(meta)) @@ -1485,6 +1553,8 @@ def test_train_pump_offloads_the_policy_between_ppo_epochs(monkeypatch) -> None: asyncio.run(asyncio.wait_for(ctrl._train_pump(), timeout=1.0)) assert calls == [ + # Neither logprob is required here, so the policy is parked up front. + "policy.offload_to_cpu", "policy.finish_inference", "critic.prepare_for_inference", "critic.get_values_from_meta", From b60aee00e855df7a2855c4b6e20db1bbab6346fa Mon Sep 17 00:00:00 2001 From: Yuki Huang Date: Tue, 25 Aug 2026 17:26:40 -0700 Subject: [PATCH 31/31] fix unit test Signed-off-by: Yuki Huang --- .../single_controller/test_rollout_pump.py | 5 +++- .../test_single_controller_actor.py | 25 +++++++++++++++---- 2 files changed, 24 insertions(+), 6 deletions(-) diff --git a/tests/unit/single_controller/test_rollout_pump.py b/tests/unit/single_controller/test_rollout_pump.py index 94fb220bb6e..97f25574e34 100644 --- a/tests/unit/single_controller/test_rollout_pump.py +++ b/tests/unit/single_controller/test_rollout_pump.py @@ -1001,7 +1001,10 @@ def test_rollout_pump_writes_expected_tq_data( dp_adapter = _SyncDPAdapter(tq_actor) master_config = MasterConfig.model_construct( - policy={"train_global_batch_size": expected_samples}, + policy={ + "train_global_batch_size": expected_samples, + "generation": {"colocated": {"enabled": False}}, + }, grpo=GRPOConfig.model_construct( num_prompts_per_step=num_prompts, num_generations_per_prompt=num_generations, diff --git a/tests/unit/single_controller/test_single_controller_actor.py b/tests/unit/single_controller/test_single_controller_actor.py index 9dc52b7f31a..739250d0a79 100644 --- a/tests/unit/single_controller/test_single_controller_actor.py +++ b/tests/unit/single_controller/test_single_controller_actor.py @@ -61,7 +61,10 @@ def _checkpointing_config(tmp_path) -> dict: def _grpo_master_config(tmp_path) -> MasterConfig: """A minimal GRPO MasterConfig the real __init__ accepts.""" return MasterConfig.model_construct( - policy={"train_global_batch_size": 8}, + policy={ + "train_global_batch_size": 8, + "generation": {"colocated": {"enabled": False}}, + }, grpo=GRPOConfig.model_construct( num_prompts_per_step=2, num_generations_per_prompt=4, @@ -114,7 +117,10 @@ def _init_controller(master_config, actor_args): def test_rejects_multiple_optimizer_steps_per_rl_step(monkeypatch) -> None: monkeypatch.setattr(single_controller, "Logger", lambda _: object()) master_config = MasterConfig.model_construct( - policy={"train_global_batch_size": 4}, + policy={ + "train_global_batch_size": 4, + "generation": {"colocated": {"enabled": False}}, + }, grpo=GRPOConfig.model_construct( num_prompts_per_step=2, num_generations_per_prompt=4, @@ -164,7 +170,10 @@ def test_logs_hyperparameters_and_concrete_weight_synchronizer( logger = MagicMock() monkeypatch.setattr(single_controller, "Logger", lambda _: logger) master_config = MasterConfig.model_construct( - policy={"train_global_batch_size": 8}, + policy={ + "train_global_batch_size": 8, + "generation": {"colocated": {"enabled": False}}, + }, grpo=GRPOConfig.model_construct( num_prompts_per_step=2, num_generations_per_prompt=4, @@ -230,7 +239,10 @@ def test_reference_logprobs_required_only_when_kl_enabled( """KL-disabled SingleController runs do not request reference logprobs.""" monkeypatch.setattr(single_controller, "Logger", lambda _: MagicMock()) master_config = MasterConfig.model_construct( - policy={"train_global_batch_size": 8}, + policy={ + "train_global_batch_size": 8, + "generation": {"colocated": {"enabled": False}}, + }, grpo=GRPOConfig.model_construct( num_prompts_per_step=2, num_generations_per_prompt=4, @@ -284,7 +296,10 @@ def test_logs_setup_timing_metrics(monkeypatch, tmp_path) -> None: logger = MagicMock() monkeypatch.setattr(single_controller, "Logger", lambda _: logger) master_config = MasterConfig.model_construct( - policy={"train_global_batch_size": 8}, + policy={ + "train_global_batch_size": 8, + "generation": {"colocated": {"enabled": False}}, + }, grpo=GRPOConfig.model_construct( num_prompts_per_step=2, num_generations_per_prompt=4,