From 14045abf0570ad4fc224eec3225232410ef2f652 Mon Sep 17 00:00:00 2001 From: Yi-Fu Wu Date: Fri, 21 Aug 2026 02:52:23 -0700 Subject: [PATCH 1/5] feat(mopd): add single-controller text path Route teacher logprob inference through TransferQueue with deduplicated per-teacher serialization and wire OPD advantages into the single controller. Add configuration validation, failure cleanup, unit coverage, and a same-model nightly recipe. Signed-off-by: Yi-Fu Wu --- ...-3n8g-megatron-pack-single-controller.yaml | 32 +++ examples/run_grpo_single_controller.py | 6 + .../algorithms/async_utils/replay_buffer.py | 24 +- nemo_rl/algorithms/metric_utils.py | 3 +- nemo_rl/algorithms/opd.py | 228 ++++++++++++++++++ nemo_rl/algorithms/single_controller.py | 47 +++- .../single_controller_utils/config.py | 30 +++ .../single_controller_utils/setup.py | 186 +++++++++++++- nemo_rl/data_plane/column_io.py | 1 + nemo_rl/data_plane/worker_mixin.py | 26 +- nemo_rl/experience/rollout_manager.py | 10 +- nemo_rl/models/policy/teacher_worker_group.py | 67 ++++- ...7b-3n8g-megatron-pack-single-controller.sh | 60 +++++ tests/test_suites/nightly.txt | 4 + tests/unit/algorithms/test_opd.py | 167 +++++++++++++ tests/unit/experience/test_rollout_manager.py | 42 ++++ .../policy/test_teacher_worker_group.py | 89 +++++++ .../test_single_controller.py | 81 +++++++ .../test_single_controller_setup.py | 116 +++++++++ .../test_tq_replay_buffer.py | 47 +++- 20 files changed, 1246 insertions(+), 20 deletions(-) create mode 100644 examples/configs/recipes/llm/mopd-qwen3-1.7b-3n8g-megatron-pack-single-controller.yaml create mode 100755 tests/test_suites/llm/mopd-qwen3-1.7b-3n8g-megatron-pack-single-controller.sh 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..dc5559d8f34 --- /dev/null +++ b/examples/configs/recipes/llm/mopd-qwen3-1.7b-3n8g-megatron-pack-single-controller.yaml @@ -0,0 +1,32 @@ +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 + val_at_start: false + val_at_end: false + # 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 + keep_top_k: 1 + +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 fef3b2294c1..e3479fa4a81 100644 --- a/examples/run_grpo_single_controller.py +++ b/examples/run_grpo_single_controller.py @@ -158,6 +158,12 @@ def main() -> None: except Exception as kill_error: print(f"Env {env_name!r} kill failed: {kill_error}") + for teacher_alias, teacher in (actor_args.teacher_worker_groups or {}).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 2b67a4ce93d..7be63fe7c45 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 @@ -36,6 +36,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. @@ -737,6 +741,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, @@ -827,6 +841,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 8c7cd7c0ffe..dbc40396856 100644 --- a/nemo_rl/algorithms/metric_utils.py +++ b/nemo_rl/algorithms/metric_utils.py @@ -41,7 +41,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..7c42783ccb2 100644 --- a/nemo_rl/algorithms/opd.py +++ b/nemo_rl/algorithms/opd.py @@ -21,15 +21,23 @@ from __future__ import annotations +import asyncio +import threading +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.distributed.virtual_cluster import ( RayVirtualCluster, prepare_segment_topology, ) +from nemo_rl.experience.interfaces import PromptGroupRecord # --------------------------------------------------------------------------- # Config schemas @@ -200,6 +208,226 @@ 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: threading.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() + self._models_seen: set[str] = set() + + def _resolve_teacher(self, record: PromptGroupRecord) -> tuple[str, str]: + extra_env_info = record.extra_env_info + if not isinstance(extra_env_info, dict): + raise ValueError( + "MOPD rollout is missing prompt-level extra_env_info with agent_ref" + ) + agent_ref = extra_env_info.get("agent_ref") + if not isinstance(agent_ref, dict): + raise ValueError( + "MOPD rollout extra_env_info must contain an agent_ref mapping" + ) + + 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, + ) -> tuple[float, 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] = [] + 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=["input_ids", "input_lengths"], + 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=["input_ids", "input_lengths"], + sequence_lengths=[meta.sequence_lengths[-1]] * pad_count, + ) + 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) + + lock_started_at = time.perf_counter() + try: + with self._teacher_locks[group_alias]: + inference_started_at = time.perf_counter() + teacher.get_logprobs_from_meta(padded_meta) + except BaseException as inference_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 inference and temporary-row cleanup both failed " + f"for group {group_alias!r}", + [inference_error, cleanup_error], + ) + raise + inference_finished_at = time.perf_counter() + if temporary_sample_ids: + self._dp_client.clear_samples( + sample_ids=temporary_sample_ids, + partition_id=meta.partition_id, + ) + return ( + inference_started_at - lock_started_at, + 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_wait_s, inference_time_s = await asyncio.to_thread( + self._enrich_sync, meta, group_alias + ) + total_time_s = time.perf_counter() - started_at + + teacher_model_by_agent_name = self._opd_cfg["teacher_model_by_agent_name"] + 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) + self._models_seen.add(teacher_model_by_agent_name[alias]) + record.rollout_metrics["teacher_logprob_time"] = total_time_s + 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.""" + alias_unique = len(self._aliases_seen) + model_unique = len(self._models_seen) + 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, + "on_policy_distillation/teacher_alias_unique": float(alias_unique), + "on_policy_distillation/teacher_model_unique": float(model_unique), + "on_policy_distillation/teacher_alias_to_model_compression": float( + model_unique / max(alias_unique, 1) + ), + } + 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() + self._models_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 f9fa674366f..e443c6865fd 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -47,6 +47,7 @@ import torch from nemo_rl.algorithms.async_utils.staleness_sampler import create_sampler +from nemo_rl.algorithms import opd as opd_module from nemo_rl.algorithms.grpo import GRPOSaveState, _write_latest_checkpoint_status from nemo_rl.algorithms.metric_utils import SetupTimingMetrics from nemo_rl.algorithms.single_controller_utils.config import ( @@ -123,6 +124,7 @@ def __init__( self._master_config = master_config self._async_cfg = master_config.async_rl + self._opd_enabled = opd_module.is_opd_enabled(master_config) self._policy_logprobs_required = not ( master_config.loss_fn.force_on_policy_ratio and master_config.grpo.seq_logprob_error_threshold is None @@ -154,6 +156,21 @@ def __init__( # Rebind so writer and sampler share one buffer instance even # when Ray deserializes rollout_manager and tq_buffer separately. self._rollout_manager._tq_buffer = self._buffer + 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. @@ -257,6 +274,7 @@ def __init__( "masked_advantages": [], "sequence_lengths": [], } + self._advantage_metric_values: dict[str, list[float]] = {} print( f"SingleControllerActor: " @@ -1055,6 +1073,16 @@ 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( + { + name: float(sum(values) / len(values)) + for name, values in self._advantage_metric_values.items() + if values + } + ) + self._advantage_metric_values.clear() + if self._teacher_coordinator is not None: + step_metrics.update(self._teacher_coordinator.drain_metrics()) self._trainer_version += 1 self._train_steps += 1 @@ -1693,15 +1721,21 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> KVBatchMeta: 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._opd_enabled: + 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._opd_enabled: + kwargs["teacher_logprobs"] = tensor_field( + data, + adv_cfg.teacher_logprobs_field, + ) advantages = self._advantage_estimator.compute_advantage( prompt_ids=prompt_ids, @@ -1715,6 +1749,9 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> KVBatchMeta: self._step_log_dict["masked_advantages"].append( response_advantages.detach().cpu() ) + estimator_metrics = getattr(self._advantage_estimator, "last_metrics", {}) + for name, value in estimator_metrics.items(): + self._advantage_metric_values.setdefault(name, []).append(float(value)) await self._call_dp( "put_samples", @@ -1742,4 +1779,6 @@ def _advantage_input_fields(self) -> list[str]: fields.append(adv_cfg.policy_logprobs_field) if self._reference_logprobs_required: fields.append(adv_cfg.reference_logprobs_field) + if self._opd_enabled: + fields.append(adv_cfg.teacher_logprobs_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 c57f7b88db1..9ad793d7ac5 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/nemo_rl/algorithms/single_controller_utils/config.py @@ -36,6 +36,8 @@ ) from nemo_rl.algorithms.grpo import GRPOConfig, GRPOLoggerConfig from nemo_rl.algorithms.loss import ClippedPGLossConfig +from nemo_rl.algorithms import opd as opd_module +from nemo_rl.algorithms.opd import OnPolicyDistillationConfig from nemo_rl.data import DataConfig from nemo_rl.data_plane.interfaces import DataPlaneConfig from nemo_rl.distributed.virtual_cluster import ClusterConfig @@ -544,6 +546,7 @@ class MasterConfig(BaseModel, extra="allow"): checkpointing: CheckpointingConfig data_plane: DataPlaneConfig async_rl: AsyncRLConfig + on_policy_distillation: Optional[OnPolicyDistillationConfig] = None def validate_sampler_buffer_capacity( @@ -758,6 +761,32 @@ def validate_single_controller_config(master_config: MasterConfig) -> None: "loss_fn.reference_policy_kl_penalty=0." ) + opd_enabled = opd_module.is_opd_enabled(master_config) + if master_config.grpo.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 master_config.grpo.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 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) + _validate_failure_settings(async_config, num_prompts_per_step) # Nesting says which knob applies to which path, but nothing stops an operator @@ -808,3 +837,4 @@ class AdvantageConfig: repeated_batch_fields: list[str] = field(default_factory=list) policy_logprobs_field: str = "prev_logprobs" reference_logprobs_field: str = "reference_policy_logprobs" + teacher_logprobs_field: str = "teacher_reference_logprobs" diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py index 0346d441c75..3d0c0f2ac4d 100644 --- a/nemo_rl/algorithms/single_controller_utils/setup.py +++ b/nemo_rl/algorithms/single_controller_utils/setup.py @@ -35,6 +35,7 @@ from transformers.tokenization_utils_base import PreTrainedTokenizerBase from nemo_rl.algorithms.async_utils.replay_buffer import TQReplayBuffer +from nemo_rl.algorithms import opd as opd_module from nemo_rl.algorithms.grpo import ( GRPOSaveState, _create_advantage_estimator, @@ -59,6 +60,7 @@ 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 @@ -124,33 +126,87 @@ 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 + + +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.""" cluster_config = master_config.cluster generation_config = master_config.policy["generation"] colocated = generation_config["colocated"]["enabled"] 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}." + ) if colocated: # Policy + 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=1 if backend == "megatron" else 2, 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", ( @@ -165,7 +221,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, ( @@ -173,11 +229,76 @@ 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 + ) + else: + gpus_per_instance = generation_config_dict["sglang_cfg"].get( + "gpus_per_server", 1 + ) + 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, @@ -186,6 +307,8 @@ def _build_clusters( max_colocated_worker_groups=1, 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", @@ -195,8 +318,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( @@ -551,6 +676,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: @@ -614,9 +744,23 @@ 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"] + # 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=master_config.cluster.get("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]] = {} generation = None @@ -752,6 +896,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 @@ -837,5 +1003,7 @@ 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, ) return actor_args, setup_timing_metrics diff --git a/nemo_rl/data_plane/column_io.py b/nemo_rl/data_plane/column_io.py index e8196866b50..e9fe9a7adb5 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", "token_mask", "sample_mask", diff --git a/nemo_rl/data_plane/worker_mixin.py b/nemo_rl/data_plane/worker_mixin.py index 1125245c98a..4f831776e5c 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,29 @@ 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", {}) + if cfg.get("sequence_packing", {}).get("enabled", False): + 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 + # ── split-API entrypoints (SC async path) ────────────────────────────── # # The split path lets SingleController drive forward/backward per diff --git a/nemo_rl/experience/rollout_manager.py b/nemo_rl/experience/rollout_manager.py index c983b587d32..df0f65f4a22 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 @@ -1286,6 +1289,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 isinstance(error, PostWriteEnrichmentError): + 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..9a67ab926fc 100644 --- a/nemo_rl/models/policy/teacher_worker_group.py +++ b/nemo_rl/models/policy/teacher_worker_group.py @@ -22,13 +22,18 @@ 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 from nemo_rl.distributed.batched_data_dict import ( BatchedDataDict, SequencePackingArgs, @@ -114,7 +119,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__( @@ -241,6 +246,64 @@ def __init__( if microbatch_order is not None: self.sequence_packing_args["microbatch_order"] = microbatch_order + 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) + ) + teacher_meta = replace( + meta, + task_name=f"teacher_lp:{self.alias}", + fields=["input_ids", "input_lengths"], + extra_info={ + **dict(meta.extra_info or {}), + GLOBAL_FORWARD_PAD_SEQLEN: round_up( + max(meta.sequence_lengths), sequence_pad_multiple + ), + }, + ) + sequence_packing_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"] + 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, + ) + 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, data: BatchedDataDict[GenerationDatumSpec], 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..58fe17ef55e --- /dev/null +++ b/tests/test_suites/llm/mopd-qwen3-1.7b-3n8g-megatron-pack-single-controller.sh @@ -0,0 +1,60 @@ +#!/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=True \ + 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' + + rm -rf "$CKPT_DIR" +fi diff --git a/tests/test_suites/nightly.txt b/tests/test_suites/nightly.txt index 001c95235cd..f822ff902ce 100644 --- a/tests/test_suites/nightly.txt +++ b/tests/test_suites/nightly.txt @@ -267,6 +267,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..2f754c5e2ef 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,171 @@ 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_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 + + # --------------------------------------------------------------------------- # 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..f3239973aa6 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,47 @@ 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_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..32c63c58dd4 100644 --- a/tests/unit/models/policy/test_teacher_worker_group.py +++ b/tests/unit/models/policy/test_teacher_worker_group.py @@ -12,6 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. +from unittest.mock import MagicMock + +import torch + +from nemo_rl.data_plane import KVBatchMeta +from nemo_rl.data_plane.schema import GLOBAL_FORWARD_PAD_SEQLEN +from nemo_rl.distributed.batched_data_dict import BatchedDataDict + def test_teacher_resource_config_defaults(): from nemo_rl.algorithms.opd import TeacherResourceConfig @@ -77,3 +85,84 @@ def test_create_teacher_configs_deduplicates(): } ) assert len(configs) == 2 + + +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.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.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_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 + + 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 data, micro_batch_size + 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], + ) + 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" diff --git a/tests/unit/single_controller/test_single_controller.py b/tests/unit/single_controller/test_single_controller.py index fc663aaf35c..3cb89971bf0 100644 --- a/tests/unit/single_controller/test_single_controller.py +++ b/tests/unit/single_controller/test_single_controller.py @@ -20,6 +20,7 @@ import pytest import torch +from tensordict import TensorDict import nemo_rl.algorithms.single_controller as single_controller from nemo_rl.algorithms.async_utils.staleness_sampler import BaseSampler @@ -285,6 +286,83 @@ def test_sync_weights_calibrates_and_forwards_fp8_kv_scales() -> None: ) +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: + last_metrics = {"on_policy_distillation/teacher_student_logprob_gap_mean": 0.25} + + 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 + return TensorDict( + { + "prompt_ids_for_adv": torch.zeros(2, 3, dtype=torch.long), + "total_reward": torch.zeros(2), + "token_mask": torch.ones(2, 3), + "sample_mask": torch.ones(2), + "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._opd_enabled = True + ctrl._dp_client = FakeDataPlane() + ctrl._step_log_dict = { + "rewards": [], + "masked_advantages": [], + "sequence_lengths": [], + } + ctrl._advantage_metric_values = {} + meta = KVBatchMeta( + partition_id="rollout_data", + task_name="train", + sample_ids=["a", "b"], + fields=[], + sequence_lengths=[3, 3], + ) + + enriched = asyncio.run(ctrl._advantage_stage(meta)) + + 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 + assert ctrl._advantage_metric_values == { + "on_policy_distillation/teacher_student_logprob_gap_mean": [0.25] + } + + class _EmptySampler: async def evict(self, *, current_train_weight: int) -> int: del current_train_weight @@ -401,6 +479,7 @@ def _train_pump_controller(*, sampler) -> object: ctrl._advantage_cfg = AdvantageConfig() ctrl._policy_logprobs_required = False ctrl._reference_logprobs_required = False + ctrl._opd_enabled = False ctrl._advantage_estimator = None ctrl._partition_id = "rollout_data" ctrl._sampler = sampler @@ -423,6 +502,8 @@ def _train_pump_controller(*, sampler) -> object: "masked_advantages": [], "sequence_lengths": [], } + ctrl._advantage_metric_values = {} + ctrl._teacher_coordinator = None return ctrl diff --git a/tests/unit/single_controller/test_single_controller_setup.py b/tests/unit/single_controller/test_single_controller_setup.py index 74f32e8b5c5..19c59d8a8c3 100644 --- a/tests/unit/single_controller/test_single_controller_setup.py +++ b/tests/unit/single_controller/test_single_controller_setup.py @@ -25,8 +25,10 @@ ReadyFirstSamplerConfig, SamplerConfig, ) +from nemo_rl.algorithms.advantage_estimator import AdvEstimatorConfig 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, @@ -139,6 +141,7 @@ def patched_factories(): return_value=( MagicMock(name="train_cluster"), MagicMock(name="inference_cluster"), + None, ), ) as mock_clusters, patch.object( @@ -223,6 +226,36 @@ def test_build_clusters_rejects_non_colocated_megatron_generation(): 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 + + class TestSetup: """setup arg validation + actor_args assembly.""" @@ -236,6 +269,89 @@ 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", "match"), + [ + (False, False, "opd", "requires on_policy_distillation.enabled=true"), + (True, True, "grpo", "requires grpo.adv_estimator.name='opd'"), + (True, False, "opd", "non_colocated_teachers.enabled=true"), + ], + ) + def test_invalid_mopd_config_fails_before_allocating_resources( + self, + opd_enabled: bool, + teacher_enabled: bool, + adv_name: str, + match: str, + patched_factories, + ): + mc = _make_master_config() + 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() + mc.cluster = {"num_nodes": 3, "gpus_per_node": 8} + 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, + ) + + 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"), [ 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) From 6af9670f727f27519d9f593685493999871b9558 Mon Sep 17 00:00:00 2001 From: Yi-Fu Wu Date: Sat, 22 Aug 2026 01:17:52 -0700 Subject: [PATCH 2/5] fix(mopd): harden single-controller execution Build teacher dynamic-batching metadata globally, warm the TransferQueue schema, and make teacher enrichment cleanup and cancellation safe. Pool OPD metrics from exact sufficient statistics and expand topology, routing, recipe, and lifecycle coverage. Signed-off-by: Yi-Fu Wu --- docs/about/algorithms/mopd.md | 31 ++- docs/guides/single-controller.md | 3 + ...-3n8g-megatron-pack-single-controller.yaml | 3 - examples/run_grpo_single_controller.py | 3 +- nemo_rl/algorithms/opd.py | 122 ++++++----- nemo_rl/algorithms/single_controller.py | 49 ++++- .../single_controller_utils/config.py | 17 +- .../single_controller_utils/setup.py | 28 ++- nemo_rl/data_plane/schema.py | 12 ++ nemo_rl/data_plane/worker_mixin.py | 16 +- nemo_rl/experience/rollout_manager.py | 13 +- nemo_rl/models/policy/teacher_worker_group.py | 19 ++ ...7b-3n8g-megatron-pack-single-controller.sh | 2 +- tests/unit/algorithms/test_opd.py | 192 ++++++++++++++++++ tests/unit/experience/test_rollout_manager.py | 49 +++++ .../policy/test_teacher_worker_group.py | 120 ++++++++++- .../test_single_controller.py | 41 +++- .../test_single_controller_setup.py | 151 +++++++++++++- tests/unit/test_recipes_and_test_suites.py | 9 +- 19 files changed, 781 insertions(+), 99 deletions(-) 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 703859d2c8c..1e2ec058fc6 100644 --- a/docs/guides/single-controller.md +++ b/docs/guides/single-controller.md @@ -170,6 +170,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. - Checkpointing and validation are not yet supported (setup raises if enabled). 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 index dc5559d8f34..fe7bdf3501f 100644 --- 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 @@ -4,8 +4,6 @@ defaults: ./mopd-qwen3-1.7b-3n8g-megatron-pack.yaml grpo: async_grpo: null val_period: 0 - val_at_start: false - val_at_end: false # OPD needs the student pass, but not a separate frozen reference-policy pass. skip_reference_policy_logprobs_calculation: true @@ -23,7 +21,6 @@ data_plane: checkpointing: checkpoint_dir: results/mopd-qwen3-1.7b-3n8g-megatron-pack-single-controller metric_name: null - keep_top_k: 1 logger: wandb: diff --git a/examples/run_grpo_single_controller.py b/examples/run_grpo_single_controller.py index e3479fa4a81..140a01d53f8 100644 --- a/examples/run_grpo_single_controller.py +++ b/examples/run_grpo_single_controller.py @@ -158,7 +158,8 @@ def main() -> None: except Exception as kill_error: print(f"Env {env_name!r} kill failed: {kill_error}") - for teacher_alias, teacher in (actor_args.teacher_worker_groups or {}).items(): + 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: diff --git a/nemo_rl/algorithms/opd.py b/nemo_rl/algorithms/opd.py index 7c42783ccb2..c73861a8830 100644 --- a/nemo_rl/algorithms/opd.py +++ b/nemo_rl/algorithms/opd.py @@ -254,14 +254,17 @@ def __init__( def _resolve_teacher(self, record: PromptGroupRecord) -> tuple[str, str]: extra_env_info = record.extra_env_info - if not isinstance(extra_env_info, dict): - raise ValueError( - "MOPD rollout is missing prompt-level extra_env_info with agent_ref" - ) - agent_ref = extra_env_info.get("agent_ref") + agent_ref = ( + extra_env_info.get("agent_ref") + if isinstance(extra_env_info, dict) + else None + ) if not isinstance(agent_ref, dict): raise ValueError( - "MOPD rollout extra_env_info must contain an agent_ref mapping" + "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( @@ -303,51 +306,54 @@ def _enrich_sync( remainder = actual_batch_size % dp_size padded_meta = meta temporary_sample_ids: list[str] = [] - 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=["input_ids", "input_lengths"], - 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=["input_ids", "input_lengths"], - sequence_lengths=[meta.sequence_lengths[-1]] * pad_count, - ) - 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) - - lock_started_at = time.perf_counter() 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=["input_ids", "input_lengths"], + 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=["input_ids", "input_lengths"], + 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) + + lock_started_at = time.perf_counter() with self._teacher_locks[group_alias]: inference_started_at = time.perf_counter() teacher.get_logprobs_from_meta(padded_meta) - except BaseException as inference_error: + inference_finished_at = time.perf_counter() + except BaseException as enrichment_error: if temporary_sample_ids: try: self._dp_client.clear_samples( @@ -356,12 +362,11 @@ def _enrich_sync( ) except BaseException as cleanup_error: raise BaseExceptionGroup( - f"teacher inference and temporary-row cleanup both failed " + f"teacher enrichment and temporary-row cleanup both failed " f"for group {group_alias!r}", - [inference_error, cleanup_error], + [enrichment_error, cleanup_error], ) raise - inference_finished_at = time.perf_counter() if temporary_sample_ids: self._dp_client.clear_samples( sample_ids=temporary_sample_ids, @@ -380,9 +385,22 @@ async def enrich( """Write teacher logprobs before the replay-buffer slot becomes ready.""" alias, group_alias = self._resolve_teacher(record) started_at = time.perf_counter() - lock_wait_s, inference_time_s = await asyncio.to_thread( - self._enrich_sync, meta, group_alias + # 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: + lock_wait_s, 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 teacher_model_by_agent_name = self._opd_cfg["teacher_model_by_agent_name"] diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index e443c6865fd..d946e3de2b0 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -85,6 +85,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. @@ -124,7 +140,6 @@ def __init__( self._master_config = master_config self._async_cfg = master_config.async_rl - self._opd_enabled = opd_module.is_opd_enabled(master_config) self._policy_logprobs_required = not ( master_config.loss_fn.force_on_policy_ratio and master_config.grpo.seq_logprob_error_threshold is None @@ -132,6 +147,7 @@ def __init__( self._reference_logprobs_required = not bool( master_config.grpo.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 @@ -275,6 +291,9 @@ def __init__( "sequence_lengths": [], } self._advantage_metric_values: dict[str, list[float]] = {} + self._opd_stat_sum = 0.0 + self._opd_stat_sumsq = 0.0 + self._opd_stat_count = 0 print( f"SingleControllerActor: " @@ -1081,6 +1100,16 @@ async def _train_pump(self) -> None: } ) self._advantage_metric_values.clear() + 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()) @@ -1722,7 +1751,7 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> KVBatchMeta: kwargs: dict[str, torch.Tensor] = {} if self._policy_logprobs_required: policy_logprobs = tensor_field(data, adv_cfg.policy_logprobs_field) - if self._opd_enabled: + if self._teacher_logprobs_required: kwargs["prev_logprobs"] = policy_logprobs else: kwargs["logprobs_policy"] = policy_logprobs @@ -1731,7 +1760,7 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> KVBatchMeta: data, adv_cfg.reference_logprobs_field, ) - if self._opd_enabled: + if self._teacher_logprobs_required: kwargs["teacher_logprobs"] = tensor_field( data, adv_cfg.teacher_logprobs_field, @@ -1749,9 +1778,15 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> KVBatchMeta: self._step_log_dict["masked_advantages"].append( response_advantages.detach().cpu() ) - estimator_metrics = getattr(self._advantage_estimator, "last_metrics", {}) - for name, value in estimator_metrics.items(): - self._advantage_metric_values.setdefault(name, []).append(float(value)) + 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()) + else: + estimator_metrics = getattr(self._advantage_estimator, "last_metrics", {}) + for name, value in estimator_metrics.items(): + self._advantage_metric_values.setdefault(name, []).append(float(value)) await self._call_dp( "put_samples", @@ -1779,6 +1814,6 @@ def _advantage_input_fields(self) -> list[str]: fields.append(adv_cfg.policy_logprobs_field) if self._reference_logprobs_required: fields.append(adv_cfg.reference_logprobs_field) - if self._opd_enabled: + if self._teacher_logprobs_required: fields.append(adv_cfg.teacher_logprobs_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 9ad793d7ac5..518acf9c1be 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/nemo_rl/algorithms/single_controller_utils/config.py @@ -761,6 +761,10 @@ 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 master_config.grpo.adv_estimator.name == "opd" and not opd_enabled: raise ValueError( @@ -780,6 +784,12 @@ def validate_single_controller_config(master_config: MasterConfig) -> None: "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 " @@ -794,13 +804,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" diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py index 3d0c0f2ac4d..db37c4bd536 100644 --- a/nemo_rl/algorithms/single_controller_utils/setup.py +++ b/nemo_rl/algorithms/single_controller_utils/setup.py @@ -34,8 +34,8 @@ from transformers import AutoProcessor from transformers.tokenization_utils_base import PreTrainedTokenizerBase -from nemo_rl.algorithms.async_utils.replay_buffer import TQReplayBuffer 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, _create_advantage_estimator, @@ -56,6 +56,10 @@ 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, @@ -749,6 +753,12 @@ def setup_single_controller( ) colocated = generation_config["colocated"]["enabled"] + # 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 master_config.cluster.get("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] = {} @@ -930,6 +940,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 + * grpo_config.num_generations_per_prompt + ), + consumer_tasks=["prev_lp", "ref_lp", "train"], + grpo_group_size=grpo_config.num_generations_per_prompt, + ) t0 = time.perf_counter() weight_synchronizer = create_weight_synchronizer( diff --git a/nemo_rl/data_plane/schema.py b/nemo_rl/data_plane/schema.py index 49cf79422e7..014742f1c0d 100644 --- a/nemo_rl/data_plane/schema.py +++ b/nemo_rl/data_plane/schema.py @@ -44,6 +44,18 @@ "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, 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", + "teacher_reference_logprobs", +) + # Subset fetched by logprob / ref-logprob workers. LP_SEED_FIELDS = ( "input_ids", diff --git a/nemo_rl/data_plane/worker_mixin.py b/nemo_rl/data_plane/worker_mixin.py index 4f831776e5c..2a9cb89a5ad 100644 --- a/nemo_rl/data_plane/worker_mixin.py +++ b/nemo_rl/data_plane/worker_mixin.py @@ -535,8 +535,20 @@ def get_teacher_logprobs_presharded( """Per-rank frozen-teacher logprob entrypoint for SingleController MOPD.""" data = self._fetch(meta) cfg = getattr(self, "cfg", {}) - if cfg.get("sequence_packing", {}).get("enabled", False): - data = self._attach_or_repack_pack_metadata(data, meta) + 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, diff --git a/nemo_rl/experience/rollout_manager.py b/nemo_rl/experience/rollout_manager.py index df0f65f4a22..2d2ee9b7654 100644 --- a/nemo_rl/experience/rollout_manager.py +++ b/nemo_rl/experience/rollout_manager.py @@ -61,6 +61,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.""" @@ -1292,7 +1303,7 @@ async def generate_and_push( # 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 isinstance(error, PostWriteEnrichmentError): + if _contains_post_write_enrichment_error(error): raise reason = type(error).__name__ diff --git a/nemo_rl/models/policy/teacher_worker_group.py b/nemo_rl/models/policy/teacher_worker_group.py index 9a67ab926fc..8651da53a19 100644 --- a/nemo_rl/models/policy/teacher_worker_group.py +++ b/nemo_rl/models/policy/teacher_worker_group.py @@ -36,6 +36,7 @@ from nemo_rl.data_plane.schema import GLOBAL_FORWARD_PAD_SEQLEN from nemo_rl.distributed.batched_data_dict import ( BatchedDataDict, + DynamicBatchingArgs, SequencePackingArgs, ) from nemo_rl.distributed.named_sharding import NamedSharding @@ -245,6 +246,17 @@ 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.""" @@ -275,16 +287,23 @@ def get_logprobs_from_meta(self, meta: KVBatchMeta) -> None: }, ) 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", 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 index 58fe17ef55e..55b721c5381 100755 --- 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 @@ -41,7 +41,7 @@ uv run examples/run_grpo_single_controller.py \ logger.wandb.name=$EXP_NAME \ logger.monitor_gpus=True \ logger.tensorboard_enabled=True \ - checkpointing.enabled=True \ + checkpointing.enabled=False \ checkpointing.checkpoint_dir=$CKPT_DIR \ "$@" \ 2>&1 | tee $RUN_LOG diff --git a/tests/unit/algorithms/test_opd.py b/tests/unit/algorithms/test_opd.py index 2f754c5e2ef..028597b741b 100644 --- a/tests/unit/algorithms/test_opd.py +++ b/tests/unit/algorithms/test_opd.py @@ -423,6 +423,198 @@ def fake_write_columns(dp_client, meta, fields): 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_enrichment_serializes_deduplicated_teacher(monkeypatch): """Two aliases sharing one physical teacher never overlap collectives.""" import asyncio diff --git a/tests/unit/experience/test_rollout_manager.py b/tests/unit/experience/test_rollout_manager.py index f3239973aa6..88d4ab6529f 100644 --- a/tests/unit/experience/test_rollout_manager.py +++ b/tests/unit/experience/test_rollout_manager.py @@ -195,6 +195,55 @@ async def _count_rollout(_sample): 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 32c63c58dd4..311bbbb9acb 100644 --- a/tests/unit/models/policy/test_teacher_worker_group.py +++ b/tests/unit/models/policy/test_teacher_worker_group.py @@ -14,10 +14,15 @@ 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 +from nemo_rl.data_plane.schema import ( + GLOBAL_FORWARD_PAD_SEQLEN, + MICRO_BATCH_INDICES, + MICRO_BATCH_LENGTHS, +) from nemo_rl.distributed.batched_data_dict import BatchedDataDict @@ -101,6 +106,7 @@ def get_axis_size(self, axis): 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 @@ -125,6 +131,71 @@ def get_axis_size(self, axis): 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) + 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": 1, + "max_tokens_per_microbatch": 999, + } + teacher.cfg = { + "dynamic_batching": { + "enabled": True, + "train_mb_tokens": 999, + "logprob_mb_tokens": 10, + "sequence_length_round": 1, + } + } + 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["sequence_packing_args"] is None + assert captured["dynamic_batching_args"]["max_tokens_per_microbatch"] == 10 + 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(shard.extra_info[MICRO_BATCH_LENGTHS]) for shard in shards}) == 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 @@ -134,6 +205,7 @@ class Worker(TQWorkerMixin): def __init__(self): self.written = None + self.received_batching_metadata = None def _fetch(self, meta): del meta @@ -145,7 +217,11 @@ def _fetch(self, meta): ) def get_logprobs(self, data, micro_batch_size=None): - del data, micro_batch_size + 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): @@ -157,6 +233,10 @@ def _write_back_result_field(self, meta, result, *, result_key, tq_field): sample_ids=["a"], fields=["input_ids", "input_lengths"], sequence_lengths=[3], + extra_info={ + MICRO_BATCH_INDICES: [[0]], + MICRO_BATCH_LENGTHS: [3], + }, ) worker = Worker() @@ -166,3 +246,39 @@ def _write_back_result_field(self, meta, result, *, result_key, tq_field): 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_single_controller.py b/tests/unit/single_controller/test_single_controller.py index 3cb89971bf0..f0a51baac6e 100644 --- a/tests/unit/single_controller/test_single_controller.py +++ b/tests/unit/single_controller/test_single_controller.py @@ -27,7 +27,10 @@ from nemo_rl.algorithms.grpo import GRPOConfig, _initial_grpo_save_state from nemo_rl.algorithms.loss import ClippedPGLossConfig from nemo_rl.algorithms.metric_utils import SetupTimingMetrics -from nemo_rl.algorithms.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, @@ -326,7 +329,7 @@ def put_samples(self, sample_ids, partition_id, fields): ctrl._advantage_estimator = FakeEstimator() ctrl._policy_logprobs_required = True ctrl._reference_logprobs_required = False - ctrl._opd_enabled = True + ctrl._teacher_logprobs_required = True ctrl._dp_client = FakeDataPlane() ctrl._step_log_dict = { "rewards": [], @@ -334,6 +337,9 @@ def put_samples(self, sample_ids, partition_id, fields): "sequence_lengths": [], } ctrl._advantage_metric_values = {} + 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", @@ -358,9 +364,29 @@ def put_samples(self, sample_ids, partition_id, fields): torch.full((2, 3), 0.25), ) assert "advantages" in enriched.fields - assert ctrl._advantage_metric_values == { - "on_policy_distillation/teacher_student_logprob_gap_mean": [0.25] - } + assert ctrl._advantage_metric_values == {} + assert ctrl._opd_stat_sum == pytest.approx(1.5) + assert ctrl._opd_stat_sumsq == pytest.approx(0.375) + assert ctrl._opd_stat_count == 6 + + +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: @@ -479,7 +505,7 @@ def _train_pump_controller(*, sampler) -> object: ctrl._advantage_cfg = AdvantageConfig() ctrl._policy_logprobs_required = False ctrl._reference_logprobs_required = False - ctrl._opd_enabled = False + ctrl._teacher_logprobs_required = False ctrl._advantage_estimator = None ctrl._partition_id = "rollout_data" ctrl._sampler = sampler @@ -503,6 +529,9 @@ def _train_pump_controller(*, sampler) -> object: "sequence_lengths": [], } ctrl._advantage_metric_values = {} + ctrl._opd_stat_sum = 0.0 + ctrl._opd_stat_sumsq = 0.0 + ctrl._opd_stat_count = 0 ctrl._teacher_coordinator = None return ctrl diff --git a/tests/unit/single_controller/test_single_controller_setup.py b/tests/unit/single_controller/test_single_controller_setup.py index 19c59d8a8c3..9ed18b17e52 100644 --- a/tests/unit/single_controller/test_single_controller_setup.py +++ b/tests/unit/single_controller/test_single_controller_setup.py @@ -16,9 +16,11 @@ 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.async_utils.staleness_sampler import ( @@ -35,6 +37,8 @@ SingleControllerActorArgs, setup_single_controller, ) +from nemo_rl.data_plane.schema import SC_ROLLOUT_SCHEMA_FIELDS +from nemo_rl.utils.config import load_config, register_omegaconf_resolvers def _make_master_config( @@ -98,6 +102,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( @@ -256,6 +261,73 @@ def __init__(self, **kwargs): 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.""" @@ -270,11 +342,30 @@ def test_multiple_dataloader_not_supported(self): setup_single_controller(mc, MagicMock(pad_token_id=0)) @pytest.mark.parametrize( - ("opd_enabled", "teacher_enabled", "adv_name", "match"), + ( + "opd_enabled", + "teacher_enabled", + "adv_name", + "use_nemo_gym", + "match", + ), [ - (False, False, "opd", "requires on_policy_distillation.enabled=true"), - (True, True, "grpo", "requires grpo.adv_estimator.name='opd'"), - (True, False, "opd", "non_colocated_teachers.enabled=true"), + ( + 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( @@ -282,10 +373,12 @@ def test_invalid_mopd_config_fails_before_allocating_resources( 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, @@ -301,8 +394,15 @@ def test_invalid_mopd_config_fails_before_allocating_resources( def test_mopd_reserves_before_models_and_initializes_teacher_last( self, patched_factories, monkeypatch ): - mc = _make_master_config() + 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, @@ -342,6 +442,12 @@ def create_teachers(*args, **kwargs): "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)) @@ -455,6 +561,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_router_replay_requires_routes_in_tq_buffer(self, patched_factories): mc = _make_master_config(colocated=True) diff --git a/tests/unit/test_recipes_and_test_suites.py b/tests/unit/test_recipes_and_test_suites.py index dd80daf6e30..090c50d2359 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_4048_hours(nightly_test_suite, tracker): +def test_nightly_compute_stays_below_4056_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,10 +288,11 @@ def test_nightly_compute_stays_below_4048_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()) - # Dynamo adds 96 GPU-hours and the two Async PPO nightlies add 24 to the + # Dynamo adds 96 GPU-hours, the two Async PPO nightlies add 24, and the + # budget reserves 8 for the SingleController MOPD nightly on top of the # former 3928-hour limit. - assert total_gpu_hours <= 4048, ( - f"Total GPU hours exceeded 4048: {last_line}. We should revisit the test suites to reduce the total GPU hours." + assert total_gpu_hours <= 4056, ( + f"Total GPU hours exceeded 4056: {last_line}. We should revisit the test suites to reduce the total GPU hours." ) tracker.track("total_nightly_gpu_hours", total_gpu_hours) From 7e7d72130a112fe336978c03ae38e28773feb546 Mon Sep 17 00:00:00 2001 From: Yi-Fu Wu Date: Sat, 22 Aug 2026 01:54:58 -0700 Subject: [PATCH 3/5] fix(lint): sort MOPD imports Signed-off-by: Yi-Fu Wu --- nemo_rl/algorithms/single_controller_utils/config.py | 2 +- tests/unit/single_controller/test_single_controller_setup.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/nemo_rl/algorithms/single_controller_utils/config.py b/nemo_rl/algorithms/single_controller_utils/config.py index 84e6baef5fc..17c6ebd6ae2 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, @@ -36,7 +37,6 @@ ) from nemo_rl.algorithms.grpo import GRPOConfig, GRPOLoggerConfig from nemo_rl.algorithms.loss import ClippedPGLossConfig -from nemo_rl.algorithms import opd as opd_module from nemo_rl.algorithms.opd import OnPolicyDistillationConfig from nemo_rl.data import DataConfig from nemo_rl.data_plane.interfaces import DataPlaneConfig diff --git a/tests/unit/single_controller/test_single_controller_setup.py b/tests/unit/single_controller/test_single_controller_setup.py index 9ed18b17e52..71c94a638ff 100644 --- a/tests/unit/single_controller/test_single_controller_setup.py +++ b/tests/unit/single_controller/test_single_controller_setup.py @@ -23,11 +23,11 @@ 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.advantage_estimator import AdvEstimatorConfig from nemo_rl.algorithms.grpo import GRPOConfig from nemo_rl.algorithms.loss import ClippedPGLossConfig from nemo_rl.algorithms.opd import OnPolicyDistillationConfig From b0b91a30293cd8afced9ccc373ad5c4af54b0d98 Mon Sep 17 00:00:00 2001 From: Yi-Fu Wu Date: Sat, 22 Aug 2026 17:41:53 -0700 Subject: [PATCH 4/5] fix(mopd): address single-controller review findings Signed-off-by: Yi-Fu Wu --- nemo_rl/algorithms/opd.py | 57 ++++--- nemo_rl/algorithms/single_controller.py | 13 -- .../single_controller_utils/setup.py | 7 +- nemo_rl/data_plane/schema.py | 3 + nemo_rl/models/policy/teacher_worker_group.py | 17 ++- tests/unit/algorithms/test_opd.py | 57 +++++++ .../policy/test_teacher_worker_group.py | 142 +++++++++++++++++- .../test_single_controller.py | 15 +- .../test_single_controller_setup.py | 28 ++++ 9 files changed, 277 insertions(+), 62 deletions(-) diff --git a/nemo_rl/algorithms/opd.py b/nemo_rl/algorithms/opd.py index c73861a8830..cbe21feb598 100644 --- a/nemo_rl/algorithms/opd.py +++ b/nemo_rl/algorithms/opd.py @@ -22,7 +22,6 @@ from __future__ import annotations import asyncio -import threading import time import uuid from typing import Any, Optional @@ -33,6 +32,7 @@ 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, @@ -242,7 +242,7 @@ def __init__( # Physical (deduplicated) groups own locks, not routing aliases. Two # aliases sharing one checkpoint therefore share one collective FIFO. self._teacher_locks = { - group_alias: threading.Lock() for group_alias in self._teacher_worker_groups + group_alias: asyncio.Lock() for group_alias in self._teacher_worker_groups } self._teacher_batches = 0 self._teacher_samples = 0 @@ -291,7 +291,7 @@ def _enrich_sync( self, meta: KVBatchMeta, group_alias: str, - ) -> tuple[float, float]: + ) -> 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") @@ -313,7 +313,7 @@ def _enrich_sync( source_data = read_columns( self._dp_client, source_meta, - select_fields=["input_ids", "input_lengths"], + select_fields=TEACHER_LP_FIELDS, pad_value_dict={"input_ids": 0}, ) input_ids = source_data["input_ids"] @@ -331,7 +331,7 @@ def _enrich_sync( partition_id=meta.partition_id, task_name=meta.task_name, sample_ids=temporary_sample_ids, - fields=["input_ids", "input_lengths"], + fields=list(TEACHER_LP_FIELDS), sequence_lengths=[meta.sequence_lengths[-1]] * pad_count, ) # Keep this write inside the cleanup lifetime. A backend may @@ -348,10 +348,8 @@ def _enrich_sync( ) padded_meta = meta.concat(pad_meta) - lock_started_at = time.perf_counter() - with self._teacher_locks[group_alias]: - inference_started_at = time.perf_counter() - teacher.get_logprobs_from_meta(padded_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: @@ -372,10 +370,7 @@ def _enrich_sync( sample_ids=temporary_sample_ids, partition_id=meta.partition_id, ) - return ( - inference_started_at - lock_started_at, - inference_finished_at - inference_started_at, - ) + return inference_finished_at - inference_started_at async def enrich( self, @@ -385,22 +380,27 @@ async def enrich( """Write teacher logprobs before the replay-buffer slot becomes ready.""" alias, group_alias = self._resolve_teacher(record) started_at = time.perf_counter() - # 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: - lock_wait_s, inference_time_s = await asyncio.shield(enrichment_task) - except asyncio.CancelledError: + 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: - 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 + 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 teacher_model_by_agent_name = self._opd_cfg["teacher_model_by_agent_name"] @@ -411,7 +411,6 @@ async def enrich( self._teacher_lock_wait_time_s += lock_wait_s self._aliases_seen.add(alias) self._models_seen.add(teacher_model_by_agent_name[alias]) - record.rollout_metrics["teacher_logprob_time"] = total_time_s print( f"[teacher_logprob] group={group_alias} samples={meta.size} " f"lock_wait={lock_wait_s:.2f}s inference={inference_time_s:.2f}s " diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index b33467e3e39..7604b873dd4 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -295,7 +295,6 @@ def __init__( "sequence_lengths": [], "seq_logprob_error_metrics": [], } - self._advantage_metric_values: dict[str, list[float]] = {} self._opd_stat_sum = 0.0 self._opd_stat_sumsq = 0.0 self._opd_stat_count = 0 @@ -1111,14 +1110,6 @@ 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( - { - name: float(sum(values) / len(values)) - for name, values in self._advantage_metric_values.items() - if values - } - ) - self._advantage_metric_values.clear() step_metrics.update( _pooled_opd_metrics( self._opd_stat_sum, @@ -1857,10 +1848,6 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> tuple[KVBatchMeta, bool]: self._opd_stat_sum += float(valid.sum()) self._opd_stat_sumsq += float((valid * valid).sum()) self._opd_stat_count += int(valid.numel()) - else: - estimator_metrics = getattr(self._advantage_estimator, "last_metrics", {}) - for name, value in estimator_metrics.items(): - self._advantage_metric_values.setdefault(name, []).append(float(value)) fields_to_put = {adv_cfg.output_field: advantages} if seq_logprob_error_threshold is not None: diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py index db37c4bd536..c588d596024 100644 --- a/nemo_rl/algorithms/single_controller_utils/setup.py +++ b/nemo_rl/algorithms/single_controller_utils/setup.py @@ -272,10 +272,15 @@ def _build_clusters( gpus_per_instance = vllm_cfg["tensor_parallel_size"] * vllm_cfg.get( "pipeline_parallel_size", 1 ) - else: + 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 diff --git a/nemo_rl/data_plane/schema.py b/nemo_rl/data_plane/schema.py index 014742f1c0d..51bf0c02e98 100644 --- a/nemo_rl/data_plane/schema.py +++ b/nemo_rl/data_plane/schema.py @@ -64,6 +64,9 @@ "sample_mask", ) +# Text-only inputs fetched by frozen MOPD teachers for logprob inference. +TEACHER_LP_FIELDS = (INPUT_IDS, INPUT_LENGTHS) + # Fields requested for KV-scale calibration. Positive include-list: # calibration only handles seq-dim tensor inputs, so we name them # explicitly. Train-side deltas (logprobs/advantages/masks) and diff --git a/nemo_rl/models/policy/teacher_worker_group.py b/nemo_rl/models/policy/teacher_worker_group.py index 8651da53a19..69de3d5e85c 100644 --- a/nemo_rl/models/policy/teacher_worker_group.py +++ b/nemo_rl/models/policy/teacher_worker_group.py @@ -33,7 +33,7 @@ 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 +from nemo_rl.data_plane.schema import GLOBAL_FORWARD_PAD_SEQLEN, TEACHER_LP_FIELDS from nemo_rl.distributed.batched_data_dict import ( BatchedDataDict, DynamicBatchingArgs, @@ -169,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: @@ -275,10 +280,18 @@ def get_logprobs_from_meta(self, meta: KVBatchMeta) -> None: 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=["input_ids", "input_lengths"], + fields=list(TEACHER_LP_FIELDS), extra_info={ **dict(meta.extra_info or {}), GLOBAL_FORWARD_PAD_SEQLEN: round_up( diff --git a/tests/unit/algorithms/test_opd.py b/tests/unit/algorithms/test_opd.py index 028597b741b..c5313548952 100644 --- a/tests/unit/algorithms/test_opd.py +++ b/tests/unit/algorithms/test_opd.py @@ -615,6 +615,63 @@ async def 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 diff --git a/tests/unit/models/policy/test_teacher_worker_group.py b/tests/unit/models/policy/test_teacher_worker_group.py index 311bbbb9acb..8ab4ab9eb5f 100644 --- a/tests/unit/models/policy/test_teacher_worker_group.py +++ b/tests/unit/models/policy/test_teacher_worker_group.py @@ -22,6 +22,7 @@ GLOBAL_FORWARD_PAD_SEQLEN, MICRO_BATCH_INDICES, MICRO_BATCH_LENGTHS, + TEACHER_LP_FIELDS, ) from nemo_rl.distributed.batched_data_dict import BatchedDataDict @@ -92,6 +93,63 @@ 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 @@ -125,6 +183,7 @@ def get_axis_size(self, axis): 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"] ) @@ -146,6 +205,7 @@ def get_axis_size(self, axis): 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) @@ -158,15 +218,85 @@ def capture_plan(meta, **kwargs): teacher.dynamic_batching_args = { "input_key": "input_ids", "input_lengths_key": "input_lengths", - "sequence_length_round": 1, + "sequence_length_round": 64, "max_tokens_per_microbatch": 999, } teacher.cfg = { "dynamic_batching": { "enabled": True, "train_mb_tokens": 999, - "logprob_mb_tokens": 10, - "sequence_length_round": 1, + "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(shard.extra_info[MICRO_BATCH_LENGTHS]) for shard in shards}) == 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() @@ -182,8 +312,9 @@ def capture_plan(meta, **kwargs): teacher.get_logprobs_from_meta(meta) - assert captured["sequence_packing_args"] is None - assert captured["dynamic_batching_args"]["max_tokens_per_microbatch"] == 10 + 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", @@ -193,7 +324,6 @@ def capture_plan(meta, **kwargs): ] 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(shard.extra_info[MICRO_BATCH_LENGTHS]) for shard in shards}) == 1 def test_teacher_worker_presharded_entrypoint_writes_teacher_tq_field(): diff --git a/tests/unit/single_controller/test_single_controller.py b/tests/unit/single_controller/test_single_controller.py index fed5db93544..2adbcd60a11 100644 --- a/tests/unit/single_controller/test_single_controller.py +++ b/tests/unit/single_controller/test_single_controller.py @@ -348,7 +348,6 @@ def test_advantage_stage_applies_seq_logprob_error_mask_before_streaming_train( ctrl._policy_logprobs_required = True ctrl._reference_logprobs_required = False ctrl._teacher_logprobs_required = False - ctrl._advantage_metric_values = {} ctrl._master_config = SimpleNamespace( grpo=SimpleNamespace(seq_logprob_error_threshold=2.0) ) @@ -416,7 +415,6 @@ def test_advantage_stage_reports_seq_logprob_metrics_without_masking() -> None: ctrl._policy_logprobs_required = True ctrl._reference_logprobs_required = False ctrl._teacher_logprobs_required = False - ctrl._advantage_metric_values = {} ctrl._master_config = SimpleNamespace( grpo=SimpleNamespace(seq_logprob_error_threshold=None) ) @@ -478,7 +476,6 @@ def test_advantage_stage_skips_estimator_when_seq_mask_removes_whole_chunk( ctrl._policy_logprobs_required = True ctrl._reference_logprobs_required = False ctrl._teacher_logprobs_required = False - ctrl._advantage_metric_values = {} ctrl._master_config = SimpleNamespace( grpo=SimpleNamespace(seq_logprob_error_threshold=2.0) ) @@ -533,7 +530,6 @@ def test_advantage_stage_skips_preexisting_empty_mask_without_seq_threshold() -> ctrl._policy_logprobs_required = False ctrl._reference_logprobs_required = False ctrl._teacher_logprobs_required = False - ctrl._advantage_metric_values = {} ctrl._master_config = SimpleNamespace( grpo=SimpleNamespace(seq_logprob_error_threshold=None) ) @@ -589,7 +585,7 @@ def get_samples(self, sample_ids, partition_id, select_fields): { "prompt_ids_for_adv": torch.zeros(2, 3, dtype=torch.long), "total_reward": torch.zeros(2), - "token_mask": torch.ones(2, 3), + "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), @@ -617,7 +613,6 @@ def put_samples(self, sample_ids, partition_id, fields): "sequence_lengths": [], "seq_logprob_error_metrics": [], } - ctrl._advantage_metric_values = {} ctrl._opd_stat_sum = 0.0 ctrl._opd_stat_sumsq = 0.0 ctrl._opd_stat_count = 0 @@ -646,10 +641,9 @@ def put_samples(self, sample_ids, partition_id, fields): torch.full((2, 3), 0.25), ) assert "advantages" in (enriched.fields or []) - assert ctrl._advantage_metric_values == {} - assert ctrl._opd_stat_sum == pytest.approx(1.5) - assert ctrl._opd_stat_sumsq == pytest.approx(0.375) - assert ctrl._opd_stat_count == 6 + 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: @@ -822,7 +816,6 @@ def _train_pump_controller(*, sampler) -> object: "sequence_lengths": [], "seq_logprob_error_metrics": [], } - ctrl._advantage_metric_values = {} ctrl._opd_stat_sum = 0.0 ctrl._opd_stat_sumsq = 0.0 ctrl._opd_stat_count = 0 diff --git a/tests/unit/single_controller/test_single_controller_setup.py b/tests/unit/single_controller/test_single_controller_setup.py index 71c94a638ff..2fb84f9d300 100644 --- a/tests/unit/single_controller/test_single_controller_setup.py +++ b/tests/unit/single_controller/test_single_controller_setup.py @@ -231,6 +231,34 @@ 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) From 112afd1e55e0cadaa25aa8f185b09398aef7eed6 Mon Sep 17 00:00:00 2001 From: Yi-Fu Wu Date: Tue, 25 Aug 2026 22:57:24 -0700 Subject: [PATCH 5/5] fix(mopd): address follow-up review findings Strengthen the SingleController MOPD nightly and batching/concurrency coverage. Reuse routing metric computation and omit cardinality values for idle enrichment intervals. Signed-off-by: Yi-Fu Wu --- nemo_rl/algorithms/opd.py | 20 +++--- ...7b-3n8g-megatron-pack-single-controller.sh | 3 +- tests/unit/algorithms/test_opd.py | 68 +++++++++++++++++++ .../policy/test_teacher_worker_group.py | 21 +++++- 4 files changed, 99 insertions(+), 13 deletions(-) diff --git a/nemo_rl/algorithms/opd.py b/nemo_rl/algorithms/opd.py index cbe21feb598..d99a64b659e 100644 --- a/nemo_rl/algorithms/opd.py +++ b/nemo_rl/algorithms/opd.py @@ -250,7 +250,6 @@ def __init__( self._teacher_inference_time_s = 0.0 self._teacher_lock_wait_time_s = 0.0 self._aliases_seen: set[str] = set() - self._models_seen: set[str] = set() def _resolve_teacher(self, record: PromptGroupRecord) -> tuple[str, str]: extra_env_info = record.extra_env_info @@ -403,14 +402,12 @@ async def enrich( raise total_time_s = time.perf_counter() - started_at - teacher_model_by_agent_name = self._opd_cfg["teacher_model_by_agent_name"] 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) - self._models_seen.add(teacher_model_by_agent_name[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 " @@ -421,27 +418,28 @@ async def enrich( def drain_metrics(self) -> dict[str, float]: """Return and reset teacher activity accumulated since the last drain.""" - alias_unique = len(self._aliases_seen) - model_unique = len(self._models_seen) 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, - "on_policy_distillation/teacher_alias_unique": float(alias_unique), - "on_policy_distillation/teacher_model_unique": float(model_unique), - "on_policy_distillation/teacher_alias_to_model_compression": float( - model_unique / max(alias_unique, 1) - ), } + 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() - self._models_seen.clear() return metrics 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 index 55b721c5381..e33726d8652 100755 --- 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 @@ -54,7 +54,8 @@ if [[ $(jq 'to_entries | .[] | select(.key == "train/token_mult_prob_error") | . '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/teacher_model_unique"]) == 1' \ + 'max(data["train/on_policy_distillation/adv_std"]) < 0.01' rm -rf "$CKPT_DIR" fi diff --git a/tests/unit/algorithms/test_opd.py b/tests/unit/algorithms/test_opd.py index c5313548952..258bd319888 100644 --- a/tests/unit/algorithms/test_opd.py +++ b/tests/unit/algorithms/test_opd.py @@ -735,6 +735,74 @@ async def 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/models/policy/test_teacher_worker_group.py b/tests/unit/models/policy/test_teacher_worker_group.py index 8ab4ab9eb5f..2fb51686503 100644 --- a/tests/unit/models/policy/test_teacher_worker_group.py +++ b/tests/unit/models/policy/test_teacher_worker_group.py @@ -252,7 +252,16 @@ def capture_plan(meta, **kwargs): ] 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(shard.extra_info[MICRO_BATCH_LENGTHS]) for shard in shards}) == 1 + 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 @@ -324,6 +333,16 @@ def capture_plan(meta, **kwargs): ] 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():