diff --git a/docs/about/algorithms/mopd.md b/docs/about/algorithms/mopd.md index 25bc53f380c..3d0b232d5e9 100644 --- a/docs/about/algorithms/mopd.md +++ b/docs/about/algorithms/mopd.md @@ -125,9 +125,26 @@ policy and generation nodes. ## Running MOPD -MOPD collects rollouts through NeMo Gym, so use the NeMo Gym GRPO entrypoint -with an MOPD recipe. The checked-in recipe uses placeholder dataset paths; -override them for your local data: +MOPD collects rollouts through NeMo Gym and supports both the legacy async GRPO +runtime and the Single-Controller runtime. The checked-in recipes use +placeholder dataset paths; override them for your local data. + +### Single-Controller text path + +The Single-Controller path moves rollout and teacher-logprob tensors through +TransferQueue. It currently supports text-only MOPD rollouts: + +```sh +uv run examples/run_grpo_single_controller.py \ + --config examples/configs/recipes/llm/mopd-qwen3-1.7b-3n8g-megatron-pack-single-controller.yaml \ + data.train.data_path=/path/to/train.jsonl \ + data.validation.data_path=/path/to/val.jsonl +``` + +See [Train with Single-Controller](../../guides/single-controller.md) for the +runtime's configuration and architecture. + +### Legacy async GRPO path ```sh uv run examples/nemo_gym/run_grpo_nemo_gym.py \ @@ -136,10 +153,10 @@ uv run examples/nemo_gym/run_grpo_nemo_gym.py \ data.validation.data_path=/path/to/val.jsonl ``` -The reference recipe self-distills `Qwen/Qwen3-1.7B` (student == teacher) across -3 nodes (1 policy + 1 vLLM + 1 teacher) with sequence packing enabled. Because -student and teacher are identical, the OPD loss stays near zero — it is a -correctness smoke test, not a demonstration of distillation gains. +Both reference recipes self-distill `Qwen/Qwen3-1.7B` (student == teacher) +across 3 nodes (1 policy + 1 vLLM + 1 teacher) with sequence packing enabled. +Because student and teacher are identical, the OPD loss stays near zero — it is +a correctness smoke test, not a demonstration of distillation gains. ## References diff --git a/docs/guides/single-controller.md b/docs/guides/single-controller.md index ace04ec811b..401b7dc78d8 100644 --- a/docs/guides/single-controller.md +++ b/docs/guides/single-controller.md @@ -175,6 +175,9 @@ SC reads its async knobs from `async_rl:` and **requires `grpo.async_grpo: null` The SC path is still under active development. Feature gaps are tracked in [issue #2625](https://github.com/NVIDIA-NeMo/RL/issues/2625). Notable items: +- Multi-Teacher On-Policy Distillation (MOPD) is supported for text-only NeMo + Gym rollouts; multimodal/VLM MOPD is not yet supported. See + [Multi-Teacher On-Policy Distillation](../about/algorithms/mopd.md#running-mopd). - 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`). diff --git a/examples/configs/recipes/llm/mopd-qwen3-1.7b-3n8g-megatron-pack-single-controller.yaml b/examples/configs/recipes/llm/mopd-qwen3-1.7b-3n8g-megatron-pack-single-controller.yaml new file mode 100644 index 00000000000..fe7bdf3501f --- /dev/null +++ b/examples/configs/recipes/llm/mopd-qwen3-1.7b-3n8g-megatron-pack-single-controller.yaml @@ -0,0 +1,29 @@ +defaults: ./mopd-qwen3-1.7b-3n8g-megatron-pack.yaml + +# SingleController consumes async_rl and requires the legacy async block to be absent. +grpo: + async_grpo: null + val_period: 0 + # OPD needs the student pass, but not a separate frozen reference-policy pass. + skip_reference_policy_logprobs_calculation: true + +async_rl: + sampler: + name: in_order + max_lookahead_versions: 1 + min_groups_for_streaming_train: ${grpo.num_prompts_per_step} + max_inflight_prompts: ${grpo.num_prompts_per_step} + max_buffered_rollouts: 64 + +data_plane: + enabled: true + +checkpointing: + checkpoint_dir: results/mopd-qwen3-1.7b-3n8g-megatron-pack-single-controller + metric_name: null + +logger: + wandb: + name: mopd-qwen3-1.7b-3n8g-megatron-pack-single-controller + mlflow: + run_name: mopd-qwen3-1.7b-3n8g-megatron-pack-single-controller diff --git a/examples/run_grpo_single_controller.py b/examples/run_grpo_single_controller.py index b54cb898e15..3a151016a2a 100644 --- a/examples/run_grpo_single_controller.py +++ b/examples/run_grpo_single_controller.py @@ -167,6 +167,13 @@ def main() -> None: except Exception as kill_error: print(f"Env {env_name!r} kill failed: {kill_error}") + teacher_worker_groups = getattr(actor_args, "teacher_worker_groups", None) or {} + for teacher_alias, teacher in teacher_worker_groups.items(): + try: + teacher.shutdown() + except Exception as e: + print(f"Teacher {teacher_alias!r} shutdown failed: {e}") + for resource_name, resource in ( ("Generation", actor_args.gen_handle), ("Trainer", actor_args.trainer_handle), diff --git a/nemo_rl/algorithms/async_utils/replay_buffer.py b/nemo_rl/algorithms/async_utils/replay_buffer.py index 6b7112a9949..f5061aa75ef 100644 --- a/nemo_rl/algorithms/async_utils/replay_buffer.py +++ b/nemo_rl/algorithms/async_utils/replay_buffer.py @@ -19,7 +19,7 @@ import uuid from collections import Counter from collections.abc import Mapping -from typing import Any, Iterable, Optional +from typing import Any, Awaitable, Callable, Iterable, Optional import ray import torch @@ -37,6 +37,10 @@ from nemo_rl.utils.r3_trace import trace_rollout_payload +class PostWriteEnrichmentError(RuntimeError): + """A rollout reached TQ but failed in required post-write processing.""" + + # Classes with @ray.remote can't be inherited from, so we split the implementation out. class ReplayBufferImpl(ReplayBufferProtocol): """Replay buffer storing per-prompt groups. @@ -774,6 +778,16 @@ def __init__( self.target_step_list: list[Optional[int]] = [] self.ready_list: list[bool] = [] self._group_ids: list[str] = [] + self._post_write_enricher: Optional[ + Callable[[KVBatchMeta, PromptGroupRecord], Awaitable[KVBatchMeta]] + ] = None + + def set_post_write_enricher( + self, + enricher: Callable[[KVBatchMeta, PromptGroupRecord], Awaitable[KVBatchMeta]], + ) -> None: + """Install the required enrichment stage run before slots become ready.""" + self._post_write_enricher = enricher def reserve( self, @@ -864,6 +878,14 @@ async def commit( tags=[dict(t) for t in tags], ) + if self._post_write_enricher is not None: + try: + meta = await self._post_write_enricher(meta, record) + except Exception as error: + raise PostWriteEnrichmentError( + f"post-write enrichment failed for group_id={group_id!r}" + ) from error + idx = self._group_ids.index(group_id) self.meta_list[idx] = meta self.end_weight_list[idx] = end_weight_version diff --git a/nemo_rl/algorithms/metric_utils.py b/nemo_rl/algorithms/metric_utils.py index 2425eed0826..8e1328bc7d5 100644 --- a/nemo_rl/algorithms/metric_utils.py +++ b/nemo_rl/algorithms/metric_utils.py @@ -40,7 +40,8 @@ class SetupTimingMetrics: parallel_wall_time_s: Optional[float] = None parallel_init_enabled: Optional[float] = None - # grpo-only phases (non-colocated OPD teachers, sparse refit, checkpoint-engine). + # Optional setup phases. OPD teacher timings are shared by legacy GRPO and SC; + # sparse refit and checkpoint-engine timings remain legacy-GRPO-only. teacher_reservation_time_s: Optional[float] = None teacher_model_init_time_s: Optional[float] = None teacher_init_time_s: Optional[float] = None diff --git a/nemo_rl/algorithms/opd.py b/nemo_rl/algorithms/opd.py index 8fd73bf9994..d99a64b659e 100644 --- a/nemo_rl/algorithms/opd.py +++ b/nemo_rl/algorithms/opd.py @@ -21,15 +21,23 @@ from __future__ import annotations +import asyncio +import time +import uuid from typing import Any, Optional import ray +import torch from pydantic import BaseModel, Field +from nemo_rl.data_plane.column_io import read_columns, write_columns +from nemo_rl.data_plane.interfaces import DataPlaneClient, KVBatchMeta +from nemo_rl.data_plane.schema import TEACHER_LP_FIELDS from nemo_rl.distributed.virtual_cluster import ( RayVirtualCluster, prepare_segment_topology, ) +from nemo_rl.experience.interfaces import PromptGroupRecord # --------------------------------------------------------------------------- # Config schemas @@ -200,6 +208,241 @@ def get_teacher_routing_metrics( } +class TQTeacherLogprobCoordinator: + """Enrich SingleController rollout rows with teacher logprobs through TQ. + + A prompt group is already present in the data plane when :meth:`enrich` is + called. The coordinator sends its metadata to the inference-only teacher; + teacher workers fetch their own DP shards and write the token column back + under the same sample IDs. The controller materializes only the final row + when unique temporary rows are needed for DP divisibility. Calls targeting + the same deduplicated teacher group are serialized to keep NCCL collective + ordering identical across that group's workers; distinct teachers remain + independent and can run concurrently. + """ + + teacher_logprobs_field = "teacher_reference_logprobs" + + def __init__( + self, + *, + dp_client: DataPlaneClient, + teacher_worker_groups: dict[str, Any], + alias_to_group_alias: dict[str, str], + on_policy_distillation_cfg: dict[str, Any], + ) -> None: + if not teacher_worker_groups: + raise ValueError( + "TQTeacherLogprobCoordinator requires at least one teacher worker group" + ) + self._dp_client = dp_client + self._teacher_worker_groups = dict(teacher_worker_groups) + self._alias_to_group_alias = dict(alias_to_group_alias) + self._opd_cfg = dict(on_policy_distillation_cfg) + # Physical (deduplicated) groups own locks, not routing aliases. Two + # aliases sharing one checkpoint therefore share one collective FIFO. + self._teacher_locks = { + group_alias: asyncio.Lock() for group_alias in self._teacher_worker_groups + } + self._teacher_batches = 0 + self._teacher_samples = 0 + self._teacher_logprob_time_s = 0.0 + self._teacher_inference_time_s = 0.0 + self._teacher_lock_wait_time_s = 0.0 + self._aliases_seen: set[str] = set() + + def _resolve_teacher(self, record: PromptGroupRecord) -> tuple[str, str]: + extra_env_info = record.extra_env_info + agent_ref = ( + extra_env_info.get("agent_ref") + if isinstance(extra_env_info, dict) + else None + ) + if not isinstance(agent_ref, dict): + raise ValueError( + "on_policy_distillation is enabled but this prompt group has no " + "extra_env_info['agent_ref'] mapping to route it to a teacher. " + "SingleController MOPD requires the NeMo-Gym rollout path, and " + "regenerating this prompt cannot repair missing routing metadata." + ) + + teacher_model_by_agent_name = dict( + self._opd_cfg.get("teacher_model_by_agent_name", {}) + ) + alias = resolve_reference_aliases( + [agent_ref], + teacher_model_by_agent_name, + default_teacher_alias=self._opd_cfg.get("default_teacher_alias"), + strict_agent_name_match=bool( + self._opd_cfg.get("strict_agent_name_match", False) + ), + )[0] + group_alias = self._alias_to_group_alias.get(alias, alias) + if group_alias not in self._teacher_worker_groups: + raise ValueError( + f"Teacher alias {alias!r} resolved to unavailable worker group " + f"{group_alias!r}; available groups: " + f"{sorted(self._teacher_worker_groups)}" + ) + return alias, group_alias + + def _enrich_sync( + self, + meta: KVBatchMeta, + group_alias: str, + ) -> float: + """Run the blocking TQ-read, teacher inference, and TQ-write sequence.""" + if not meta.sequence_lengths: + raise ValueError("MOPD teacher enrichment requires sequence_lengths") + + teacher = self._teacher_worker_groups[group_alias] + dp_size = teacher.sharding_annotations.get_axis_size("data_parallel") + if dp_size <= 0: + raise ValueError( + f"Teacher {group_alias!r} has invalid data-parallel size {dp_size}" + ) + actual_batch_size = meta.size + remainder = actual_batch_size % dp_size + padded_meta = meta + temporary_sample_ids: list[str] = [] + try: + if remainder: + pad_count = dp_size - remainder + source_meta = meta.slice(actual_batch_size - 1, actual_batch_size) + source_data = read_columns( + self._dp_client, + source_meta, + select_fields=TEACHER_LP_FIELDS, + pad_value_dict={"input_ids": 0}, + ) + input_ids = source_data["input_ids"] + input_lengths = source_data["input_lengths"] + if not isinstance(input_ids, torch.Tensor) or not isinstance( + input_lengths, torch.Tensor + ): + raise TypeError("MOPD teacher padding inputs must be tensors") + temporary_prefix = uuid.uuid4().hex + temporary_sample_ids = [ + f"{meta.sample_ids[-1]}__teacher_pad_{temporary_prefix}_{index}" + for index in range(pad_count) + ] + pad_meta = KVBatchMeta( + partition_id=meta.partition_id, + task_name=meta.task_name, + sample_ids=temporary_sample_ids, + fields=list(TEACHER_LP_FIELDS), + sequence_lengths=[meta.sequence_lengths[-1]] * pad_count, + ) + # Keep this write inside the cleanup lifetime. A backend may + # write only some rows before reporting failure. + write_columns( + self._dp_client, + pad_meta, + fields={ + "input_ids": input_ids.expand(pad_count, *input_ids.shape[1:]), + "input_lengths": input_lengths.expand( + pad_count, *input_lengths.shape[1:] + ), + }, + ) + padded_meta = meta.concat(pad_meta) + + inference_started_at = time.perf_counter() + teacher.get_logprobs_from_meta(padded_meta) + inference_finished_at = time.perf_counter() + except BaseException as enrichment_error: + if temporary_sample_ids: + try: + self._dp_client.clear_samples( + sample_ids=temporary_sample_ids, + partition_id=meta.partition_id, + ) + except BaseException as cleanup_error: + raise BaseExceptionGroup( + f"teacher enrichment and temporary-row cleanup both failed " + f"for group {group_alias!r}", + [enrichment_error, cleanup_error], + ) + raise + if temporary_sample_ids: + self._dp_client.clear_samples( + sample_ids=temporary_sample_ids, + partition_id=meta.partition_id, + ) + return inference_finished_at - inference_started_at + + async def enrich( + self, + meta: KVBatchMeta, + record: PromptGroupRecord, + ) -> KVBatchMeta: + """Write teacher logprobs before the replay-buffer slot becomes ready.""" + alias, group_alias = self._resolve_teacher(record) + started_at = time.perf_counter() + lock_started_at = time.perf_counter() + # Wait for a physical teacher without occupying a default-executor + # thread. Only the active inference consumes a thread-pool slot. + async with self._teacher_locks[group_alias]: + lock_wait_s = time.perf_counter() - lock_started_at + # asyncio cannot cancel a running thread. Shield and explicitly + # drain it so replay-buffer rollback never clears rows while teacher + # workers are still fetching from or writing to those rows. + enrichment_task = asyncio.create_task( + asyncio.to_thread(self._enrich_sync, meta, group_alias) + ) + try: + inference_time_s = await asyncio.shield(enrichment_task) + except asyncio.CancelledError: + try: + await enrichment_task + except BaseException as drain_error: + raise asyncio.CancelledError( + f"cancelled while draining teacher enrichment for {group_alias!r}" + ) from drain_error + raise + total_time_s = time.perf_counter() - started_at + + self._teacher_batches += 1 + self._teacher_samples += meta.size + self._teacher_logprob_time_s += total_time_s + self._teacher_inference_time_s += inference_time_s + self._teacher_lock_wait_time_s += lock_wait_s + self._aliases_seen.add(alias) + print( + f"[teacher_logprob] group={group_alias} samples={meta.size} " + f"lock_wait={lock_wait_s:.2f}s inference={inference_time_s:.2f}s " + f"total={total_time_s:.2f}s", + flush=True, + ) + return meta.with_fields([self.teacher_logprobs_field]) + + def drain_metrics(self) -> dict[str, float]: + """Return and reset teacher activity accumulated since the last drain.""" + metrics = { + "on_policy_distillation/teacher_batches": float(self._teacher_batches), + "on_policy_distillation/teacher_samples": float(self._teacher_samples), + "on_policy_distillation/teacher_logprob_time_s": self._teacher_logprob_time_s, + "on_policy_distillation/teacher_inference_time_s": self._teacher_inference_time_s, + "on_policy_distillation/teacher_lock_wait_time_s": self._teacher_lock_wait_time_s, + } + if self._teacher_batches: + # Cardinality describes what ran. On an idle step, zero reads as + # "zero teacher models" rather than "no teacher activity." + metrics.update( + get_teacher_routing_metrics( + sorted(self._aliases_seen), + self._opd_cfg["teacher_model_by_agent_name"], + ) + ) + self._teacher_batches = 0 + self._teacher_samples = 0 + self._teacher_logprob_time_s = 0.0 + self._teacher_inference_time_s = 0.0 + self._teacher_lock_wait_time_s = 0.0 + self._aliases_seen.clear() + return metrics + + # --------------------------------------------------------------------------- # Setup helper — teacher worker group creation # --------------------------------------------------------------------------- diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 8edd5028d86..c3e910998a7 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -51,6 +51,7 @@ import ray import torch +from nemo_rl.algorithms import opd as opd_module from nemo_rl.algorithms.async_utils.staleness_sampler import create_sampler from nemo_rl.algorithms.grpo import ( GRPOSaveState, @@ -97,6 +98,22 @@ log = logging.getLogger(__name__) +def _pooled_opd_metrics( + stat_sum: float, stat_sumsq: float, count: int +) -> dict[str, float]: + """Compute whole-step OPD metrics from exact pooled sufficient statistics.""" + if count <= 0: + return {} + mean = stat_sum / count + # OPDAdvantageEstimator uses torch.std's default unbiased estimator. + variance = (stat_sumsq - count * mean * mean) / (count - 1) if count > 1 else 0.0 + return { + "on_policy_distillation/teacher_student_logprob_gap_mean": mean, + "on_policy_distillation/adv_mean": mean, + "on_policy_distillation/adv_std": math.sqrt(max(variance, 0.0)), + } + + @ray.remote(num_cpus=1, num_gpus=0) # pragma: no cover class SingleControllerActor: """CPU-only Ray actor that orchestrates the RL training loop. @@ -149,6 +166,7 @@ def __init__( master_config.loss_fn.reference_policy_kl_penalty > 0 and not self._algo_cfg.skip_reference_policy_logprobs_calculation ) + self._teacher_logprobs_required = opd_module.is_opd_enabled(master_config) self._dp_client = actor_args.dp_client self._gen: Generation = actor_args.gen_handle self._trainer: TQPolicy = actor_args.trainer_handle @@ -176,6 +194,21 @@ 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) + teacher_worker_groups = getattr(actor_args, "teacher_worker_groups", None) or {} + if teacher_worker_groups: + self._teacher_coordinator: Optional[ + opd_module.TQTeacherLogprobCoordinator + ] = opd_module.TQTeacherLogprobCoordinator( + dp_client=self._dp_client, + teacher_worker_groups=teacher_worker_groups, + alias_to_group_alias=( + getattr(actor_args, "alias_to_group_alias", None) or {} + ), + on_policy_distillation_cfg=opd_module._opd_cfg(master_config), + ) + self._buffer.set_post_write_enricher(self._teacher_coordinator.enrich) + else: + self._teacher_coordinator = None # Built here, not on the driver: Logger backends (wandb/tb/...) hold # _thread.lock that Ray can't cloudpickle into the actor. @@ -280,6 +313,9 @@ def __init__( "sequence_lengths": [], "seq_logprob_error_metrics": [], } + self._opd_stat_sum = 0.0 + self._opd_stat_sumsq = 0.0 + self._opd_stat_count = 0 # Seeded here rather than in run(): on resume _trainer_version is the # checkpoint's step, so a run resuming mid-warmup needs the widened @@ -1205,6 +1241,18 @@ async def _train_pump(self) -> None: reduce_advantage_pump_metrics(**self._step_log_dict) ) self._step_log_dict = {k: [] for k in self._step_log_dict} + step_metrics.update( + _pooled_opd_metrics( + self._opd_stat_sum, + self._opd_stat_sumsq, + self._opd_stat_count, + ) + ) + self._opd_stat_sum = 0.0 + self._opd_stat_sumsq = 0.0 + self._opd_stat_count = 0 + if self._teacher_coordinator is not None: + step_metrics.update(self._teacher_coordinator.drain_metrics()) self._trainer_version += 1 self._train_steps += 1 @@ -1980,15 +2028,21 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> tuple[KVBatchMeta, bool]: kwargs: dict[str, torch.Tensor] = {} if self._policy_logprobs_required: - kwargs["logprobs_policy"] = tensor_field( - data, - adv_cfg.policy_logprobs_field, - ) + policy_logprobs = tensor_field(data, adv_cfg.policy_logprobs_field) + if self._teacher_logprobs_required: + kwargs["prev_logprobs"] = policy_logprobs + else: + kwargs["logprobs_policy"] = policy_logprobs if self._reference_logprobs_required: kwargs["logprobs_reference"] = tensor_field( data, adv_cfg.reference_logprobs_field, ) + if self._teacher_logprobs_required: + kwargs["teacher_logprobs"] = tensor_field( + data, + adv_cfg.teacher_logprobs_field, + ) if self._is_ppo: kwargs["values"] = tensor_field(data, adv_cfg.values_field) @@ -2020,6 +2074,11 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> tuple[KVBatchMeta, bool]: self._step_log_dict["masked_advantages"].append( response_advantages.detach().cpu() ) + if self._teacher_logprobs_required: + valid = response_advantages.detach().double() + self._opd_stat_sum += float(valid.sum()) + self._opd_stat_sumsq += float((valid * valid).sum()) + self._opd_stat_count += int(valid.numel()) fields_to_put = {adv_cfg.output_field: advantages} if seq_logprob_error_threshold is not None: @@ -2057,6 +2116,8 @@ 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._teacher_logprobs_required: + fields.append(adv_cfg.teacher_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 d336f53d787..ce28051b297 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/nemo_rl/algorithms/single_controller_utils/config.py @@ -28,6 +28,7 @@ model_validator, ) +from nemo_rl.algorithms import opd as opd_module from nemo_rl.algorithms.async_utils.staleness_sampler import ( InOrderSamplerConfig, ReadyFirstSamplerConfig, @@ -37,6 +38,7 @@ 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.opd import OnPolicyDistillationConfig from nemo_rl.algorithms.ppo import PPOConfig from nemo_rl.data import DataConfig from nemo_rl.data_plane.interfaces import DataPlaneConfig @@ -552,6 +554,7 @@ class MasterConfig(BaseModel, extra="allow"): checkpointing: CheckpointingConfig data_plane: DataPlaneConfig async_rl: AsyncRLConfig + on_policy_distillation: Optional[OnPolicyDistillationConfig] = None @model_validator(mode="after") def validate_algorithm_block(self) -> "MasterConfig": @@ -973,6 +976,46 @@ def validate_single_controller_config(master_config: MasterConfig) -> None: "loss_fn.reference_policy_kl_penalty=0." ) + # ``env`` is required in production configs, but model_construct-based unit + # configs can omit it. Only apply rollout-path validation when it is present. + env_config = getattr(master_config, "env", None) + + opd_enabled = opd_module.is_opd_enabled(master_config) + if opd_enabled and is_ppo_run(master_config): + raise ValueError( + "on_policy_distillation is only supported with the `grpo` algorithm block." + ) + if algo_cfg.adv_estimator.name == "opd" and not opd_enabled: + raise ValueError( + "grpo.adv_estimator.name='opd' requires " + "on_policy_distillation.enabled=true." + ) + if opd_enabled: + opd_config = master_config.on_policy_distillation + assert opd_config is not None + if algo_cfg.adv_estimator.name != "opd": + raise ValueError( + "on_policy_distillation.enabled=true requires " + "grpo.adv_estimator.name='opd'." + ) + if not opd_module.is_non_colocated_teachers_enabled(master_config): + raise ValueError( + "SingleController MOPD currently requires " + "on_policy_distillation.non_colocated_teachers.enabled=true." + ) + if env_config is not None and not bool(env_config.get("should_use_nemo_gym")): + raise ValueError( + "on_policy_distillation requires env.should_use_nemo_gym=true: " + "teacher routing keys off the gym rollout path's per-agent " + "agent_ref." + ) + if not opd_config.teacher_model_by_agent_name: + raise ValueError( + "on_policy_distillation.teacher_model_by_agent_name must contain " + "at least one teacher mapping." + ) + opd_module.assert_prev_logprobs_available(master_config) + if ( reference_policy_kl_penalty == 0 and not algo_cfg.skip_reference_policy_logprobs_calculation @@ -990,13 +1033,6 @@ def validate_single_controller_config(master_config: MasterConfig) -> None: # wrong-path block is still a silent no-op, which is the failure this whole # restructure exists to remove. Only a check at setup actually closes it. # - # ``env`` is a required field, so a run built the production way -- through - # MasterConfig(**cfg), which validates -- always carries it. model_construct - # skips validation and only fills fields that have defaults, so a config - # assembled that way can genuinely lack the attribute, and without it the - # rollout path is unknowable. This check only reads it to decide which half of - # the block is inert, so skip rather than fail a construction over it. - env_config = getattr(master_config, "env", None) if env_config is not None: use_nemo_gym = bool(env_config.get("should_use_nemo_gym")) unused_name = "native" if use_nemo_gym else "nemo_gym" @@ -1034,6 +1070,7 @@ class AdvantageConfig: policy_logprobs_field: str = "prev_logprobs" generation_logprobs_field: str = "generation_logprobs" reference_logprobs_field: str = "reference_policy_logprobs" + teacher_logprobs_field: str = "teacher_reference_logprobs" # PPO only: the critic's pre-update prediction (input) and GAE's # regression target for it (output). values_field: str = "values" diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py index e56550d044b..313c3a94ab6 100644 --- a/nemo_rl/algorithms/single_controller_utils/setup.py +++ b/nemo_rl/algorithms/single_controller_utils/setup.py @@ -34,6 +34,7 @@ from transformers import AutoProcessor from transformers.tokenization_utils_base import PreTrainedTokenizerBase +from nemo_rl.algorithms import opd as opd_module from nemo_rl.algorithms.async_utils.replay_buffer import TQReplayBuffer from nemo_rl.algorithms.grpo import ( GRPOSaveState, @@ -59,10 +60,15 @@ from nemo_rl.data.collate_fn import rl_collate_fn from nemo_rl.data.utils import load_dataloader_state, setup_response_data from nemo_rl.data_plane import DataPlaneClient, build_data_plane_client +from nemo_rl.data_plane.schema import ( + SC_ROLLOUT_SCHEMA_FIELDS, + fields_with_optional_routed_experts, +) from nemo_rl.distributed.virtual_cluster import ( RayVirtualCluster, _get_free_port_local, _get_node_ip_local, + prepare_segment_topology, ) from nemo_rl.environments.interfaces import EnvironmentInterface from nemo_rl.environments.nemo_gym import should_use_nemo_gym, spinup_nemo_gym_actor @@ -129,16 +135,49 @@ 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 + # Populated only for text MOPD. Aliases may outnumber worker groups when + # multiple agents share one deduplicated teacher checkpoint. + teacher_worker_groups: Optional[dict[str, Any]] = None + alias_to_group_alias: Optional[dict[str, str]] = 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 _non_colocated_teacher_node_count(master_config: MasterConfig) -> int: + """Validate teacher GPU geometry and return its deduplicated node count.""" + if not opd_module.is_non_colocated_teachers_enabled(master_config): + return 0 + + # Lazy to preserve teacher_worker_group's existing import cycle boundary: + # that module imports the OPD config schemas. + from nemo_rl.models.policy.teacher_worker_group import ( + create_teacher_configs_from_opd_config, + ) + + teacher_configs = create_teacher_configs_from_opd_config( + opd_module._opd_cfg(master_config) + ) + cluster_gpus_per_node = master_config.cluster["gpus_per_node"] + for teacher_config in teacher_configs: + if teacher_config.gpus_per_node > cluster_gpus_per_node: + raise ValueError( + f"OPD teacher {teacher_config.alias!r} requests " + f"gpus_per_node={teacher_config.gpus_per_node}, which exceeds " + f"cluster.gpus_per_node={cluster_gpus_per_node}." + ) + return sum(config.num_nodes for config in teacher_configs) + + def _build_clusters( master_config: MasterConfig, -) -> tuple[RayVirtualCluster, RayVirtualCluster]: - """Allocate train + inference clusters; one shared cluster when colocated. +) -> tuple[ + RayVirtualCluster, + RayVirtualCluster, + Optional[dict[str, tuple[str, int]]], +]: + """Allocate student clusters while leaving validated nodes for teachers. The colocated branch is unreachable on a real run -- validation rejects colocated.enabled=true -- and is kept for when SC can support that mode. @@ -149,17 +188,37 @@ def _build_clusters( backend = generation_config["backend"] num_nodes = cluster_config["num_nodes"] gpus_per_node = cluster_config["gpus_per_node"] + segment_size = cluster_config.get("segment_size") port_range_low = cluster_config.get("master_port_range_low") port_range_high = cluster_config.get("master_port_range_high") + teacher_nodes = _non_colocated_teacher_node_count(master_config) + policy_nodes = num_nodes - teacher_nodes + if policy_nodes <= 0: + raise ValueError( + "cluster.num_nodes must leave at least one node for the student after " + f"reserving {teacher_nodes} non-colocated teacher node(s); got " + f"cluster.num_nodes={num_nodes}." + ) + # 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 (+ critic) + generation share GPUs — one cluster. + node_constraints, remaining_ids, topology = prepare_segment_topology( + segment_size, + policy_nodes, + role="policy", + ) + teacher_topology = ( + {node_id: topology[node_id] for node_id in remaining_ids} + if segment_size is not None + else None + ) cluster = RayVirtualCluster( name="sc_policy_cluster", - bundle_ct_per_node_list=[gpus_per_node] * num_nodes, + bundle_ct_per_node_list=[gpus_per_node] * policy_nodes, use_gpus=True, num_gpus_per_node=gpus_per_node, max_colocated_worker_groups=( @@ -169,8 +228,10 @@ def _build_clusters( ), port_range_low=port_range_low, port_range_high=port_range_high, + segment_size=segment_size, + node_resource_constraints=node_constraints, ) - return cluster, cluster + return cluster, cluster, teacher_topology # Non-colocated: split node into train + inference clusters. assert backend != "megatron", ( @@ -185,7 +246,7 @@ def _build_clusters( "policy.generation.colocated.resources.gpus_per_node." ) inference_nodes = inference_resources["num_nodes"] or 1 - if num_nodes == 1: + if policy_nodes == 1: train_gpus_per_node = gpus_per_node - inference_gpus_per_node train_nodes = 1 assert train_gpus_per_node > 0, ( @@ -193,11 +254,81 @@ def _build_clusters( ) else: train_gpus_per_node = gpus_per_node - train_nodes = num_nodes - inference_nodes + train_nodes = policy_nodes - inference_nodes assert train_nodes > 0, ( - f"train_nodes must be > 0: {num_nodes} - {inference_nodes} = {train_nodes}" + f"train_nodes must be > 0: {policy_nodes} - {inference_nodes} = {train_nodes}" ) + train_constraints = None + inference_constraints = None + train_segment_size = None + inference_segment_size = None + teacher_topology = None + if segment_size is not None: + if policy_nodes == 1: + # Train and inference intentionally split one physical node by GPU. + shared_constraints, remaining_ids, topology = prepare_segment_topology( + segment_size, + 1, + role="student", + ) + train_constraints = shared_constraints + inference_constraints = shared_constraints + train_segment_size = segment_size + inference_segment_size = segment_size + teacher_topology = {node_id: topology[node_id] for node_id in remaining_ids} + else: + train_constraints, remaining_ids, topology = prepare_segment_topology( + segment_size, + train_nodes, + role="training", + ) + train_segment_size = segment_size + remaining_topology = { + node_id: topology[node_id] for node_id in remaining_ids + } + generation_config_dict = cast(dict[str, Any], generation_config) + if backend == "vllm": + vllm_cfg = generation_config_dict["vllm_cfg"] + gpus_per_instance = vllm_cfg["tensor_parallel_size"] * vllm_cfg.get( + "pipeline_parallel_size", 1 + ) + elif backend == "sglang": + gpus_per_instance = generation_config_dict["sglang_cfg"].get( + "gpus_per_server", 1 + ) + else: + raise ValueError( + "single_controller_utils.setup only supports vllm or sglang " + f"generation; got {backend!r}" + ) + nodes_per_instance = ( + gpus_per_instance + inference_gpus_per_node - 1 + ) // inference_gpus_per_node + if inference_nodes % nodes_per_instance == 0: + inference_segment_size = nodes_per_instance + ( + inference_constraints, + inference_remaining_ids, + _, + ) = prepare_segment_topology( + inference_segment_size, + inference_nodes, + topology=remaining_topology, + role="inference", + ) + teacher_topology = { + node_id: topology[node_id] for node_id in inference_remaining_ids + } + else: + print( + f" ⚠ inference_nodes={inference_nodes} is not divisible by " + f"nodes_per_instance={nodes_per_instance}; skipping inference " + "topology constraints", + flush=True, + ) + teacher_topology = remaining_topology + train_cluster = RayVirtualCluster( name="sc_train_cluster", bundle_ct_per_node_list=[train_gpus_per_node] * train_nodes, @@ -206,6 +337,8 @@ def _build_clusters( max_colocated_worker_groups=train_worker_groups, port_range_low=port_range_low, port_range_high=port_range_high, + segment_size=train_segment_size, + node_resource_constraints=train_constraints, ) inference_cluster = RayVirtualCluster( name="sc_inference_cluster", @@ -215,8 +348,10 @@ def _build_clusters( max_colocated_worker_groups=1, port_range_low=port_range_low, port_range_high=port_range_high, + segment_size=inference_segment_size, + node_resource_constraints=inference_constraints, ) - return train_cluster, inference_cluster + return train_cluster, inference_cluster, teacher_topology def _build_generation( @@ -632,6 +767,11 @@ def setup_single_controller( "single_controller_utils does not support " "data.use_multiple_dataloader=True yet." ) + if opd_module.is_opd_enabled(master_config) and processor is not None: + raise NotImplementedError( + "SingleController MOPD currently supports text-only teacher inputs. " + "Use the legacy controller for multimodal MOPD." + ) checkpointing_pretrained = master_config.checkpointing.get("pretrained_checkpoint") if checkpointing_pretrained is not None: @@ -699,8 +839,29 @@ def setup_single_controller( setup_timing_metrics = SetupTimingMetrics() # Create clusters - train_cluster, inference_cluster = _build_clusters(master_config) + train_cluster, inference_cluster, teacher_segment_topology = _build_clusters( + master_config + ) colocated = generation_config["colocated"]["enabled"] + segment_size = getattr(master_config, "cluster", {}).get("segment_size") + + # Claim constrained training nodes before unconstrained inference or Gym + # tasks can consume them. This matters when inference topology alignment + # falls back while the training cluster remains topology-constrained. + if not colocated and segment_size is not None: + train_cluster.get_placement_groups() + + # Claim teacher placement groups before deferred generation starts NeMo-Gym, + # whose resource servers may otherwise opportunistically consume those GPUs. + teacher_clusters: dict[str, RayVirtualCluster] = {} + if opd_module.is_non_colocated_teachers_enabled(master_config): + t0 = time.perf_counter() + teacher_clusters = opd_module.reserve_teacher_clusters( + master_config, + segment_size=segment_size, + teacher_segment_topology=teacher_segment_topology, + ) + setup_timing_metrics.teacher_reservation_time_s = time.perf_counter() - t0 # Create build tasks for generation / trainer / (nemo-gym) workers build_tasks: dict[str, Callable[[], Any]] = {} @@ -863,6 +1024,28 @@ def _build_generation_then_trainer( setup_timing_metrics.generation_init_reserve_time_s = gen_reserve_time setup_timing_metrics.generation_init_load_time_s = gen_load_time + # Loading a teacher with the same checkpoint as the student must happen only + # after student initialization finishes: both use the same HF-to-Megatron + # cache path, and concurrent conversion can expose a partial checkpoint. + teacher_worker_groups: dict[str, Any] = {} + alias_to_group_alias: dict[str, str] = {} + if teacher_clusters: + t0 = time.perf_counter() + teacher_worker_groups, alias_to_group_alias = ( + opd_module.create_teacher_worker_groups( + master_config, + cast(dict[str, Any], policy_config), + tokenizer, + teacher_clusters=teacher_clusters, + ) + ) + for teacher in teacher_worker_groups.values(): + teacher.setup_data_plane(dp_config) + setup_timing_metrics.teacher_model_init_time_s = time.perf_counter() - t0 + setup_timing_metrics.teacher_init_time_s = ( + setup_timing_metrics.teacher_reservation_time_s or 0.0 + ) + setup_timing_metrics.teacher_model_init_time_s + worker_setup_time = time.perf_counter() - setup_start_time setup_timing_metrics.worker_setup_time_s = worker_setup_time @@ -875,6 +1058,22 @@ def _build_generation_then_trainer( # ========================== # Connect-only DP client; TQPolicy already bootstrapped the controller. dp_client = build_data_plane_client(dp_config, bootstrap=False) + # SingleController reuses one partition for the run. Warm every known + # tensor field before rollout, policy, and teacher writers become + # concurrent; TransferQueue otherwise registers field names lazily. + dp_client.register_partition( + partition_id=partition_id, + fields=fields_with_optional_routed_experts( + SC_ROLLOUT_SCHEMA_FIELDS, + enabled=router_replay_enabled(policy_config), + ), + num_samples=( + master_config.async_rl.max_buffered_rollouts + * algo_cfg.num_generations_per_prompt + ), + consumer_tasks=["prev_lp", "ref_lp", "train"], + grpo_group_size=algo_cfg.num_generations_per_prompt, + ) t0 = time.perf_counter() weight_synchronizer = create_weight_synchronizer( @@ -952,6 +1151,8 @@ def _build_generation_then_trainer( last_checkpoint_path=last_checkpoint_path, fleet_monitor=fleet_monitor, generation_router=generation_router, + teacher_worker_groups=teacher_worker_groups, + alias_to_group_alias=alias_to_group_alias, # PPO extras value_handle=value, value_loss_fn=value_loss_fn, diff --git a/nemo_rl/data_plane/column_io.py b/nemo_rl/data_plane/column_io.py index 5242f282263..9bbdb5616c0 100644 --- a/nemo_rl/data_plane/column_io.py +++ b/nemo_rl/data_plane/column_io.py @@ -46,6 +46,7 @@ "generation_logprobs", "prev_logprobs", "reference_policy_logprobs", + "teacher_reference_logprobs", "advantages", "returns", "values", diff --git a/nemo_rl/data_plane/schema.py b/nemo_rl/data_plane/schema.py index 56464e4c4d1..9c1f9d1e2ef 100644 --- a/nemo_rl/data_plane/schema.py +++ b/nemo_rl/data_plane/schema.py @@ -44,6 +44,20 @@ "sample_mask", ) +# Full known tensor schema for SingleController's long-lived rollout partition. +# The initial rollout put writes the first seven payload fields; later stages add +# student/reference logprobs, advantages, PPO critic columns, and the MOPD teacher +# column. Registering their names once before concurrent producers start avoids +# TransferQueue's lazy field-name registration race. +SC_ROLLOUT_SCHEMA_FIELDS = ( + *DP_TRAIN_FIELDS, + "prompt_ids_for_adv", + "total_reward", + "values", + "returns", + "teacher_reference_logprobs", +) + # Subset fetched by logprob / ref-logprob workers. LP_SEED_FIELDS = ( "input_ids", @@ -52,6 +66,9 @@ "sample_mask", ) +# Text-only inputs fetched by frozen MOPD teachers for logprob inference. +TEACHER_LP_FIELDS = (INPUT_IDS, INPUT_LENGTHS) + # 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 = ( diff --git a/nemo_rl/data_plane/worker_mixin.py b/nemo_rl/data_plane/worker_mixin.py index 88329e73697..e545ae48370 100644 --- a/nemo_rl/data_plane/worker_mixin.py +++ b/nemo_rl/data_plane/worker_mixin.py @@ -15,7 +15,8 @@ Mix into a worker class to add per-rank TQ-mediated entrypoints (:meth:`train_presharded`, :meth:`get_logprobs_presharded`, -:meth:`get_reference_policy_logprobs_presharded`) without touching +:meth:`get_reference_policy_logprobs_presharded`, and the frozen-teacher +variant) without touching ``BasePolicyWorker``. Subclasses that don't need TQ keep their bare inheritance and stay zero-cost. @@ -525,6 +526,41 @@ def get_reference_policy_logprobs_presharded( ) del result + @wrap_with_nvtx_name("policy_worker/get_teacher_logprobs_presharded") + def get_teacher_logprobs_presharded( + self, + meta: "KVBatchMeta", + micro_batch_size: Optional[int] = None, + ) -> None: + """Per-rank frozen-teacher logprob entrypoint for SingleController MOPD.""" + data = self._fetch(meta) + cfg = getattr(self, "cfg", {}) + batching_enabled = bool( + cfg.get("sequence_packing", {}).get("enabled", False) + or cfg.get("dynamic_batching", {}).get("enabled", False) + ) + extra = meta.extra_info or {} + if batching_enabled and not ( + MICRO_BATCH_INDICES in extra and MICRO_BATCH_LENGTHS in extra + ): + raise RuntimeError( + "SingleController teacher batching requires driver-provided global " + "micro_batch_indices and micro_batch_lengths; local worker planning " + "can desynchronize data-parallel collectives." + ) + data = self._attach_or_repack_pack_metadata(data, meta) + result: BatchedDataDict[Any] = self.get_logprobs( # type: ignore[attr-defined] + data=data, + micro_batch_size=micro_batch_size, + ) + self._write_back_result_field( + meta, + result, + result_key="logprobs", + tq_field="teacher_reference_logprobs", + ) + del result + @wrap_with_nvtx_name("value_worker/get_values_presharded") def get_values_presharded( self, diff --git a/nemo_rl/experience/rollout_manager.py b/nemo_rl/experience/rollout_manager.py index 784ae5a7042..6717542d642 100644 --- a/nemo_rl/experience/rollout_manager.py +++ b/nemo_rl/experience/rollout_manager.py @@ -24,7 +24,10 @@ from transformers import PreTrainedTokenizerBase from wandb import Table -from nemo_rl.algorithms.async_utils.replay_buffer import TQReplayBuffer +from nemo_rl.algorithms.async_utils.replay_buffer import ( + PostWriteEnrichmentError, + TQReplayBuffer, +) from nemo_rl.data.interfaces import DatumSpec, LLMMessageLogType from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.environments.interfaces import EnvironmentInterface @@ -61,6 +64,17 @@ TokenizerType = PreTrainedTokenizerBase +def _contains_post_write_enrichment_error(error: BaseException) -> bool: + """Whether an error, including a rollback ExceptionGroup, is post-write.""" + if isinstance(error, PostWriteEnrichmentError): + return True + if isinstance(error, BaseExceptionGroup): + return any( + _contains_post_write_enrichment_error(child) for child in error.exceptions + ) + return False + + class RolloutOutcome(str, enum.Enum): """How :meth:`RolloutManager.generate_and_push` finished for one prompt.""" @@ -1300,6 +1314,11 @@ async def generate_and_push( f" warn: remove_group({group_id}) cleanup failed: {cleanup_exc!r}", flush=True, ) + # The rollout itself succeeded. Re-running generation cannot repair + # a required downstream stage (for example MOPD teacher inference), + # and would spend the rollout retry budget on the wrong subsystem. + if _contains_post_write_enrichment_error(error): + raise reason = type(error).__name__ if classify_rollout_failure(error) is FailureClass.INFRA: diff --git a/nemo_rl/models/policy/teacher_worker_group.py b/nemo_rl/models/policy/teacher_worker_group.py index 1657fffbd9e..69de3d5e85c 100644 --- a/nemo_rl/models/policy/teacher_worker_group.py +++ b/nemo_rl/models/policy/teacher_worker_group.py @@ -22,15 +22,21 @@ import warnings from copy import deepcopy -from dataclasses import dataclass +from dataclasses import dataclass, replace from typing import Any, Optional import numpy as np +import ray from transformers import PreTrainedTokenizerBase from nemo_rl.algorithms.opd import TeacherResourceConfig +from nemo_rl.data_plane import DataPlaneConfig, KVBatchMeta +from nemo_rl.data_plane.column_io import round_up +from nemo_rl.data_plane.preshard import shard_meta_for_dp +from nemo_rl.data_plane.schema import GLOBAL_FORWARD_PAD_SEQLEN, TEACHER_LP_FIELDS from nemo_rl.distributed.batched_data_dict import ( BatchedDataDict, + DynamicBatchingArgs, SequencePackingArgs, ) from nemo_rl.distributed.named_sharding import NamedSharding @@ -114,7 +120,7 @@ class TeacherWorkerGroup: - Never initializes an optimizer - Never initializes a reference model - Loads the checkpoint once at startup - - Only exposes get_logprobs() + - Exposes logprobs through either the legacy tensor API or TQ metadata API """ def __init__( @@ -163,6 +169,11 @@ def __init__( cfg["megatron_cfg"]["peft"]["enabled"] = False if "draft" in cfg: cfg["draft"]["enabled"] = False + # Router replay keeps the student's rollout and training logprobs + # consistent. A frozen teacher has no training pass, and its text-only + # TQ fetch does not carry routed_experts, so replay must stay off. + if "router_replay" in cfg: + cfg["router_replay"]["enabled"] = False # The teacher uses the plain Megatron worker, so a student-side quant_cfg # would be silently ignored. Drop it explicitly and warn instead. if cfg.get("quant_cfg") is not None: @@ -240,6 +251,90 @@ def __init__( microbatch_order = cfg["sequence_packing"].get("microbatch_order") if microbatch_order is not None: self.sequence_packing_args["microbatch_order"] = microbatch_order + if self.use_dynamic_batches: + self.dynamic_batching_args: DynamicBatchingArgs = { + "input_key": "input_ids", + "input_lengths_key": "input_lengths", + "sequence_length_round": cfg["dynamic_batching"][ + "sequence_length_round" + ], + "max_tokens_per_microbatch": cfg["dynamic_batching"][ + "logprob_mb_tokens" + ], + } + + def setup_data_plane(self, dp_cfg: DataPlaneConfig) -> None: + """Attach every teacher worker to the already-bootstrapped TQ controller.""" + ray.get( + self.worker_group.run_all_workers_single_data( + "setup_data_plane", cfg=dp_cfg + ) + ) + + def get_logprobs_from_meta(self, meta: KVBatchMeta) -> None: + """Dispatch a TQ metadata batch; workers fetch and write teacher logprobs.""" + if not meta.sequence_lengths: + raise ValueError( + f"Teacher {self.alias!r} requires sequence_lengths in TQ metadata" + ) + sequence_pad_multiple = ( + 1 if self.use_sequence_packing else (self.sequence_length_pad_multiple) + ) + if self.use_dynamic_batches: + # The driver rounds dynamic microbatch lengths to this value. Match + # that rounding in the stamped fetch width so worker-side narrowing + # cannot request more tokens than TQ materialized. + sequence_pad_multiple = max( + sequence_pad_multiple, + int(self.cfg["dynamic_batching"]["sequence_length_round"]), + ) + teacher_meta = replace( + meta, + task_name=f"teacher_lp:{self.alias}", + fields=list(TEACHER_LP_FIELDS), + extra_info={ + **dict(meta.extra_info or {}), + GLOBAL_FORWARD_PAD_SEQLEN: round_up( + max(meta.sequence_lengths), sequence_pad_multiple + ), + }, + ) + sequence_packing_args = None + dynamic_batching_args = None + if self.use_sequence_packing: + sequence_packing_args = dict(self.sequence_packing_args) + sequence_packing_args["max_tokens_per_microbatch"] = self.cfg[ + "sequence_packing" + ]["logprob_mb_tokens"] + elif self.use_dynamic_batches: + dynamic_batching_args = dict(self.dynamic_batching_args) + dynamic_batching_args["max_tokens_per_microbatch"] = self.cfg[ + "dynamic_batching" + ]["logprob_mb_tokens"] + dp_metas, _ = shard_meta_for_dp( + teacher_meta, + dp_world=self.sharding_annotations.get_axis_size("data_parallel"), + batch_size=None, + sequence_packing_args=sequence_packing_args, + dynamic_batching_args=dynamic_batching_args, + ) + futures = self.worker_group.run_all_workers_sharded_data( + "get_teacher_logprobs_presharded", + meta=dp_metas, + in_sharded_axes=["data_parallel"], + replicate_on_axes=[ + "context_parallel", + "tensor_parallel", + "pipeline_parallel", + ], + output_is_replicated=[ + "context_parallel", + "tensor_parallel", + "pipeline_parallel", + ], + common_kwargs={"micro_batch_size": self._micro_batch_size}, + ) + self.worker_group.get_all_worker_results(futures) def get_logprobs( self, diff --git a/tests/test_suites/llm/mopd-qwen3-1.7b-3n8g-megatron-pack-single-controller.sh b/tests/test_suites/llm/mopd-qwen3-1.7b-3n8g-megatron-pack-single-controller.sh new file mode 100755 index 00000000000..e33726d8652 --- /dev/null +++ b/tests/test_suites/llm/mopd-qwen3-1.7b-3n8g-megatron-pack-single-controller.sh @@ -0,0 +1,61 @@ +#!/bin/bash +# 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. +# +# Text-only SingleController MOPD sanity test. Student and teacher are the same +# Qwen3-1.7B checkpoint, so the distillation loss should remain close to zero. +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source $SCRIPT_DIR/common.env + +# ===== BEGIN CONFIG ===== +NUM_NODES=3 +GPUS_PER_NODE=8 +STEPS_PER_RUN=5 +MAX_STEPS=5 +NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) # Round up +NUM_MINUTES=18 +USES_SANDBOX=1 +USE_GYM_CONTAINER=true +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +cd $PROJECT_ROOT +uv run examples/run_grpo_single_controller.py \ + --config $CONFIG_PATH \ + grpo.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=False \ + checkpointing.checkpoint_dir=$CKPT_DIR \ + "$@" \ + 2>&1 | tee $RUN_LOG + +uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS + +if [[ $(jq 'to_entries | .[] | select(.key == "train/token_mult_prob_error") | .value | keys | map(tonumber) | max' $JSON_METRICS) -ge $MAX_STEPS ]]; then + uv run tests/check_metrics.py $JSON_METRICS \ + 'abs(median(data["train/loss"])) < 0.05' \ + 'median(data["train/token_mult_prob_error"]) < 1.1' \ + 'max(data["train/on_policy_distillation/teacher_batches"]) > 0' \ + 'max(data["train/on_policy_distillation/teacher_samples"]) > 0' \ + 'max(data["train/on_policy_distillation/teacher_model_unique"]) == 1' \ + 'max(data["train/on_policy_distillation/adv_std"]) < 0.01' + + rm -rf "$CKPT_DIR" +fi diff --git a/tests/test_suites/nightly.txt b/tests/test_suites/nightly.txt index 894600d37d6..e8b652a84c2 100644 --- a/tests/test_suites/nightly.txt +++ b/tests/test_suites/nightly.txt @@ -270,6 +270,10 @@ tests/test_suites/llm/distillation-qwen3-1.7b-1n8g-megatron-qa-nvfp4.sh # NRL_TRAIN_PATH / NRL_VAL_PATH (nemo_gym jsonl) in the environment. tests/test_suites/llm/mopd-qwen3-1.7b-3n8g-megatron-pack.sh +# SingleController/TQ sibling of the self-distillation MOPD sanity check above. +# In addition to near-zero loss, it asserts positive teacher-stage activity. +tests/test_suites/llm/mopd-qwen3-1.7b-3n8g-megatron-pack-single-controller.sh + # Cross-tokenizer off-policy distillation (xtoken): Qwen3-4B -> Llama-3.2-1B P-KL, student TP4xCP2 <- teacher TP2xCP2, guards sharded-loss parallelism-invariance. tests/test_suites/llm/distillation-xtoken-off-policy-qwen3-4b-to-llama3.2-1b-1n8g-dtensor-tp4cp2.sh diff --git a/tests/unit/algorithms/test_opd.py b/tests/unit/algorithms/test_opd.py index 73077687e00..258bd319888 100644 --- a/tests/unit/algorithms/test_opd.py +++ b/tests/unit/algorithms/test_opd.py @@ -39,6 +39,8 @@ class _MockTeacherWorkerGroup: def __init__(self, fill_value=1.0, dp_size=4): self._fill_value = fill_value self.sharding_annotations = _MockShardingAnnotations(dp_size) + self.use_sequence_packing = False + self.sequence_length_pad_multiple = 1 def get_logprobs(self, data): input_ids = data["input_ids"] @@ -319,6 +321,488 @@ def test_compute_teacher_logprobs_default_alias_fallback_routes(): assert torch.allclose(result, torch.tensor(7.0)) +# --------------------------------------------------------------------------- +# SingleController TQ teacher enrichment +# --------------------------------------------------------------------------- + + +def _teacher_record(agent_name: str): + from nemo_rl.experience.interfaces import PromptGroupRecord + + return PromptGroupRecord( + prompt_idx=0, + prompt=[], + extra_env_info={"agent_ref": {"name": agent_name}}, + metadata={}, + completions=[], + rollout_metrics={}, + ) + + +def _teacher_meta(prefix: str, batch_size: int, seq_len: int): + from nemo_rl.data_plane import KVBatchMeta + + return KVBatchMeta( + partition_id="rollout_data", + task_name="train", + sample_ids=[f"{prefix}_{index}" for index in range(batch_size)], + fields=["input_ids", "input_lengths"], + sequence_lengths=[seq_len] * batch_size, + ) + + +def test_tq_teacher_enrichment_pads_dp_and_writes_teacher_column(monkeypatch): + """The SC coordinator preserves DP padding while TQ remains source and sink.""" + import asyncio + + from nemo_rl.algorithms import opd + + class MetaTeacher: + def __init__(self): + self.sharding_annotations = _MockShardingAnnotations(4) + self.received_meta = None + + def get_logprobs_from_meta(self, meta): + self.received_meta = meta + + class FakeDataPlane: + def __init__(self): + self.clear_calls = [] + + def clear_samples(self, sample_ids, partition_id): + self.clear_calls.append((list(sample_ids), partition_id)) + + teacher = MetaTeacher() + dp_client = FakeDataPlane() + writes = [] + + def fake_read_columns(dp_client, meta, select_fields, pad_value_dict): + del dp_client, select_fields, pad_value_dict + batch_size = len(meta.sample_ids) + seq_len = max(meta.sequence_lengths) + return BatchedDataDict( + { + "input_ids": torch.arange(batch_size * seq_len).reshape( + batch_size, seq_len + ), + "input_lengths": torch.tensor(meta.sequence_lengths), + } + ) + + def fake_write_columns(dp_client, meta, fields): + del dp_client + writes.append((meta, fields)) + + monkeypatch.setattr(opd, "read_columns", fake_read_columns) + monkeypatch.setattr(opd, "write_columns", fake_write_columns) + coordinator = opd.TQTeacherLogprobCoordinator( + dp_client=dp_client, + teacher_worker_groups={"primary": teacher}, + alias_to_group_alias={"math": "primary"}, + on_policy_distillation_cfg={ + "teacher_model_by_agent_name": {"math": "/ckpt/shared"} + }, + ) + meta = _teacher_meta("group", batch_size=3, seq_len=5) + + enriched = asyncio.run(coordinator.enrich(meta, _teacher_record("math"))) + + assert teacher.received_meta is not None + assert teacher.received_meta.size == 4 + assert teacher.received_meta.sample_ids[:3] == meta.sample_ids + assert "__teacher_pad_" in teacher.received_meta.sample_ids[-1] + assert len(writes) == 1 + pad_meta, padding_fields = writes[0] + assert pad_meta.size == 1 + assert padding_fields["input_ids"].shape == (1, 5) + assert dp_client.clear_calls == [(pad_meta.sample_ids, "rollout_data")] + assert "teacher_reference_logprobs" in enriched.fields + metrics = coordinator.drain_metrics() + assert metrics["on_policy_distillation/teacher_batches"] == 1.0 + assert metrics["on_policy_distillation/teacher_samples"] == 3.0 + assert metrics["on_policy_distillation/teacher_model_unique"] == 1.0 + + +def test_tq_teacher_enrichment_skips_padding_for_dp_divisible_batch(monkeypatch): + """DP-divisible teacher batches do not create or clean temporary TQ rows.""" + import asyncio + + from nemo_rl.algorithms import opd + + teacher = _MockTeacherWorkerGroup(dp_size=2) + teacher.received_meta = None + teacher.get_logprobs_from_meta = lambda meta: setattr( + teacher, "received_meta", meta + ) + + class FakeDataPlane: + def clear_samples(self, **kwargs): + raise AssertionError(f"unexpected temporary-row cleanup: {kwargs}") + + monkeypatch.setattr( + opd, + "read_columns", + lambda *args, **kwargs: (_ for _ in ()).throw( + AssertionError("DP-divisible batches must not read a padding source") + ), + ) + monkeypatch.setattr( + opd, + "write_columns", + lambda *args, **kwargs: (_ for _ in ()).throw( + AssertionError("DP-divisible batches must not write padding rows") + ), + ) + coordinator = opd.TQTeacherLogprobCoordinator( + dp_client=FakeDataPlane(), + teacher_worker_groups={"teacher": teacher}, + alias_to_group_alias={"math": "teacher"}, + on_policy_distillation_cfg={ + "teacher_model_by_agent_name": {"math": "/ckpt/teacher"} + }, + ) + meta = _teacher_meta("group", batch_size=2, seq_len=5) + + enriched = asyncio.run(coordinator.enrich(meta, _teacher_record("math"))) + + assert teacher.received_meta is meta + assert "teacher_reference_logprobs" in enriched.fields + + +def test_tq_teacher_routing_rejects_missing_agent_ref_without_retry_hint(): + """Missing Gym routing metadata gets a diagnostic that identifies the cause.""" + from nemo_rl.algorithms import opd + from nemo_rl.experience.interfaces import PromptGroupRecord + + coordinator = opd.TQTeacherLogprobCoordinator( + dp_client=object(), + teacher_worker_groups={"teacher": _MockTeacherWorkerGroup(dp_size=1)}, + alias_to_group_alias={}, + on_policy_distillation_cfg={ + "teacher_model_by_agent_name": {"teacher": "/ckpt/teacher"} + }, + ) + record = PromptGroupRecord( + prompt_idx=0, + prompt=[], + extra_env_info={}, + metadata={}, + completions=[], + rollout_metrics={}, + ) + + with pytest.raises( + ValueError, + match="requires the NeMo-Gym rollout path.*cannot repair", + ): + coordinator._resolve_teacher(record) + + +def test_tq_teacher_routing_uses_default_teacher_for_unmapped_agent(): + """An unmapped Gym agent follows the configured default teacher alias.""" + from nemo_rl.algorithms import opd + + coordinator = opd.TQTeacherLogprobCoordinator( + dp_client=object(), + teacher_worker_groups={"default": _MockTeacherWorkerGroup(dp_size=1)}, + alias_to_group_alias={"default": "default"}, + on_policy_distillation_cfg={ + "teacher_model_by_agent_name": {"default": "/ckpt/default"}, + "default_teacher_alias": "default", + }, + ) + + assert coordinator._resolve_teacher(_teacher_record("unmapped")) == ( + "default", + "default", + ) + + +def test_tq_teacher_padding_rows_are_cleaned_when_write_partially_fails(monkeypatch): + """Temporary IDs enter cleanup scope before their first TQ write.""" + import asyncio + + from nemo_rl.algorithms import opd + + cleared = [] + + class FakeDataPlane: + def clear_samples(self, sample_ids, partition_id): + cleared.append((list(sample_ids), partition_id)) + + monkeypatch.setattr( + opd, + "read_columns", + lambda *args, **kwargs: BatchedDataDict( + { + "input_ids": torch.ones(1, 3, dtype=torch.long), + "input_lengths": torch.tensor([3]), + } + ), + ) + + def partially_failing_write(_dp_client, meta, fields): + del fields + assert meta.sample_ids + raise RuntimeError("partial pad write") + + monkeypatch.setattr(opd, "write_columns", partially_failing_write) + coordinator = opd.TQTeacherLogprobCoordinator( + dp_client=FakeDataPlane(), + teacher_worker_groups={"teacher": _MockTeacherWorkerGroup(dp_size=4)}, + alias_to_group_alias={"math": "teacher"}, + on_policy_distillation_cfg={ + "teacher_model_by_agent_name": {"math": "/ckpt/teacher"} + }, + ) + + with pytest.raises(RuntimeError, match="partial pad write"): + asyncio.run( + coordinator.enrich( + _teacher_meta("group", batch_size=3, seq_len=3), + _teacher_record("math"), + ) + ) + + assert len(cleared) == 1 + assert len(cleared[0][0]) == 1 + assert "__teacher_pad_" in cleared[0][0][0] + + +def test_tq_teacher_enrichment_drains_background_thread_before_cancellation(): + """Cancellation waits for teacher TQ activity to finish before propagating.""" + import asyncio + import threading + + from nemo_rl.algorithms import opd + + started = threading.Event() + release = threading.Event() + finished = threading.Event() + + class BlockingTeacher(_MockTeacherWorkerGroup): + def get_logprobs_from_meta(self, meta): + del meta + started.set() + release.wait(timeout=2) + finished.set() + + coordinator = opd.TQTeacherLogprobCoordinator( + dp_client=object(), + teacher_worker_groups={"teacher": BlockingTeacher(dp_size=1)}, + alias_to_group_alias={"math": "teacher"}, + on_policy_distillation_cfg={ + "teacher_model_by_agent_name": {"math": "/ckpt/teacher"} + }, + ) + + async def cancel_during_inference(): + task = asyncio.create_task( + coordinator.enrich( + _teacher_meta("group", batch_size=1, seq_len=3), + _teacher_record("math"), + ) + ) + assert await asyncio.to_thread(started.wait, 1) + task.cancel() + await asyncio.sleep(0) + assert not task.done() + release.set() + with pytest.raises(asyncio.CancelledError): + await task + + asyncio.run(cancel_during_inference()) + assert finished.is_set() + + +def test_tq_teacher_waiters_do_not_occupy_executor_threads(monkeypatch): + """Only the active inference for a physical teacher enters to_thread.""" + import asyncio + import threading + + from nemo_rl.algorithms import opd + + started = threading.Event() + release = threading.Event() + submissions = 0 + real_to_thread = asyncio.to_thread + + class BlockingTeacher(_MockTeacherWorkerGroup): + def get_logprobs_from_meta(self, meta): + del meta + started.set() + release.wait(timeout=2) + + async def counted_to_thread(func, /, *args, **kwargs): + nonlocal submissions + submissions += 1 + return await real_to_thread(func, *args, **kwargs) + + monkeypatch.setattr(opd.asyncio, "to_thread", counted_to_thread) + coordinator = opd.TQTeacherLogprobCoordinator( + dp_client=object(), + teacher_worker_groups={"teacher": BlockingTeacher(dp_size=1)}, + alias_to_group_alias={"math": "teacher"}, + on_policy_distillation_cfg={ + "teacher_model_by_agent_name": {"math": "/ckpt/teacher"} + }, + ) + + async def run_waiters(): + first = asyncio.create_task( + coordinator.enrich( + _teacher_meta("first", batch_size=1, seq_len=3), + _teacher_record("math"), + ) + ) + while not started.is_set(): + await asyncio.sleep(0.001) + second = asyncio.create_task( + coordinator.enrich( + _teacher_meta("second", batch_size=1, seq_len=3), + _teacher_record("math"), + ) + ) + await asyncio.sleep(0.01) + assert submissions == 1 + release.set() + await asyncio.gather(first, second) + + asyncio.run(run_waiters()) + assert submissions == 2 + + +def test_tq_teacher_enrichment_serializes_deduplicated_teacher(monkeypatch): + """Two aliases sharing one physical teacher never overlap collectives.""" + import asyncio + import threading + import time + + from nemo_rl.algorithms import opd + + active = 0 + max_active = 0 + active_lock = threading.Lock() + + class SlowTeacher(_MockTeacherWorkerGroup): + def get_logprobs_from_meta(self, meta): + del meta + nonlocal active, max_active + with active_lock: + active += 1 + max_active = max(max_active, active) + time.sleep(0.05) + try: + return None + finally: + with active_lock: + active -= 1 + + def fake_read_columns(dp_client, meta, select_fields, pad_value_dict): + del dp_client, select_fields, pad_value_dict + return BatchedDataDict( + { + "input_ids": torch.ones(len(meta.sample_ids), 4, dtype=torch.long), + "input_lengths": torch.full( + (len(meta.sample_ids),), 4, dtype=torch.long + ), + } + ) + + monkeypatch.setattr(opd, "read_columns", fake_read_columns) + monkeypatch.setattr(opd, "write_columns", lambda *args, **kwargs: None) + teacher = SlowTeacher(dp_size=1) + coordinator = opd.TQTeacherLogprobCoordinator( + dp_client=object(), + teacher_worker_groups={"primary": teacher}, + alias_to_group_alias={"math": "primary", "code": "primary"}, + on_policy_distillation_cfg={ + "teacher_model_by_agent_name": { + "math": "/ckpt/shared", + "code": "/ckpt/shared", + } + }, + ) + + async def run_both(): + await asyncio.gather( + coordinator.enrich(_teacher_meta("math", 1, 4), _teacher_record("math")), + coordinator.enrich(_teacher_meta("code", 1, 4), _teacher_record("code")), + ) + + asyncio.run(run_both()) + + assert max_active == 1 + + +def test_tq_teacher_enrichment_runs_distinct_teachers_concurrently(): + """Distinct physical teachers hold distinct locks, so they overlap.""" + import asyncio + import threading + + from nemo_rl.algorithms import opd + + barrier = threading.Barrier(2) + + class BarrierTeacher(_MockTeacherWorkerGroup): + def get_logprobs_from_meta(self, meta): + del meta + # Both teachers must enter inference concurrently. One shared lock + # would serialize them and trip BrokenBarrierError. + barrier.wait(timeout=5) + + coordinator = opd.TQTeacherLogprobCoordinator( + dp_client=object(), + teacher_worker_groups={ + "primary": BarrierTeacher(dp_size=1), + "secondary": BarrierTeacher(dp_size=1), + }, + alias_to_group_alias={"math": "primary", "code": "secondary"}, + on_policy_distillation_cfg={ + "teacher_model_by_agent_name": { + "math": "/ckpt/math", + "code": "/ckpt/code", + } + }, + ) + + async def run_both(): + await asyncio.gather( + coordinator.enrich(_teacher_meta("math", 1, 4), _teacher_record("math")), + coordinator.enrich(_teacher_meta("code", 1, 4), _teacher_record("code")), + ) + + asyncio.run(run_both()) + + metrics = coordinator.drain_metrics() + assert metrics["on_policy_distillation/teacher_model_unique"] == 2.0 + + +def test_tq_teacher_metrics_omit_routing_cardinality_on_idle_drain(): + """An idle interval reports zero activity without claiming zero teachers.""" + from nemo_rl.algorithms import opd + + coordinator = opd.TQTeacherLogprobCoordinator( + dp_client=object(), + teacher_worker_groups={"teacher": _MockTeacherWorkerGroup(dp_size=1)}, + alias_to_group_alias={"math": "teacher"}, + on_policy_distillation_cfg={ + "teacher_model_by_agent_name": {"math": "/ckpt/teacher"} + }, + ) + + metrics = coordinator.drain_metrics() + + assert metrics["on_policy_distillation/teacher_batches"] == 0.0 + assert metrics["on_policy_distillation/teacher_samples"] == 0.0 + assert metrics["on_policy_distillation/teacher_logprob_time_s"] == 0.0 + assert metrics["on_policy_distillation/teacher_inference_time_s"] == 0.0 + assert metrics["on_policy_distillation/teacher_lock_wait_time_s"] == 0.0 + assert "on_policy_distillation/teacher_alias_unique" not in metrics + assert "on_policy_distillation/teacher_model_unique" not in metrics + assert "on_policy_distillation/teacher_alias_to_model_compression" not in metrics + + # --------------------------------------------------------------------------- # Unsort / reorder_data regression test # --------------------------------------------------------------------------- diff --git a/tests/unit/experience/test_rollout_manager.py b/tests/unit/experience/test_rollout_manager.py index 0a73393e7db..88d4ab6529f 100644 --- a/tests/unit/experience/test_rollout_manager.py +++ b/tests/unit/experience/test_rollout_manager.py @@ -33,6 +33,7 @@ import pytest import torch +from nemo_rl.algorithms.async_utils.replay_buffer import PostWriteEnrichmentError from nemo_rl.data.collate_fn import rl_collate_fn from nemo_rl.data.datasets.response_datasets import NemoGymDataset from nemo_rl.data.interfaces import DatumSpec @@ -153,6 +154,96 @@ def _make_manager( class TestGenerateAndPushFlow: + def test_post_write_failure_does_not_regenerate_the_rollout(self): + class _EnrichmentFailBuffer(_FakeBuffer): + async def commit( + self, + group_id: str, + record, + start_weight_version: int, + end_weight_version: int, + ): + await super().commit( + group_id, + record, + start_weight_version, + end_weight_version, + ) + raise PostWriteEnrichmentError("teacher stage failed") + + rollout_calls = 0 + + async def _count_rollout(_sample): + nonlocal rollout_calls + rollout_calls += 1 + + buf = _EnrichmentFailBuffer() + mgr = _make_manager( + buf, + _FakeImpl(on_run=_count_rollout), + retry_policy=RolloutRetryPolicy( + max_infra_attempts=3, + max_data_attempts=3, + max_gym_row_attempts=1, + ), + ) + + with pytest.raises(PostWriteEnrichmentError, match="teacher stage failed"): + _run(mgr.generate_and_push({"prompt": "p"})) + + assert rollout_calls == 1 + assert len(buf.reserve_calls) == 1 + assert len(buf.remove_calls) == 1 + + def test_grouped_post_write_failure_does_not_regenerate_the_rollout(self): + """Rollback failures do not hide the post-write failure classification.""" + + class _GroupedEnrichmentFailBuffer(_FakeBuffer): + async def commit( + self, + group_id: str, + record, + start_weight_version: int, + end_weight_version: int, + ): + await super().commit( + group_id, + record, + start_weight_version, + end_weight_version, + ) + raise ExceptionGroup( + "commit and rollback both failed", + [ + PostWriteEnrichmentError("teacher stage failed"), + RuntimeError("rollback failed"), + ], + ) + + rollout_calls = 0 + + async def _count_rollout(_sample): + nonlocal rollout_calls + rollout_calls += 1 + + buf = _GroupedEnrichmentFailBuffer() + mgr = _make_manager( + buf, + _FakeImpl(on_run=_count_rollout), + retry_policy=RolloutRetryPolicy( + max_infra_attempts=3, + max_data_attempts=3, + max_gym_row_attempts=1, + ), + ) + + with pytest.raises(ExceptionGroup, match="commit and rollback"): + _run(mgr.generate_and_push({"prompt": "p"})) + + assert rollout_calls == 1 + assert len(buf.reserve_calls) == 1 + assert len(buf.remove_calls) == 1 + def test_explicit_registry_tracks_only_inflight_generation(self): registry: dict[str, tuple[asyncio.Task[None], int]] = {} buf = _FakeBuffer() diff --git a/tests/unit/models/policy/test_teacher_worker_group.py b/tests/unit/models/policy/test_teacher_worker_group.py index 94c25c5f8e5..2fb51686503 100644 --- a/tests/unit/models/policy/test_teacher_worker_group.py +++ b/tests/unit/models/policy/test_teacher_worker_group.py @@ -12,6 +12,20 @@ # See the License for the specific language governing permissions and # limitations under the License. +from unittest.mock import MagicMock + +import pytest +import torch + +from nemo_rl.data_plane import KVBatchMeta +from nemo_rl.data_plane.schema import ( + GLOBAL_FORWARD_PAD_SEQLEN, + MICRO_BATCH_INDICES, + MICRO_BATCH_LENGTHS, + TEACHER_LP_FIELDS, +) +from nemo_rl.distributed.batched_data_dict import BatchedDataDict + def test_teacher_resource_config_defaults(): from nemo_rl.algorithms.opd import TeacherResourceConfig @@ -77,3 +91,343 @@ def test_create_teacher_configs_deduplicates(): } ) assert len(configs) == 2 + + +def test_teacher_worker_group_disables_student_router_replay(monkeypatch): + """Frozen teachers do not require rollout-to-training route consistency.""" + import nemo_rl.distributed.worker_groups as worker_groups + from nemo_rl.models.policy.teacher_worker_group import ( + TeacherConfig, + TeacherWorkerGroup, + ) + + captured = {} + + class FakeWorkerBuilder: + def __init__(self, worker_path, cfg, **kwargs): + del worker_path, kwargs + captured["cfg"] = cfg + + class FakeWorkerGroup: + def __init__(self, cluster, worker_builder, **kwargs): + del cluster, worker_builder, kwargs + + monkeypatch.setattr(worker_groups, "RayWorkerBuilder", FakeWorkerBuilder) + monkeypatch.setattr(worker_groups, "RayWorkerGroup", FakeWorkerGroup) + cluster = MagicMock() + cluster.world_size.return_value = 1 + policy_config = { + "model_name": "/ckpt/student", + "megatron_cfg": {"enabled": True}, + "dtensor_cfg": {"enabled": False}, + "sequence_packing": {"enabled": False}, + "dynamic_batching": {"enabled": False}, + "router_replay": {"enabled": True}, + } + teacher_config = TeacherConfig( + alias="teacher", + model_name="/ckpt/teacher", + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + context_parallel_size=1, + expert_model_parallel_size=1, + num_nodes=1, + gpus_per_node=1, + precision="bf16", + micro_batch_size=1, + megatron_cfg_overrides={}, + ) + + teacher = TeacherWorkerGroup( + teacher_config, + cluster, + policy_config, + MagicMock(), + ) + + assert captured["cfg"]["router_replay"]["enabled"] is False + assert teacher.cfg["router_replay"]["enabled"] is False + assert policy_config["router_replay"]["enabled"] is True + + +def test_get_logprobs_from_meta_dispatches_tq_shards_to_teacher_workers(): + """TeacherWorkerGroup sends metadata, not token tensors, to each DP rank.""" + from nemo_rl.models.policy.teacher_worker_group import TeacherWorkerGroup + + class Sharding: + def get_axis_size(self, axis): + assert axis == "data_parallel" + return 2 + + worker_group = MagicMock() + worker_group.run_all_workers_sharded_data.return_value = "futures" + teacher = object.__new__(TeacherWorkerGroup) + teacher.alias = "teacher" + teacher.use_sequence_packing = False + teacher.use_dynamic_batches = False + teacher.sequence_length_pad_multiple = 2 + teacher.sharding_annotations = Sharding() + teacher.worker_group = worker_group + teacher._micro_batch_size = 1 + meta = KVBatchMeta( + partition_id="rollout_data", + task_name="train", + sample_ids=["a", "b"], + fields=["input_ids", "input_lengths"], + sequence_lengths=[3, 5], + ) + + teacher.get_logprobs_from_meta(meta) + + call = worker_group.run_all_workers_sharded_data.call_args + kwargs = call.kwargs + assert call.args == ("get_teacher_logprobs_presharded",) + assert [shard.sample_ids for shard in kwargs["meta"]] == [["a"], ["b"]] + assert all(shard.fields == list(TEACHER_LP_FIELDS) for shard in kwargs["meta"]) + assert all( + shard.extra_info[GLOBAL_FORWARD_PAD_SEQLEN] == 6 for shard in kwargs["meta"] + ) + worker_group.get_all_worker_results.assert_called_once_with("futures") + + +def test_get_logprobs_from_meta_builds_global_dynamic_batch_plan(monkeypatch): + """Teacher presharding uses logprob tokens and ships one balanced global plan.""" + import nemo_rl.models.policy.teacher_worker_group as teacher_module + from nemo_rl.models.policy.teacher_worker_group import TeacherWorkerGroup + + class Sharding: + def get_axis_size(self, axis): + assert axis == "data_parallel" + return 2 + + captured = {} + real_shard_meta_for_dp = teacher_module.shard_meta_for_dp + + def capture_plan(meta, **kwargs): + captured.update(kwargs) + captured["meta"] = meta + return real_shard_meta_for_dp(meta, **kwargs) + + monkeypatch.setattr(teacher_module, "shard_meta_for_dp", capture_plan) + worker_group = MagicMock() + teacher = object.__new__(TeacherWorkerGroup) + teacher.alias = "teacher" + teacher.use_sequence_packing = False + teacher.use_dynamic_batches = True + teacher.sequence_length_pad_multiple = 1 + teacher.dynamic_batching_args = { + "input_key": "input_ids", + "input_lengths_key": "input_lengths", + "sequence_length_round": 64, + "max_tokens_per_microbatch": 999, + } + teacher.cfg = { + "dynamic_batching": { + "enabled": True, + "train_mb_tokens": 999, + "logprob_mb_tokens": 256, + "sequence_length_round": 64, + } + } + teacher.sharding_annotations = Sharding() + teacher.worker_group = worker_group + teacher._micro_batch_size = 1 + meta = KVBatchMeta( + partition_id="rollout_data", + task_name="train", + sample_ids=["a", "b"], + fields=["input_ids", "input_lengths"], + sequence_lengths=[100, 60], + ) + + teacher.get_logprobs_from_meta(meta) + + assert captured["sequence_packing_args"] is None + assert captured["dynamic_batching_args"]["max_tokens_per_microbatch"] == 256 + assert captured["meta"].extra_info[GLOBAL_FORWARD_PAD_SEQLEN] == 128 + shards = worker_group.run_all_workers_sharded_data.call_args.kwargs["meta"] + assert sorted(sample_id for shard in shards for sample_id in shard.sample_ids) == [ + "a", + "b", + ] + assert all(MICRO_BATCH_INDICES in shard.extra_info for shard in shards) + assert all(MICRO_BATCH_LENGTHS in shard.extra_info for shard in shards) + assert ( + len( + { + len(microbatch_lengths) + for shard in shards + for microbatch_lengths in shard.extra_info[MICRO_BATCH_LENGTHS] + } + ) + == 1 + ) + assert all( + microbatch_length <= 128 + for shard in shards + for microbatch_lengths in shard.extra_info[MICRO_BATCH_LENGTHS] + for microbatch_length in microbatch_lengths + ) + + +def test_get_logprobs_from_meta_builds_global_sequence_packing_plan(monkeypatch): + """Teacher packing plans use logprob tokens and skip global pad rounding.""" + import nemo_rl.models.policy.teacher_worker_group as teacher_module + from nemo_rl.models.policy.teacher_worker_group import TeacherWorkerGroup + + class Sharding: + def get_axis_size(self, axis): + assert axis == "data_parallel" + return 2 + + captured = {} + real_shard_meta_for_dp = teacher_module.shard_meta_for_dp + + def capture_plan(meta, **kwargs): + captured.update(kwargs) + captured["meta"] = meta + return real_shard_meta_for_dp(meta, **kwargs) + + monkeypatch.setattr(teacher_module, "shard_meta_for_dp", capture_plan) + worker_group = MagicMock() + teacher = object.__new__(TeacherWorkerGroup) + teacher.alias = "teacher" + teacher.use_sequence_packing = True + teacher.use_dynamic_batches = False + teacher.sequence_length_pad_multiple = 16 + teacher.sequence_packing_args = { + "algorithm": "modified_first_fit_decreasing", + "input_key": "input_ids", + "input_lengths_key": "input_lengths", + "sequence_length_pad_multiple": 16, + } + teacher.cfg = { + "sequence_packing": { + "enabled": True, + "train_mb_tokens": 999, + "logprob_mb_tokens": 16, + } + } + teacher.sharding_annotations = Sharding() + teacher.worker_group = worker_group + teacher._micro_batch_size = 1 + meta = KVBatchMeta( + partition_id="rollout_data", + task_name="train", + sample_ids=["a", "b", "c", "d"], + fields=["input_ids", "input_lengths"], + sequence_lengths=[8, 1, 7, 2], + ) + + teacher.get_logprobs_from_meta(meta) + + assert captured["dynamic_batching_args"] is None + assert captured["sequence_packing_args"]["max_tokens_per_microbatch"] == 16 + assert captured["meta"].extra_info[GLOBAL_FORWARD_PAD_SEQLEN] == 8 + shards = worker_group.run_all_workers_sharded_data.call_args.kwargs["meta"] + assert sorted(sample_id for shard in shards for sample_id in shard.sample_ids) == [ + "a", + "b", + "c", + "d", + ] + assert all(MICRO_BATCH_INDICES in shard.extra_info for shard in shards) + assert all(MICRO_BATCH_LENGTHS in shard.extra_info for shard in shards) + assert ( + len( + { + len(microbatch_lengths) + for shard in shards + for microbatch_lengths in shard.extra_info[MICRO_BATCH_LENGTHS] + } + ) + == 1 + ) + + +def test_teacher_worker_presharded_entrypoint_writes_teacher_tq_field(): + """The worker consumes its TQ shard and writes only the teacher delta.""" + from nemo_rl.data_plane.worker_mixin import TQWorkerMixin + + class Worker(TQWorkerMixin): + cfg = {"sequence_packing": {"enabled": False}} + + def __init__(self): + self.written = None + self.received_batching_metadata = None + + def _fetch(self, meta): + del meta + return BatchedDataDict( + { + "input_ids": torch.ones(1, 3, dtype=torch.long), + "input_lengths": torch.tensor([3]), + } + ) + + def get_logprobs(self, data, micro_batch_size=None): + del micro_batch_size + self.received_batching_metadata = ( + data.micro_batch_indices, + data.micro_batch_lengths, + ) + return BatchedDataDict({"logprobs": torch.full((1, 3), 0.25)}) + + def _write_back_result_field(self, meta, result, *, result_key, tq_field): + self.written = (meta, result[result_key], tq_field) + + meta = KVBatchMeta( + partition_id="rollout_data", + task_name="teacher_lp:teacher", + sample_ids=["a"], + fields=["input_ids", "input_lengths"], + sequence_lengths=[3], + extra_info={ + MICRO_BATCH_INDICES: [[0]], + MICRO_BATCH_LENGTHS: [3], + }, + ) + worker = Worker() + + worker.get_teacher_logprobs_presharded(meta) + + assert worker.written is not None + assert worker.written[0] is meta + assert torch.allclose(worker.written[1], torch.full((1, 3), 0.25)) + assert worker.written[2] == "teacher_reference_logprobs" + assert worker.received_batching_metadata == ([[0]], [3]) + + +def test_teacher_worker_rejects_local_dynamic_batch_planning(): + """A missing driver plan fails instead of independently repacking each DP rank.""" + from nemo_rl.data_plane.worker_mixin import TQWorkerMixin + + class Worker(TQWorkerMixin): + cfg = { + "sequence_packing": {"enabled": False}, + "dynamic_batching": { + "enabled": True, + "sequence_length_round": 1, + "train_mb_tokens": 8, + }, + } + + def _fetch(self, meta): + del meta + return BatchedDataDict( + { + "input_ids": torch.ones(1, 3, dtype=torch.long), + "input_lengths": torch.tensor([3]), + } + ) + + meta = KVBatchMeta( + partition_id="rollout_data", + task_name="teacher_lp:teacher", + sample_ids=["a"], + fields=["input_ids", "input_lengths"], + sequence_lengths=[3], + ) + + with pytest.raises(RuntimeError, match="driver-provided global"): + Worker().get_teacher_logprobs_presharded(meta) diff --git a/tests/unit/single_controller/test_ppo_setup.py b/tests/unit/single_controller/test_ppo_setup.py index d63de950e1f..6fe7853f3bf 100644 --- a/tests/unit/single_controller/test_ppo_setup.py +++ b/tests/unit/single_controller/test_ppo_setup.py @@ -134,6 +134,7 @@ def _make_master_config( }, loss_fn=ClippedPGLossConfig(reference_policy_kl_penalty=0.0), env={}, + cluster={"num_nodes": 2, "gpus_per_node": 8, "segment_size": None}, async_rl=AsyncRLConfig( min_groups_for_streaming_train=min_groups_for_streaming_train, max_buffered_rollouts=_NUM_PROMPTS_PER_STEP * 2, @@ -484,7 +485,11 @@ def patched_ppo_factories(): patch.object( sc_setup_mod, "_build_clusters", - return_value=(MagicMock(name="train"), MagicMock(name="inference")), + return_value=( + MagicMock(name="train"), + MagicMock(name="inference"), + None, + ), ), patch.object( sc_setup_mod, "_build_generation", return_value=(MagicMock(), 0.0) @@ -619,7 +624,7 @@ def test_worker_group_slots( ): mc = _cluster_config(make_config(), colocated=colocated, backend=backend) - train, inference = sc_setup_mod._build_clusters(mc) + train, inference, _teacher_topology = sc_setup_mod._build_clusters(mc) assert train.kwargs["max_colocated_worker_groups"] == expected_train_groups if colocated: diff --git a/tests/unit/single_controller/test_setup.py b/tests/unit/single_controller/test_setup.py index 1a026b5b955..cd56bcf7653 100644 --- a/tests/unit/single_controller/test_setup.py +++ b/tests/unit/single_controller/test_setup.py @@ -16,24 +16,30 @@ from __future__ import annotations +from pathlib import Path from unittest.mock import MagicMock, patch import pytest +from omegaconf import OmegaConf import nemo_rl.algorithms.single_controller_utils.setup as sc_setup_mod +from nemo_rl.algorithms.advantage_estimator import AdvEstimatorConfig from nemo_rl.algorithms.async_utils.staleness_sampler import ( ReadyFirstSamplerConfig, SamplerConfig, ) from nemo_rl.algorithms.grpo import GRPOConfig from nemo_rl.algorithms.loss import ClippedPGLossConfig +from nemo_rl.algorithms.opd import OnPolicyDistillationConfig from nemo_rl.algorithms.single_controller_utils import ( AsyncRLConfig, MasterConfig, SingleControllerActorArgs, setup_single_controller, ) +from nemo_rl.data_plane.schema import SC_ROLLOUT_SCHEMA_FIELDS from nemo_rl.experience.rollouts import EffortLevelsConfig +from nemo_rl.utils.config import load_config, register_omegaconf_resolvers def _make_master_config( @@ -97,6 +103,7 @@ def _make_master_config( "save_period": 10, "save_optimizer": False, }, + cluster={"num_nodes": 2, "gpus_per_node": 8, "segment_size": None}, loss_fn=loss_cfg if loss_cfg is not None else ClippedPGLossConfig(), env=env if env is not None else {}, async_rl=AsyncRLConfig( @@ -140,6 +147,7 @@ def patched_factories(): return_value=( MagicMock(name="train_cluster"), MagicMock(name="inference_cluster"), + None, ), ) as mock_clusters, patch.object( @@ -223,6 +231,131 @@ def test_build_clusters_rejects_non_colocated_megatron_generation(): sc_setup_mod._build_clusters(master_config) +def test_build_clusters_rejects_unsupported_topology_backend(monkeypatch): + """Topology planning reports the supported SC backends instead of KeyError.""" + master_config = _make_master_config(colocated=False, backend="trtllm") + master_config.cluster = {"num_nodes": 2, "gpus_per_node": 8, "segment_size": 1} + master_config.policy["generation"]["colocated"]["resources"] = { + "gpus_per_node": 8, + "num_nodes": 1, + } + monkeypatch.setattr( + sc_setup_mod, + "prepare_segment_topology", + lambda *args, **kwargs: ( + [{"nvlink_domain": 0.001}], + ["inference"], + { + "training": ("nvlink_domain", 0), + "inference": ("nvlink_domain", 1), + }, + ), + ) + + with pytest.raises( + ValueError, + match="only supports vllm or sglang generation; got 'trtllm'", + ): + sc_setup_mod._build_clusters(master_config) + + +def test_build_clusters_leaves_dedicated_teacher_nodes(monkeypatch): + """Teacher nodes are removed before the student train/inference split.""" + master_config = _make_master_config(colocated=False) + master_config.cluster = {"num_nodes": 3, "gpus_per_node": 8} + master_config.policy["generation"]["colocated"]["resources"] = { + "gpus_per_node": 8, + "num_nodes": 1, + } + master_config.on_policy_distillation = OnPolicyDistillationConfig( + enabled=True, + teacher_model_by_agent_name={"default_teacher": "Qwen/Qwen3-1.7B"}, + default_teacher_alias="default_teacher", + non_colocated_teachers={"enabled": True}, + ) + constructed = [] + + class FakeCluster: + def __init__(self, **kwargs): + self.kwargs = kwargs + constructed.append(kwargs) + + monkeypatch.setattr(sc_setup_mod, "RayVirtualCluster", FakeCluster) + + _, _, teacher_topology = sc_setup_mod._build_clusters(master_config) + + assert constructed[0]["bundle_ct_per_node_list"] == [8] + assert constructed[1]["bundle_ct_per_node_list"] == [8] + assert teacher_topology is None + + +def test_build_clusters_supports_two_node_shared_student_layout(monkeypatch): + """One student node can split train/inference while node two hosts teacher.""" + master_config = _make_master_config(colocated=False) + master_config.cluster = {"num_nodes": 2, "gpus_per_node": 8} + master_config.policy["generation"]["colocated"]["resources"] = { + "gpus_per_node": 4, + "num_nodes": 1, + } + master_config.on_policy_distillation = OnPolicyDistillationConfig( + enabled=True, + teacher_model_by_agent_name={"default": "/ckpt/teacher"}, + default_teacher_alias="default", + non_colocated_teachers={ + "enabled": True, + "default_teacher_cfg": {"num_nodes": 1, "gpus_per_node": 8}, + }, + ) + constructed = [] + + class FakeCluster: + def __init__(self, **kwargs): + self.kwargs = kwargs + constructed.append(self) + + monkeypatch.setattr(sc_setup_mod, "RayVirtualCluster", FakeCluster) + + train_cluster, inference_cluster, teacher_topology = sc_setup_mod._build_clusters( + master_config + ) + + assert train_cluster.kwargs["bundle_ct_per_node_list"] == [4] + assert inference_cluster.kwargs["bundle_ct_per_node_list"] == [4] + assert teacher_topology is None + + +def test_single_controller_mopd_recipe_resolves_to_runtime_contract(): + """The inherited recipe resolves exactly as the SC entrypoint consumes it.""" + register_omegaconf_resolvers() + repo_root = Path(__file__).resolve().parents[3] + recipe = repo_root / ( + "examples/configs/recipes/llm/" + "mopd-qwen3-1.7b-3n8g-megatron-pack-single-controller.yaml" + ) + resolved = OmegaConf.to_container(load_config(recipe), resolve=True) + + assert isinstance(resolved, dict) + config = MasterConfig.model_validate(resolved) + assert config.grpo.async_grpo is None + assert config.grpo.adv_estimator.name == "opd" + assert config.grpo.skip_reference_policy_logprobs_calculation is True + assert config.async_rl.min_groups_for_streaming_train == ( + config.grpo.num_prompts_per_step + ) + assert config.policy["train_global_batch_size"] == ( + config.grpo.num_prompts_per_step * config.grpo.num_generations_per_prompt + ) + assert config.data_plane["enabled"] is True + assert config.env["should_use_nemo_gym"] is True + assert config.on_policy_distillation.enabled is True + assert config.on_policy_distillation.non_colocated_teachers is not None + assert config.on_policy_distillation.non_colocated_teachers.enabled is True + assert ( + config.on_policy_distillation.teacher_model_by_agent_name["default_teacher"] + == config.policy["model_name"] + ) + + class TestSetup: """setup arg validation + actor_args assembly.""" @@ -236,6 +369,123 @@ def test_multiple_dataloader_not_supported(self): with pytest.raises(NotImplementedError, match="use_multiple_dataloader"): setup_single_controller(mc, MagicMock(pad_token_id=0)) + @pytest.mark.parametrize( + ( + "opd_enabled", + "teacher_enabled", + "adv_name", + "use_nemo_gym", + "match", + ), + [ + ( + False, + False, + "opd", + True, + "requires on_policy_distillation.enabled=true", + ), + ( + True, + True, + "grpo", + True, + "requires grpo.adv_estimator.name='opd'", + ), + (True, False, "opd", True, "non_colocated_teachers.enabled=true"), + (True, True, "opd", False, "requires env.should_use_nemo_gym=true"), + ], + ) + def test_invalid_mopd_config_fails_before_allocating_resources( + self, + opd_enabled: bool, + teacher_enabled: bool, + adv_name: str, + use_nemo_gym: bool, + match: str, + patched_factories, + ): + mc = _make_master_config() + mc.env["should_use_nemo_gym"] = use_nemo_gym + mc.grpo.adv_estimator = AdvEstimatorConfig(name=adv_name) + mc.on_policy_distillation = OnPolicyDistillationConfig( + enabled=opd_enabled, + teacher_model_by_agent_name={"teacher": "/ckpt/teacher"}, + non_colocated_teachers={"enabled": teacher_enabled}, + ) + + with pytest.raises(ValueError, match=match): + setup_single_controller(mc, MagicMock(pad_token_id=0)) + + patched_factories["_build_clusters"].assert_not_called() + + def test_mopd_reserves_before_models_and_initializes_teacher_last( + self, patched_factories, monkeypatch + ): + mc = _make_master_config(env={"should_use_nemo_gym": True}) + mc.cluster = {"num_nodes": 3, "gpus_per_node": 8} + mc.policy["generation"]["vllm_cfg"] = { + "async_engine": True, + "expose_http_server": True, + } + mc.policy["generation"].update( + {"stop_strings": None, "stop_token_ids": None, "top_k": None} + ) + mc.grpo.adv_estimator = AdvEstimatorConfig(name="opd") + mc.on_policy_distillation = OnPolicyDistillationConfig( + enabled=True, + teacher_model_by_agent_name={"teacher": "/ckpt/teacher"}, + default_teacher_alias="teacher", + non_colocated_teachers={"enabled": True}, + ) + events = [] + teacher_cluster = MagicMock(name="teacher_cluster") + teacher_group = MagicMock(name="teacher_group") + + def reserve_teachers(*args, **kwargs): + del args, kwargs + events.append("reserve_teacher") + return {"teacher": teacher_cluster} + + original_build_generation = patched_factories["_build_generation"].return_value + + def build_generation(*args, **kwargs): + del args, kwargs + events.append("build_generation") + return original_build_generation + + def create_teachers(*args, **kwargs): + del args, kwargs + events.append("create_teacher") + return {"teacher": teacher_group}, {"teacher": "teacher"} + + patched_factories["_build_generation"].side_effect = build_generation + monkeypatch.setattr( + sc_setup_mod.opd_module, + "reserve_teacher_clusters", + reserve_teachers, + ) + monkeypatch.setattr( + sc_setup_mod.opd_module, + "create_teacher_worker_groups", + create_teachers, + ) + patched_factories["setup_response_data"].return_value = (list(range(8)), None) + monkeypatch.setattr( + sc_setup_mod, + "_spinup_gym", + lambda **_kwargs: (MagicMock(name="nemo_gym_actor"), 0.0), + ) + + actor_args, timings = setup_single_controller(mc, MagicMock(pad_token_id=17)) + + assert events == ["reserve_teacher", "build_generation", "create_teacher"] + assert actor_args.teacher_worker_groups == {"teacher": teacher_group} + assert actor_args.alias_to_group_alias == {"teacher": "teacher"} + teacher_group.setup_data_plane.assert_called_once_with(mc.data_plane) + assert timings.teacher_reservation_time_s is not None + assert timings.teacher_model_init_time_s is not None + @pytest.mark.parametrize( ("invalid_case", "match"), [ @@ -339,6 +589,41 @@ def test_returns_actor_args(self, patched_factories): assert actor_args.partition_id == "rollout_data" assert actor_args.tq_buffer._partition_id == "rollout_data" assert actor_args.tq_buffer._require_routed_experts is False + actor_args.dp_client.register_partition.assert_called_once() + warmup = actor_args.dp_client.register_partition.call_args.kwargs + assert warmup["partition_id"] == "rollout_data" + assert set(SC_ROLLOUT_SCHEMA_FIELDS) <= set(warmup["fields"]) + assert "teacher_reference_logprobs" in warmup["fields"] + assert warmup["num_samples"] == 16 + assert warmup["grpo_group_size"] == 2 + + def test_reserves_topology_constrained_training_before_builds( + self, patched_factories + ): + mc = _make_master_config(colocated=False) + mc.cluster = {"num_nodes": 2, "gpus_per_node": 8, "segment_size": 1} + mc.policy["generation"]["colocated"]["resources"] = { + "gpus_per_node": 4, + "num_nodes": 1, + } + train_cluster = patched_factories["_build_clusters"].return_value[0] + events = [] + train_cluster.get_placement_groups.side_effect = lambda: events.append( + "reserve_train" + ) + original_build = patched_factories["_build_generation"].return_value + + def build_generation(*args, **kwargs): + del args, kwargs + events.append("build_generation") + return original_build + + patched_factories["_build_generation"].side_effect = build_generation + + setup_single_controller(mc, MagicMock(pad_token_id=0)) + + assert events[0] == "reserve_train" + train_cluster.get_placement_groups.assert_called_once_with() def test_effort_levels_reach_the_rollout_manager(self, patched_factories): """env.nemo_gym.effort_levels is resolved into RolloutManager's kwarg. diff --git a/tests/unit/single_controller/test_single_controller_actor.py b/tests/unit/single_controller/test_single_controller_actor.py index 739250d0a79..6c711a22f36 100644 --- a/tests/unit/single_controller/test_single_controller_actor.py +++ b/tests/unit/single_controller/test_single_controller_actor.py @@ -29,7 +29,10 @@ 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 import ( + SingleControllerActor, + _pooled_opd_metrics, +) from nemo_rl.algorithms.single_controller_utils.config import ( AdvantageConfig, AsyncRLConfig, @@ -574,6 +577,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._teacher_logprobs_required = False ctrl._is_ppo = False ctrl._master_config = SimpleNamespace( grpo=SimpleNamespace(seq_logprob_error_threshold=2.0) @@ -642,6 +646,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._teacher_logprobs_required = False ctrl._is_ppo = False ctrl._master_config = SimpleNamespace( grpo=SimpleNamespace(seq_logprob_error_threshold=None) @@ -704,6 +709,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._teacher_logprobs_required = False ctrl._is_ppo = False ctrl._master_config = SimpleNamespace( grpo=SimpleNamespace(seq_logprob_error_threshold=2.0) @@ -759,6 +765,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._teacher_logprobs_required = False ctrl._is_ppo = False ctrl._master_config = SimpleNamespace( grpo=SimpleNamespace(seq_logprob_error_threshold=None) @@ -793,6 +800,111 @@ def test_advantage_stage_skips_preexisting_empty_mask_without_seq_threshold() -> assert "advantages" in (result_meta.fields or []) +def test_opd_advantage_stage_reads_teacher_and_student_logprobs() -> None: + """SC passes the TQ teacher column under OPD's estimator contract.""" + controller_cls = SingleControllerActor.__ray_metadata__.modified_class + ctrl = object.__new__(controller_cls) + captured_kwargs = {} + + class FakeEstimator: + def compute_advantage(self, **kwargs): + captured_kwargs.update(kwargs) + return kwargs["teacher_logprobs"] - kwargs["prev_logprobs"] + + class FakeDataPlane: + def __init__(self): + self.put_fields = None + + def get_samples(self, sample_ids, partition_id, select_fields): + del sample_ids, partition_id + assert "teacher_reference_logprobs" in select_fields + assert "generation_logprobs" in select_fields + return TensorDict( + { + "prompt_ids_for_adv": torch.zeros(2, 3, dtype=torch.long), + "total_reward": torch.zeros(2), + "token_mask": torch.tensor([[1.0, 1.0, 1.0], [1.0, 0.0, 0.0]]), + "sample_mask": torch.ones(2), + "generation_logprobs": torch.full((2, 3), 0.5), + "prev_logprobs": torch.full((2, 3), 0.5), + "teacher_reference_logprobs": torch.full((2, 3), 0.75), + }, + batch_size=(2,), + ) + + def put_samples(self, sample_ids, partition_id, fields): + del sample_ids, partition_id + self.put_fields = fields + + ctrl._advantage_cfg = AdvantageConfig() + ctrl._advantage_estimator = FakeEstimator() + ctrl._policy_logprobs_required = True + ctrl._reference_logprobs_required = False + ctrl._teacher_logprobs_required = True + ctrl._is_ppo = False + ctrl._dp_client = FakeDataPlane() + 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": [], + "sequence_lengths": [], + "seq_logprob_error_metrics": [], + } + ctrl._opd_stat_sum = 0.0 + ctrl._opd_stat_sumsq = 0.0 + ctrl._opd_stat_count = 0 + meta = KVBatchMeta( + partition_id="rollout_data", + task_name="train", + sample_ids=["a", "b"], + fields=[], + sequence_lengths=[3, 3], + ) + + enriched, has_valid_training_tokens = asyncio.run(ctrl._advantage_stage(meta)) + + assert has_valid_training_tokens + assert set(captured_kwargs) >= { + "teacher_logprobs", + "prev_logprobs", + "prompt_ids", + "rewards", + "mask", + "repeated_batch", + } + assert "logprobs_policy" not in captured_kwargs + assert torch.allclose( + captured_kwargs["teacher_logprobs"] - captured_kwargs["prev_logprobs"], + torch.full((2, 3), 0.25), + ) + assert "advantages" in (enriched.fields or []) + assert ctrl._opd_stat_sum == pytest.approx(1.0) + assert ctrl._opd_stat_sumsq == pytest.approx(0.25) + assert ctrl._opd_stat_count == 4 + + +def test_pooled_opd_metrics_weight_unequal_chunks_by_valid_token_count() -> None: + """A small streaming chunk cannot receive the same weight as a large one.""" + # Chunk 1 has values [0, 2]; chunk 2 has [4]. Averaging chunk means + # would incorrectly produce 2.5. Exact pooling produces mean=2, std=2. + metrics = _pooled_opd_metrics( + stat_sum=6.0, + stat_sumsq=20.0, + count=3, + ) + + assert metrics == pytest.approx( + { + "on_policy_distillation/teacher_student_logprob_gap_mean": 2.0, + "on_policy_distillation/adv_mean": 2.0, + "on_policy_distillation/adv_std": 2.0, + } + ) + + class _EmptySampler: async def evict(self, *, current_train_weight: int) -> int: del current_train_weight @@ -974,6 +1086,7 @@ def _train_pump_controller(*, sampler) -> object: ctrl._advantage_cfg = AdvantageConfig() ctrl._policy_logprobs_required = False ctrl._reference_logprobs_required = False + ctrl._teacher_logprobs_required = False ctrl._advantage_estimator = None ctrl._partition_id = "rollout_data" ctrl._sampler = sampler @@ -1000,7 +1113,12 @@ def _train_pump_controller(*, sampler) -> object: "rewards": [], "masked_advantages": [], "sequence_lengths": [], + "seq_logprob_error_metrics": [], } + ctrl._opd_stat_sum = 0.0 + ctrl._opd_stat_sumsq = 0.0 + ctrl._opd_stat_count = 0 + ctrl._teacher_coordinator = None return ctrl @@ -1631,6 +1749,7 @@ def compute_advantage(self, *, rewards, mask, **kwargs): ctrl._advantage_estimator = estimator ctrl._policy_logprobs_required = False ctrl._reference_logprobs_required = False + ctrl._teacher_logprobs_required = False ctrl._is_ppo = True ctrl._master_config = SimpleNamespace( ppo=SimpleNamespace(seq_logprob_error_threshold=None) diff --git a/tests/unit/single_controller/test_tq_replay_buffer.py b/tests/unit/single_controller/test_tq_replay_buffer.py index 35475caeb5a..80081cbdd8c 100644 --- a/tests/unit/single_controller/test_tq_replay_buffer.py +++ b/tests/unit/single_controller/test_tq_replay_buffer.py @@ -25,7 +25,10 @@ from tensordict import TensorDict import nemo_rl.algorithms.async_utils.replay_buffer as _replay_buffer_module -from nemo_rl.algorithms.async_utils.replay_buffer import TQReplayBuffer +from nemo_rl.algorithms.async_utils.replay_buffer import ( + PostWriteEnrichmentError, + TQReplayBuffer, +) from nemo_rl.data_plane import KVBatchMeta from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.experience.interfaces import PromptGroupRecord @@ -186,6 +189,48 @@ def _add_group( class TestTQReplayBufferReserveCommit: + def test_commit_enriches_after_put_before_slot_becomes_ready(self): + dp = FakeDataPlaneClient() + buf = _make_buffer(dp) + observations = [] + + async def enrich(meta, record): + del record + observations.append((dp.depth(), list(buf.ready_list))) + return meta.with_fields(["teacher_reference_logprobs"]) + + buf.set_post_write_enricher(enrich) + meta = _add_group(buf, weight=3) + + assert observations == [(_N_GENS, [False])] + assert "teacher_reference_logprobs" in meta.fields + assert buf.ready_list == [True] + + def test_commit_rolls_back_when_post_write_enrichment_fails(self): + dp = FakeDataPlaneClient() + buf = _make_buffer(dp) + + async def fail_enrichment(meta, record): + del meta, record + raise RuntimeError("teacher unavailable") + + buf.set_post_write_enricher(fail_enrichment) + group_id = buf.reserve(weight_version=3) + + with pytest.raises(PostWriteEnrichmentError, match="post-write enrichment"): + _run( + buf.commit( + group_id, + _make_record(), + start_weight_version=3, + end_weight_version=3, + ) + ) + + assert dp.depth() == 0 + assert buf.ready_list == [False] + assert buf.meta_list == [None] + def test_commit_clears_rows_when_put_raises_after_writing(self): dp = FailAfterPutDataPlaneClient() buf = _make_buffer(dp) diff --git a/tests/unit/test_recipes_and_test_suites.py b/tests/unit/test_recipes_and_test_suites.py index 35f14c0c22b..175942e5d51 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_4149_hours(nightly_test_suite, tracker): +def test_nightly_compute_stays_below_4157_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,10 @@ def test_nightly_compute_stays_below_4149_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 <= 4149, ( - f"Total GPU hours exceeded 4149: {last_line}. We should revisit the test suites to reduce the total GPU hours." + # Reserve 8 GPU-hours for the SingleController MOPD nightly on top of + # main's 4149-hour limit. + assert total_gpu_hours <= 4157, ( + f"Total GPU hours exceeded 4157: {last_line}. We should revisit the test suites to reduce the total GPU hours." ) tracker.track("total_nightly_gpu_hours", total_gpu_hours)