diff --git a/docs/about/algorithms/index.md b/docs/about/algorithms/index.md index 21d4155b00..238aa61705 100644 --- a/docs/about/algorithms/index.md +++ b/docs/about/algorithms/index.md @@ -10,6 +10,7 @@ NeMo RL supports multiple training algorithms for post-training large language m | [DAPO](dapo.md) | [DAPO Single Node](dapo.md#dapo-single-node) | [DAPO Multi-node](dapo.md#dapo-multi-node) | | [CISPO](cispo.md) | [CISPO Configuration](cispo.md#configuration) | [CISPO Async Lag-1 Recipe](cispo.md#async-lag-1-recipe) | | [On-policy Distillation](on-policy-distillation.md) | [Distillation Single Node](on-policy-distillation.md#on-policy-distillation-single-node) | [Distillation Multi-node](on-policy-distillation.md#on-policy-distillation-multi-node) | +| [Multi-Teacher On-Policy Distillation (MOPD)](mopd.md) | — | [MOPD Configuration](mopd.md#configuration) | | [Supervised Fine-Tuning (SFT)](sft.md) | [SFT Single Node](sft.md#sft-single-node) | [SFT Multi-node](sft.md#sft-multi-node) | | [DPO](dpo.md) | [DPO Single Node](dpo.md#dpo-single-node) | [DPO Multi-node](dpo.md#dpo-multi-node) | | [PPO](ppo.md) | [PPO Single Node](ppo.md#ppo-single-node) | [PPO Multi-node](ppo.md#ppo-multi-node) | @@ -25,6 +26,7 @@ dapo cispo ppo on-policy-distillation +mopd sft dpo rm diff --git a/docs/about/algorithms/mopd.md b/docs/about/algorithms/mopd.md new file mode 100644 index 0000000000..dd803a34ce --- /dev/null +++ b/docs/about/algorithms/mopd.md @@ -0,0 +1,148 @@ +# Multi-Teacher On-Policy Distillation (MOPD) + +Multi-Teacher On-Policy Distillation (MOPD) distills one or more teacher models +into the policy by replacing GRPO's reward-based advantage with a token-level +distillation advantage ([MiMo-V2-Flash Technical Report](https://arxiv.org/abs/2601.02780)). +MOPD runs on async GRPO and collects rollouts through NeMo Gym, so the agent +loop drives multi-turn / multi-step interaction. Each token of the resulting +student rollout is scored by a teacher, and the policy is updated to close the +gap with the teacher. + +Unlike the teacher-logit knowledge distillation in +[On-policy Distillation](on-policy-distillation.md) (`run_distillation.py`), MOPD +runs on top of the GRPO trainer: it is selected with `adv_estimator: opd` and +serves teachers from dedicated, non-colocated worker groups during async +collection. + +## Advantage + +For each token `t`, the distillation advantage is the stop-gradient +teacher-minus-student log-probability gap: + +``` +Â_t = sg[ log π_teacher(t) − log π_student(t) ] +``` + +`log π_student` is the policy's `prev_logprobs` and `log π_teacher` is computed +by the teacher worker group at collection time. Maximizing this advantage is +reverse-KL minimization — it pushes the student toward the teacher's token +distribution — but, unlike forward-KL logit distillation, it needs only the +teacher's log-probability for the *sampled* token rather than the full +vocabulary distribution. + +The advantage is applied only to trained (assistant) tokens via the loss mask; +tool / environment tokens contribute zero. Because the advantage subtracts a +real `prev_logprobs`, MOPD requires the student log-probabilities to actually be +computed — see [Configuration](#configuration). + +## Configuration + +Enable MOPD in two places: select the advantage estimator and add the +`on_policy_distillation` block. + +```yaml +grpo: + # MOPD runs on async GRPO with NeMo Gym rollouts. + async_grpo: + enabled: true + adv_estimator: + name: opd + # OPD subtracts a real prev_logprobs, so it must not be skipped. + seq_logprob_error_threshold: 2.0 + +loss_fn: + # REINFORCE form (drop the PPO probability-ratio clipping); on-policy + # correction is handled by the ICE-POP gate below instead. + disable_ppo_ratio: true + # ICE-POP hard gate: zero tokens whose train/inference importance-sampling + # weight falls outside bounds, correcting async off-policy drift. + use_importance_sampling_correction: true + truncated_importance_sampling_type: icepop + # Teacher distillation is the entire learning signal — no reference-policy KL. + reference_policy_kl_penalty: 0.0 + +on_policy_distillation: + enabled: true + # Map each NeMo Gym agent name to a teacher checkpoint. + teacher_model_by_agent_name: + default_teacher: Qwen/Qwen3-1.7B + # Agents not present in the map fall back to this alias (must be a mapped key). + default_teacher_alias: default_teacher + # If true, an unmapped agent raises instead of falling back. + strict_agent_name_match: false + # Aliases that share one checkpoint reuse a single teacher worker group. + deduplicate_shared_teacher_checkpoints: true + non_colocated_teachers: + enabled: true + # Resourcing for each teacher worker group. + default_teacher_cfg: + tensor_model_parallel_size: 2 + pipeline_model_parallel_size: 1 + context_parallel_size: 1 + num_nodes: 1 + gpus_per_node: 8 + precision: bf16 + micro_batch_size: 1 + # Optional per-alias overrides on top of default_teacher_cfg. + teacher_overrides: {} +``` + +> [!NOTE] +> Teachers run the Megatron backend in inference-only mode. A DTensor-configured +> policy is rejected for the teacher; PEFT / draft modules are stripped so +> adapters are never attached to the frozen teacher; and teachers run +> unquantized (a policy `quant_cfg` is ignored, with a warning). + +> [!NOTE] +> `adv_estimator: opd` fails fast at setup if the config would zero +> `prev_logprobs` (`loss_fn.force_on_policy_ratio: true` with no +> `grpo.seq_logprob_error_threshold`), because the advantage would silently +> degrade to `teacher_logprobs − 0`. + +### Teacher routing + +Each rollout sample carries its NeMo Gym `agent_ref`. At collection time the +agent name is resolved to a teacher alias (`teacher_model_by_agent_name`, falling +back to `default_teacher_alias`), samples are grouped by teacher, and each group +is scored by exactly one teacher — there is no ensemble averaging across +teachers. When several aliases map to the same checkpoint, +`deduplicate_shared_teacher_checkpoints` collapses them onto a single worker +group so they share GPUs. + +### Resourcing + +Non-colocated teachers each get their own Ray cluster on dedicated GPUs (they +are queried every rollout group, so time-sharing with the policy/generation +would serialize and destroy the async overlap). Their nodes are reserved from +the policy's budget: with `total_nodes` total, the teacher groups take +`sum(num_nodes)` and the policy uses the remainder (setup fails if nothing is +left for the policy). Deduplicated teachers share one group's nodes. + +For example, the reference 3-node recipe lays out: 1 node policy (student, +trainable) + 1 node vLLM generation (frozen) + 1 node teacher (frozen). Ten +distinct teachers at 1 node each would instead add 10 nodes on top of the +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: + +```sh +uv run examples/nemo_gym/run_grpo_nemo_gym.py \ + --config examples/configs/recipes/llm/mopd-qwen3-1.7b-3n8g-megatron-pack.yaml \ + data.train.data_path=/path/to/train.jsonl \ + data.val.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. + +## References + +- LLM-Core Xiaomi, *MiMo-V2-Flash Technical Report*, which introduces the + multi-teacher on-policy distillation paradigm: + [arxiv.org/abs/2601.02780](https://arxiv.org/abs/2601.02780) diff --git a/docs/index.md b/docs/index.md index 0a701ce7c6..f83a5fb57b 100644 --- a/docs/index.md +++ b/docs/index.md @@ -62,7 +62,7 @@ Learn about DTensor and Megatron Core training backends, their capabilities, and :link: about/algorithms/index :link-type: doc -Discover supported algorithms including GRPO, PPO, SFT, DPO, RM, and on-policy distillation with detailed guides and examples. +Discover supported algorithms including GRPO, PPO, SFT, DPO, RM, on-policy distillation, and multi-teacher on-policy distillation (MOPD) with detailed guides and examples. ::: :::{grid-item-card} {octicon}`graph` Evaluation diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml index ee2cca7422..53c77b8683 100644 --- a/examples/configs/grpo_math_1B.yaml +++ b/examples/configs/grpo_math_1B.yaml @@ -458,3 +458,11 @@ data_plane: local_buffer_size: 68719476736 # 64 GiB — used when backend == "mooncake_cpu" # observability: # NotRequired # enabled: false + +# Multi-Teacher On-Policy Distillation (MOPD): distills from one or more teacher +# models into the policy via token-level teacher-minus-student logprob advantages, +# served by non-colocated teacher worker groups (OPD advantage estimator + +# nemo_gym). null = disabled (default). See +# examples/configs/recipes/llm/mopd-qwen3-1.7b-3n8g-megatron-pack.yaml for a full +# enabled example. +on_policy_distillation: null diff --git a/examples/configs/recipes/llm/mopd-qwen3-1.7b-3n8g-megatron-pack.yaml b/examples/configs/recipes/llm/mopd-qwen3-1.7b-3n8g-megatron-pack.yaml new file mode 100644 index 0000000000..d967ea8651 --- /dev/null +++ b/examples/configs/recipes/llm/mopd-qwen3-1.7b-3n8g-megatron-pack.yaml @@ -0,0 +1,144 @@ +defaults: ../../grpo_math_1B.yaml +grpo: + num_prompts_per_step: 8 + num_generations_per_prompt: 4 + num_val_generations_per_prompt: 1 + max_num_steps: 5 + val_period: 1000 + overlong_filtering: true + max_val_samples: null + val_batch_size: 32 + async_grpo: + enabled: true + adv_estimator: + name: opd + seq_logprob_error_threshold: 2.0 +loss_fn: + reference_policy_kl_penalty: 0.0 + kl_input_clamp_value: null + kl_output_clamp_value: null + ratio_clip_max: 0.28 + use_on_policy_kl_approximation: true + disable_ppo_ratio: true + use_importance_sampling_correction: true + truncated_importance_sampling_ratio: 5.0 + truncated_importance_sampling_ratio_min: 0.2 + truncated_importance_sampling_type: icepop +checkpointing: + enabled: false + checkpoint_dir: results/mopd-qwen3-1.7b-3n8g-megatron-pack + metric_name: val:total_reward/mean + keep_top_k: 1 + save_period: 1000 + checkpoint_must_save_by: 00:03:40:00 + save_optimizer: false +policy: + model_name: Qwen/Qwen3-1.7B + train_global_batch_size: 32 + train_micro_batch_size: 1 + generation_batch_size: 64 + logprob_batch_size: 1 + max_total_sequence_length: 32768 + logprob_chunk_size: 2048 + dtensor_cfg: + enabled: false + megatron_cfg: + enabled: true + activation_checkpointing: true + bias_activation_fusion: false + tensor_model_parallel_size: 2 + sequence_parallel: true + defer_fp32_logits: true + optimizer: + lr: 3.0e-06 + min_lr: 3.0e-06 + weight_decay: 0.0 + scheduler: + lr_decay_iters: null + lr_warmup_iters: 10 + lr_warmup_init: 3.0e-07 + distributed_data_parallel_config: + average_in_collective: false + make_sequence_length_divisible_by: 8 + optimizer: null + scheduler: null + generation: + max_new_tokens: 2048 + vllm_cfg: + async_engine: true + tensor_parallel_size: 2 + gpu_memory_utilization: 0.5 + expose_http_server: true + http_server_serving_chat_kwargs: + enable_auto_tools: true + tool_parser: hermes + reasoning_parser: qwen3 + colocated: + enabled: false + resources: + gpus_per_node: 8 + num_nodes: 1 +data: + max_input_seq_length: null + train: + data_path: "${oc.env:HF_HOME}/nanov3_data/train-split.jsonl" + dataset_name: NemoGymDataset + validation: + data_path: "${oc.env:HF_HOME}/nanov3_data/val-split.jsonl" + dataset_name: NemoGymDataset + default: + dataset_name: NemoGymDataset + env_name: nemo_gym + prompt_file: null + processor: nemo_gym_data_processor +env: + should_use_nemo_gym: true + nemo_gym: + config_paths: + - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml + - resources_servers/math_with_judge/configs/math_with_judge.yaml + - resources_servers/code_gen/configs/code_gen.yaml + - resources_servers/workplace_assistant/configs/workplace_assistant.yaml + - resources_servers/mcqa/configs/mcqa.yaml + - resources_servers/instruction_following/configs/instruction_following.yaml + - resources_servers/structured_outputs/configs/structured_outputs_json.yaml + math_with_judge: + resources_servers: + math_with_judge: + judge_model_server: + name: policy_model + should_use_judge: false + code_gen: + resources_servers: + code_gen: + num_processes: 1024 + unit_test_timeout_secs: 10 + debug: false +logger: + tensorboard_enabled: true + wandb: + project: mopd + name: mopd-qwen3-1.7b-3n8g-megatron-pack + mlflow: + experiment_name: mopd + run_name: mopd-qwen3-1.7b-3n8g-megatron-pack +cluster: + gpus_per_node: 8 + num_nodes: 3 +on_policy_distillation: + enabled: true + teacher_model_by_agent_name: + default_teacher: Qwen/Qwen3-1.7B + default_teacher_alias: default_teacher + strict_agent_name_match: false + deduplicate_shared_teacher_checkpoints: true + non_colocated_teachers: + enabled: true + default_teacher_cfg: + tensor_model_parallel_size: 2 + pipeline_model_parallel_size: 1 + context_parallel_size: 1 + num_nodes: 1 + gpus_per_node: 8 + precision: bf16 + micro_batch_size: 1 diff --git a/examples/nemo_gym/run_grpo_nemo_gym.py b/examples/nemo_gym/run_grpo_nemo_gym.py index 2853857d07..25a2c18493 100644 --- a/examples/nemo_gym/run_grpo_nemo_gym.py +++ b/examples/nemo_gym/run_grpo_nemo_gym.py @@ -208,6 +208,8 @@ def main() -> None: checkpointer, grpo_state, master_config, + teacher_worker_groups, + alias_to_group_alias, ) = setup(config, tokenizer, train_dataset, val_dataset) # NeMo-Gym is spun up inside setup() (overlapped with vLLM model load). @@ -276,6 +278,8 @@ def main() -> None: grpo_save_state=grpo_state, master_config=master_config, max_trajectory_age_steps=async_config["max_trajectory_age_steps"], + teacher_worker_groups=teacher_worker_groups, + alias_to_group_alias=alias_to_group_alias, ) else: print("🚀 Running synchronous GRPO training") diff --git a/examples/run_grpo.py b/examples/run_grpo.py index d7118188f9..2c2f31fda6 100644 --- a/examples/run_grpo.py +++ b/examples/run_grpo.py @@ -142,6 +142,8 @@ def _make_policy(**kwargs): checkpointer, grpo_state, master_config, + teacher_worker_groups, + alias_to_group_alias, ) = setup( config, tokenizer, @@ -200,6 +202,8 @@ def _make_policy(**kwargs): grpo_save_state=grpo_state, master_config=master_config, max_trajectory_age_steps=async_config["max_trajectory_age_steps"], + teacher_worker_groups=teacher_worker_groups, + alias_to_group_alias=alias_to_group_alias, ) else: # Two parallel synchronous trainers (verl-style — main_ppo.py vs diff --git a/examples/run_grpo_sliding_puzzle.py b/examples/run_grpo_sliding_puzzle.py index 6cd2ee568f..86995939ae 100644 --- a/examples/run_grpo_sliding_puzzle.py +++ b/examples/run_grpo_sliding_puzzle.py @@ -262,6 +262,8 @@ def main(): checkpointer, grpo_state, master_config, + _teacher_worker_groups, + _alias_to_group_alias, ) = setup(config, tokenizer, dataset, val_dataset) grpo_train( diff --git a/examples/run_vlm_grpo.py b/examples/run_vlm_grpo.py index e9a7cf1f8f..8f27549f00 100644 --- a/examples/run_vlm_grpo.py +++ b/examples/run_vlm_grpo.py @@ -116,6 +116,8 @@ def main() -> None: checkpointer, grpo_state, master_config, + _teacher_worker_groups, + _alias_to_group_alias, ) = setup(config, tokenizer, dataset, val_dataset, processor=processor) grpo_train( diff --git a/nemo_rl/algorithms/advantage_estimator.py b/nemo_rl/algorithms/advantage_estimator.py index 789e6cd5ce..59fd0f1ed0 100644 --- a/nemo_rl/algorithms/advantage_estimator.py +++ b/nemo_rl/algorithms/advantage_estimator.py @@ -20,10 +20,12 @@ - ReinforcePlusPlusAdvantageEstimator: Reinforce++ with optional baseline subtraction (minus_baseline) and KL penalty in reward - RawRewardAdvantageEstimator: Raw reward as advantage with optional batch normalization (no baseline, no value model) - GeneralizedAdvantageEstimator: Generalized Advantage Estimation (GAE) with temporal bootstrapping +- OPDAdvantageEstimator: Multi-Teacher On-Policy Distillation (MOPD) token-level distillation advantages Reference papers: - ProRLv2: https://developer.nvidia.com/blog/scaling-llm-reinforcement-learning-with-prolonged-training-using-prorl-v2/ - Reinforce++: https://arxiv.org/abs/2501.03262 - GAE: https://arxiv.org/abs/1506.02438 (High-Dimensional Continuous Control Using Generalized Advantage Estimation) +- MOPD: https://arxiv.org/abs/2601.02780 """ import torch @@ -507,3 +509,82 @@ def _compute_gae( advantages = torch.stack(advantages_reversed[::-1], dim=1) returns = advantages + values return advantages, returns + + +class OPDAdvantageEstimator: + """Multi-Teacher On-Policy Distillation (MOPD) advantage estimator (arXiv:2601.02780). + + Computes token-level distillation advantages: + Â_MOPD,t = sg[log π_teacher - log π_student] + + This is Equation 8 from the MOPD paper. The IS truncation (w_t, the + hard gate on the training-to-inference ratio) is handled separately by + ICE-POP mode in ClippedPGLoss — not here. + + The loss function should be configured with: + disable_ppo_ratio: true (REINFORCE, no PPO ratio) + use_importance_sampling_correction: true + truncated_importance_sampling_type: icepop + truncated_importance_sampling_ratio_min: + truncated_importance_sampling_ratio: + + Required kwargs in compute_advantage: + teacher_logprobs: [B, S] teacher model log probabilities + prev_logprobs: [B, S] student training-engine log probabilities + """ + + def __init__(self, estimator_config: dict, loss_config: dict): + self.last_metrics: dict[str, float] = {} + + def compute_advantage( + self, + prompt_ids, + rewards, + mask, + teacher_logprobs=None, + prev_logprobs=None, + **kwargs, + ): + """Compute OPD distillation advantages. + + Args: + prompt_ids: [B] prompt IDs (unused, kept for interface compatibility) + rewards: [B] rewards (unused for pure distillation) + mask: [B, S] token mask + teacher_logprobs: [B, S] teacher model logprobs (required) + prev_logprobs: [B, S] student training-engine logprobs (required) + + Returns: + [B, S] token-level distillation advantages (stop-gradient) + """ + if teacher_logprobs is None: + raise ValueError("OPD requires teacher_logprobs") + if prev_logprobs is None: + raise ValueError("OPD requires prev_logprobs") + + # Â_MOPD,t = sg[log π_teacher - log π_student] (Equation 8) + distill_advantages = (teacher_logprobs - prev_logprobs).detach() + + # Apply mask + advantages = distill_advantages * mask + + # Metrics + self._compute_metrics(distill_advantages, advantages, mask) + + return advantages + + def _compute_metrics(self, distill_advantages, advantages, mask): + """Compute OPD logging metrics and store in self.last_metrics.""" + valid_bool = mask.bool() + distill_valid = torch.masked_select(distill_advantages, valid_bool) + adv_valid = torch.masked_select(advantages, valid_bool) + + distill_mean = distill_valid.mean().item() if distill_valid.numel() > 0 else 0.0 + adv_mean = adv_valid.mean().item() if adv_valid.numel() > 0 else 0.0 + adv_std = adv_valid.std().item() if adv_valid.numel() > 1 else 0.0 + + self.last_metrics = { + "on_policy_distillation/teacher_student_logprob_gap_mean": distill_mean, + "on_policy_distillation/adv_mean": adv_mean, + "on_policy_distillation/adv_std": adv_std, + } diff --git a/nemo_rl/algorithms/async_utils/trajectory_collector.py b/nemo_rl/algorithms/async_utils/trajectory_collector.py index 12d1b8e3af..b15db43fee 100644 --- a/nemo_rl/algorithms/async_utils/trajectory_collector.py +++ b/nemo_rl/algorithms/async_utils/trajectory_collector.py @@ -12,15 +12,21 @@ # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations + +import concurrent.futures import threading as _threading import time +from collections import defaultdict from typing import Any, Optional import ray +import torch from torchdata.stateful_dataloader import StatefulDataLoader from transformers import PreTrainedTokenizerBase from nemo_rl.algorithms.grpo import MasterConfig +from nemo_rl.algorithms.opd import resolve_reference_aliases, teacher_seq_pad_multiple from nemo_rl.data.interfaces import DatumSpec from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.environments.interfaces import EnvironmentInterface @@ -44,12 +50,31 @@ def __init__( master_config: MasterConfig, replay_buffer: Any, start_step: int = 0, + teacher_worker_groups: Optional[dict[str, Any]] = None, + alias_to_group_alias: Optional[dict[str, str]] = None, + on_policy_distillation_cfg: Optional[dict[str, Any]] = None, ): self.policy_generation = policy_generation self.tokenizer = tokenizer self.task_to_env = task_to_env self.master_config = master_config self.replay_buffer = replay_buffer + self.teacher_worker_groups = teacher_worker_groups or {} + self.alias_to_group_alias = alias_to_group_alias or {} + self.on_policy_distillation_cfg = on_policy_distillation_cfg or {} + self._has_distillation_teachers = bool(self.teacher_worker_groups) + self._teacher_seq_pad_multiple = teacher_seq_pad_multiple( + self.teacher_worker_groups, + self.master_config.policy["make_sequence_length_divisible_by"], + ) + # Per-teacher locks to serialize get_logprobs calls. Concurrent calls + # to the same teacher cause NCCL collective desync across workers + # (different workers may receive requests in different order → SeqNum + # mismatch → 600s timeout → crash). Different teachers can still run + # in parallel since they use separate NCCL groups on separate nodes. + self._teacher_locks: dict[str, _threading.Lock] = { + k: _threading.Lock() for k in self.teacher_worker_groups + } self.running = False self._pg_lock: _threading.Lock = _threading.Lock() @@ -579,6 +604,117 @@ def _maybe_release_target(self, target_weight_version: int) -> None: f"{buffered} buffered)" ) + def _compute_teacher_logprobs( + self, + input_ids: torch.Tensor, + agent_refs: list[dict[str, Any]], + input_lengths: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, float]: + """Compute teacher logprobs for non-colocated teachers. + + Groups samples by teacher, fans out in parallel, stitches results. + + Args: + input_ids: [B, S] tokenized input tensor + agent_refs: list of B agent reference dicts + input_lengths: [B] per-sample lengths (required for sequence packing) + + Returns: + ([B, S] teacher logprobs tensor, total_time_seconds) + """ + opd_cfg = self.on_policy_distillation_cfg + teacher_model_by_agent_name = opd_cfg.get("teacher_model_by_agent_name", {}) + default_teacher_alias = opd_cfg.get("default_teacher_alias") + strict = opd_cfg.get("strict_agent_name_match", False) + + # Resolve each sample's agent -> the teacher alias it should be distilled + # from: the agent name is looked up in teacher_model_by_agent_name; unmapped + # agents fall back to default_teacher_alias (or raise if strict_agent_name_match). + # Returns one alias per sample, index-aligned with agent_refs. + reference_aliases = resolve_reference_aliases( + agent_refs, + teacher_model_by_agent_name, + default_teacher_alias=default_teacher_alias, + strict_agent_name_match=strict, + ) + + # Map aliases to actual group keys via deduplication mapping + group_keys = [self.alias_to_group_alias.get(a, a) for a in reference_aliases] + + # Group sample indices by teacher group + group_to_indices: dict[str, list[int]] = defaultdict(list) + for i, gk in enumerate(group_keys): + group_to_indices[gk].append(i) + + B, S = input_ids.shape + result = torch.zeros(B, S, dtype=torch.float32) + if ( + not group_to_indices + ): # 0-sample batch: nothing to route (avoid max_workers=0) + return result, 0.0 + + def _get_logprobs_for_group(group_key, indices): + twg = self.teacher_worker_groups[group_key] + sub_input_ids = input_ids[indices] + sub_lengths = input_lengths[indices] if input_lengths is not None else None + + # Pad batch to multiple of dp_size (required for DP sharding) + dp_size = twg.sharding_annotations.get_axis_size("data_parallel") + actual_batch_size = sub_input_ids.shape[0] + remainder = actual_batch_size % dp_size + if remainder != 0: + pad_count = dp_size - remainder + # Repeat last row to fill — can't slice [:pad_count] when + # actual_batch_size < pad_count (e.g., 1 sample, dp_size=4) + pad_rows = sub_input_ids[-1:].expand(pad_count, -1) + sub_input_ids = torch.cat([sub_input_ids, pad_rows], dim=0) + if sub_lengths is not None: + sub_lengths = torch.cat( + [sub_lengths, sub_lengths[-1:].expand(pad_count)], dim=0 + ) + + sub_data = BatchedDataDict({"input_ids": sub_input_ids}) + if sub_lengths is not None: + sub_data["input_lengths"] = sub_lengths + + # Serialize calls per teacher to prevent NCCL collective desync + t_lock_start = time.time() + with self._teacher_locks[group_key]: + t_inference_start = time.time() + logprobs_result = twg.get_logprobs(sub_data) + t_done = time.time() + lock_wait = t_inference_start - t_lock_start + inference_time = t_done - t_inference_start + print( + f"[teacher_logprob] group={group_key} samples={actual_batch_size} " + f"lock_wait={lock_wait:.2f}s inference={inference_time:.2f}s" + ) + logprobs = logprobs_result["reference_logprobs"] + + # Trim DP padding + logprobs = logprobs[:actual_batch_size] + + return indices, logprobs + + # Fan out to teachers in parallel + t_total_start = time.time() + with concurrent.futures.ThreadPoolExecutor( + max_workers=len(group_to_indices) + ) as executor: + futures = { + executor.submit(_get_logprobs_for_group, gk, idxs): gk + for gk, idxs in group_to_indices.items() + } + for future in concurrent.futures.as_completed(futures): + indices, logprobs = future.result() + result[indices] = logprobs + total_time = time.time() - t_total_start + print( + f"[teacher_logprob] total={total_time:.2f}s for {B} samples across {len(group_to_indices)} teacher(s)" + ) + + return result, total_time + def _run_prompt_group_worker( self, repeated_batch: BatchedDataDict[DatumSpec], @@ -623,6 +759,34 @@ def _run_prompt_group_worker( final_batch_cpu = final_batch.to("cpu") del final_batch + # Compute teacher logprobs at collection time (overlapped with async rollouts) + if self._has_distillation_teachers and "agent_ref" in final_batch_cpu: + agent_refs = final_batch_cpu["agent_ref"] + if isinstance(agent_refs, list): + from nemo_rl.data.llm_message_utils import ( + batched_message_log_to_flat_message, + ) + + flat_for_teacher, teacher_input_lengths = ( + batched_message_log_to_flat_message( + final_batch_cpu["message_log"], + pad_value_dict={"token_ids": self.tokenizer.pad_token_id}, + make_sequence_length_divisible_by=self._teacher_seq_pad_multiple, + ) + ) + teacher_logprobs, teacher_logprob_time = ( + self._compute_teacher_logprobs( + flat_for_teacher["token_ids"], + agent_refs, + input_lengths=teacher_input_lengths, + ) + ) + # Store inside batch dict so from_batches handles + # variable-length padding across prompt groups + final_batch_cpu["teacher_reference_logprobs"] = teacher_logprobs + rollout_metrics = dict(rollout_metrics) + rollout_metrics["teacher_logprob_time"] = teacher_logprob_time + trajectory_group = { "batch": final_batch_cpu, "rollout_metrics": rollout_metrics, diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index 801970cfcf..7b71680e26 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -28,9 +28,11 @@ from transformers import AutoProcessor from transformers.tokenization_utils_base import PreTrainedTokenizerBase +from nemo_rl.algorithms import opd as opd_module from nemo_rl.algorithms.advantage_estimator import ( GDPOAdvantageEstimator, GRPOAdvantageEstimator, + OPDAdvantageEstimator, ReinforcePlusPlusAdvantageEstimator, ) from nemo_rl.algorithms.loss import ( @@ -39,6 +41,7 @@ ClippedPGLossFn, ) from nemo_rl.algorithms.loss.interfaces import LossFunction +from nemo_rl.algorithms.opd import OnPolicyDistillationConfig from nemo_rl.algorithms.reward_functions import ( RewardShapingConfig, apply_reward_shaping, @@ -240,6 +243,7 @@ class MasterConfig(BaseModel, extra="allow"): cluster: ClusterConfig checkpointing: CheckpointingConfig data_plane: Optional[DataPlaneConfig] = None + on_policy_distillation: Optional[OnPolicyDistillationConfig] = None # =============================================================================== @@ -266,14 +270,17 @@ def setup( CheckpointManager, GRPOSaveState, MasterConfig, + dict[str, Any], + dict[str, str], ]: """Main entry point for running GRPO algorithm. Returns: - An 11-tuple, in order: + A 13-tuple, in order: policy, policy_generation, nemo_gym (the NeMo-Gym env actor, or None when not enabled), cluster, dataloader, val_dataloader, loss_fn, - logger, checkpointer, grpo_save_state, master_config. + logger, checkpointer, grpo_save_state, master_config, + teacher_worker_groups, alias_to_group_alias. """ # Start timing the entire setup process setup_start_time = time.perf_counter() @@ -511,6 +518,9 @@ def _spinup_nemo_gym(base_urls, model_name): total_nodes = cluster_config["num_nodes"] segment_size = cluster_config.get("segment_size") + # Topology of nodes left over after policy/inference placement; non-colocated + # OPD teachers are placed within it so their collectives stay on NVLink. + teacher_segment_topology: Optional[dict[str, tuple[str, int]]] = None if rm_env_enabled: rm_resource = env_configs["reward_model"]["resources"] rm_nodes = rm_resource["num_nodes"] @@ -528,6 +538,31 @@ def _spinup_nemo_gym(base_urls, model_name): f"policy_nodes:{policy_nodes} + rm_nodes:{rm_nodes} = total_nodes:{total_nodes}" ) + # Reserve nodes for non-colocated OPD teachers so training doesn't claim them. + opd_teacher_nodes = 0 + enable_opd_teachers = opd_module.is_non_colocated_teachers_enabled(master_config) + if enable_opd_teachers: + assert _should_use_async_rollouts(master_config), ( + "Non-colocated OPD teachers require async GRPO (vLLM backend with async_engine enabled)." + ) + from nemo_rl.models.policy.teacher_worker_group import ( + create_teacher_configs_from_opd_config, + ) + + opd_cfg = opd_module._opd_cfg(master_config) + teacher_configs = create_teacher_configs_from_opd_config(opd_cfg) + for tcfg in teacher_configs: + opd_teacher_nodes += tcfg.num_nodes + policy_nodes -= opd_teacher_nodes + assert policy_nodes > 0, ( + "policy_nodes must be > 0 after reserving OPD teacher nodes, but got " + f"policy_nodes:{policy_nodes} + rm_nodes:{rm_nodes} + opd_teacher_nodes:{opd_teacher_nodes} = total_nodes:{total_nodes}" + ) + print( + f"policy_nodes:{policy_nodes} + rm_nodes:{rm_nodes} + opd_teacher_nodes:{opd_teacher_nodes} = total_nodes:{total_nodes}", + flush=True, + ) + if colocated_inference: if total_nodes == 1: policy_gpus_per_node = cluster_config["gpus_per_node"] - rm_gpus_per_node @@ -539,9 +574,13 @@ def _spinup_nemo_gym(base_urls, model_name): else: policy_gpus_per_node = cluster_config["gpus_per_node"] - node_resource_constraints, _, _ = prepare_segment_topology( - segment_size, policy_nodes + node_resource_constraints, policy_remaining_ids, policy_topology = ( + prepare_segment_topology(segment_size, policy_nodes) ) + if segment_size is not None: + teacher_segment_topology = { + nid: policy_topology[nid] for nid in policy_remaining_ids + } cluster = RayVirtualCluster( name="grpo_policy_cluster", bundle_ct_per_node_list=[policy_gpus_per_node] * policy_nodes, @@ -650,6 +689,11 @@ def _spinup_nemo_gym(base_urls, model_name): segment_size, train_nodes, topology=topology, role="training" ) ) + # Teachers default to the nodes left after training; narrowed further + # below if a non-colocated inference cluster is also pinned. + teacher_segment_topology = { + nid: topology[nid] for nid in remaining_node_ids + } # Warn if any selected training node lacks topo_rank — domain pinning # still works but intra-domain rank ordering will be arbitrary. if node_resource_constraints is not None: @@ -689,15 +733,20 @@ def _spinup_nemo_gym(base_urls, model_name): remaining_topology = { nid: topology[nid] for nid in remaining_node_ids } - inference_node_resource_constraints, _, _ = ( - prepare_segment_topology( - nodes_per_instance, - inference_nodes, - topology=remaining_topology, - role="inference", - ) + ( + inference_node_resource_constraints, + inference_remaining_ids, + _, + ) = prepare_segment_topology( + nodes_per_instance, + inference_nodes, + topology=remaining_topology, + role="inference", ) inference_segment_size = nodes_per_instance + teacher_segment_topology = { + nid: topology[nid] for nid in inference_remaining_ids + } elif nodes_per_instance > 1: print( f" ⚠ inference_nodes={inference_nodes} is not divisible by " @@ -1126,6 +1175,25 @@ def init_vllm_then_policy(): if policy_generation is not None: policy_generation.prepare_refit_info(state_dict_info) + # Spin up non-colocated OPD teacher worker groups AFTER policy / vLLM are + # ready. Parallelizing with policy init races on Megatron-Bridge's HF->mcore + # cache (shared key when student == teacher) — both workers write to the + # same iter_0000000/ path and the second reader gets a truncated file. + teacher_worker_groups: dict[str, Any] = {} + alias_to_group_alias: dict[str, str] = {} + if enable_opd_teachers: + t0 = time.perf_counter() + teacher_worker_groups, alias_to_group_alias = ( + opd_module.create_teacher_worker_groups( + master_config, + policy_config, + tokenizer, + segment_size=segment_size, + teacher_segment_topology=teacher_segment_topology, + ) + ) + worker_init_timing_metrics["teacher_init_time_s"] = time.perf_counter() - t0 + # Calculate total setup time total_setup_time = time.perf_counter() - setup_start_time worker_init_timing_metrics["total_setup_time_s"] = total_setup_time @@ -1144,6 +1212,10 @@ def init_vllm_then_policy(): if policy_time: print(f" Policy init: {policy_time:.1f}s") + teacher_time = worker_init_timing_metrics.get("teacher_init_time_s", 0) + if teacher_time: + print(f" Teacher init: {teacher_time:.1f}s") + # Calculate "other" time (time after worker init completes) other_time = total_setup - worker_init_complete_time worker_init_timing_metrics["other_setup_time_s"] = other_time @@ -1171,6 +1243,8 @@ def init_vllm_then_policy(): checkpointer, grpo_save_state, master_config, + teacher_worker_groups, + alias_to_group_alias, ) @@ -1614,6 +1688,28 @@ def _get_effort_config(master_config: MasterConfig) -> Optional[EffortLevelsConf return EffortLevelsConfig.model_validate(effort_dict) +def _pad_teacher_logprobs(teacher_logprobs: torch.Tensor, train_S: int) -> torch.Tensor: + """Right-zero-pad teacher logprobs ``[B, teacher_S]`` to ``train_S``. + + ``from_batches`` pads teacher logprobs to ``max(S_i)``; ``train_data`` may be + longer due to ``make_sequence_length_divisible_by``. Zero-pad is safe because + the mask zeros padding in advantage computation. ``teacher_S > train_S`` is + unexpected (teacher pads to a finer grid than the student) and raises. + """ + teacher_S = teacher_logprobs.shape[1] + if teacher_S > train_S: + raise ValueError( + f"Teacher logprobs seq length ({teacher_S}) > train_data seq length ({train_S}). " + "Teacher logprobs are padded to max(S_i) by from_batches, " + "and train_data is padded to roundup(max(S_i), make_sequence_length_divisible_by)." + ) + if teacher_S < train_S: + teacher_logprobs = torch.nn.functional.pad( + teacher_logprobs, (0, train_S - teacher_S), value=0.0 + ) + return teacher_logprobs + + def _create_advantage_estimator(master_config: MasterConfig): """Create and return an advantage estimator based on configuration. @@ -1651,6 +1747,24 @@ def _create_advantage_estimator(master_config: MasterConfig): elif adv_estimator_name == "grpo": adv_estimator = GRPOAdvantageEstimator(adv_estimator_config, loss_config) print(" ✓ Using GRPO advantage estimator") + elif adv_estimator_name == "opd": + opd_module.assert_prev_logprobs_available(master_config) + adv_estimator = OPDAdvantageEstimator({"name": "opd"}, loss_config) + print(" ✓ Using OPD advantage estimator") + # Warn if loss_fn is not configured per MOPD paper recommendations. + if not loss_config.disable_ppo_ratio: + warnings.warn( + "OPD recommends loss_fn.disable_ppo_ratio: true (REINFORCE-style, MOPD Eq. 7)" + ) + if not loss_config.use_importance_sampling_correction: + warnings.warn( + "OPD recommends loss_fn.use_importance_sampling_correction: true (MOPD Eq. 8 w_t)" + ) + if loss_config.truncated_importance_sampling_type != "icepop": + warnings.warn( + "OPD recommends loss_fn.truncated_importance_sampling_type: 'icepop' " + "(hard gate, MOPD Eq. 8)" + ) elif adv_estimator_name == "reinforce_plus_plus": adv_estimator = ReinforcePlusPlusAdvantageEstimator( adv_estimator_config, loss_config @@ -2399,9 +2513,7 @@ def grpo_train( "seq_logprob_error_threshold", None ) force_on_policy_ratio = master_config.loss_fn.force_on_policy_ratio - skip_prev_logprobs = ( - force_on_policy_ratio and seq_logprob_error_threshold is None - ) + skip_prev_logprobs = opd_module._skip_prev_logprobs(master_config) # todo @jiaqi: is there a better way to skip prev_logprobs computation while still computing the seq-level error metrics? if force_on_policy_ratio and seq_logprob_error_threshold is not None: warnings.warn( @@ -3152,6 +3264,8 @@ def async_grpo_train( grpo_save_state: GRPOSaveState, master_config: MasterConfig, max_trajectory_age_steps: int = 1, + teacher_worker_groups: Optional[dict[str, Any]] = None, + alias_to_group_alias: Optional[dict[str, str]] = None, ) -> None: """Run asynchronous GRPO training with replay buffer. @@ -3335,6 +3449,9 @@ def async_grpo_train( master_config=master_config, replay_buffer=replay_buffer, start_step=step, + teacher_worker_groups=teacher_worker_groups, + alias_to_group_alias=alias_to_group_alias, + on_policy_distillation_cfg=opd_module._opd_cfg(master_config), ) # Start trajectory collection in background @@ -3516,6 +3633,16 @@ def async_grpo_train( # Concatenate per-prompt groups into a single training batch per_prompt_batches = [t["batch"] for t in trajectories] repeated_batch = BatchedDataDict.from_batches(per_prompt_batches) + + # Teacher logprobs are stored in batch dict by collection-time + # computation and padded by from_batches. Extract here. + trajectory_teacher_logprobs = None + if opd_module.is_opd_enabled(master_config): + if "teacher_reference_logprobs" in repeated_batch: + trajectory_teacher_logprobs = repeated_batch[ + "teacher_reference_logprobs" + ] + # Aggregate rollout metrics across groups with proper aggregation per metric type per_group_metrics = {} for t in trajectories: @@ -3620,9 +3747,7 @@ def async_grpo_train( "seq_logprob_error_threshold", None ) force_on_policy_ratio = master_config.loss_fn.force_on_policy_ratio - skip_prev_logprobs = ( - force_on_policy_ratio and seq_logprob_error_threshold is None - ) + skip_prev_logprobs = opd_module._skip_prev_logprobs(master_config) # todo @jiaqi: is there a better way to skip prev_logprobs computation while still computing the seq-level error metrics? if force_on_policy_ratio and seq_logprob_error_threshold is not None: @@ -3693,6 +3818,12 @@ def async_grpo_train( "num_masked_seqs_by_logprob_error" ] = seq_logprob_error_metrics.pop("num_masked_seqs") + # Pad teacher logprobs to match train_data sequence length. + if trajectory_teacher_logprobs is not None: + trajectory_teacher_logprobs = _pad_teacher_logprobs( + trajectory_teacher_logprobs, train_data["input_ids"].shape[1] + ) + # Compute advantages with adv_estimator using correct mask and logprobs with timer.time("advantage_calculation"): print("▶ Computing advantages...", flush=True) @@ -3708,7 +3839,21 @@ def async_grpo_train( repeated_batch=repeated_batch, logprobs_policy=train_data["prev_logprobs"], logprobs_reference=train_data.get("reference_policy_logprobs"), + # OPD kwargs (ignored by non-OPD estimators via **kwargs) + teacher_logprobs=trajectory_teacher_logprobs.to( + train_data["prev_logprobs"].device + ) + if trajectory_teacher_logprobs is not None + else None, + prev_logprobs=train_data["prev_logprobs"], + generation_logprobs=train_data["generation_logprobs"], + sample_mask=train_data["sample_mask"], ) + if ( + hasattr(adv_estimator, "last_metrics") + and adv_estimator.last_metrics + ): + rollout_metrics.update(adv_estimator.last_metrics) del prompt_ids_for_adv # Log advantages stats diff --git a/nemo_rl/algorithms/opd.py b/nemo_rl/algorithms/opd.py new file mode 100644 index 0000000000..21087d49e3 --- /dev/null +++ b/nemo_rl/algorithms/opd.py @@ -0,0 +1,358 @@ +# 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. + +"""On-policy distillation (OPD) helpers for async GRPO. + +Teacher routing, config helpers, and teacher worker group creation. +Advantage computation lives in advantage_estimator.OPDAdvantageEstimator. +IS truncation lives in loss_functions.ClippedPGLoss (ICE-POP mode). +""" + +from __future__ import annotations + +from typing import Any, Optional + +import ray +from pydantic import BaseModel, Field + +# --------------------------------------------------------------------------- +# Config schemas +# --------------------------------------------------------------------------- + + +class TeacherResourceConfig(BaseModel, extra="allow"): + """Per-teacher resourcing for a non-colocated teacher worker group. + + ``extra="allow"`` keeps the escape hatch for arbitrary megatron settings: + any unknown top-level key is folded into ``megatron_cfg_overrides``. + """ + + tensor_model_parallel_size: int = 1 + pipeline_model_parallel_size: int = 1 + context_parallel_size: int = 1 + expert_model_parallel_size: int = 1 + num_nodes: int = 1 + gpus_per_node: int = 8 + precision: str = "bf16" + micro_batch_size: int = 4 + megatron_cfg_overrides: dict[str, Any] = Field(default_factory=dict) + + +class NonColocatedTeachersConfig(BaseModel, extra="allow"): + """Non-colocated (separate-GPU) teacher resourcing for on-policy distillation.""" + + enabled: bool = False + default_teacher_cfg: TeacherResourceConfig = Field( + default_factory=TeacherResourceConfig + ) + teacher_overrides: dict[str, TeacherResourceConfig] = Field(default_factory=dict) + + +class OnPolicyDistillationConfig(BaseModel, extra="allow"): + """User-facing config for the top-level ``on_policy_distillation`` block.""" + + enabled: bool = False + teacher_model_by_agent_name: dict[str, str] = Field(default_factory=dict) + default_teacher_alias: Optional[str] = None + strict_agent_name_match: bool = False + deduplicate_shared_teacher_checkpoints: bool = True + non_colocated_teachers: Optional[NonColocatedTeachersConfig] = None + + +# --------------------------------------------------------------------------- +# Config helpers +# --------------------------------------------------------------------------- + + +def _opd_cfg(master_config: Any) -> dict[str, Any]: + """Return the on_policy_distillation sub-config as a plain dict. + + Accepts a MasterConfig (where the field is an OnPolicyDistillationConfig + BaseModel), a plain dict, or a config object missing the field (non-OPD + recipes like math). Downstream code reads the result dict-style. + """ + if isinstance(master_config, dict): + cfg = master_config.get("on_policy_distillation") + else: + cfg = getattr(master_config, "on_policy_distillation", None) + if cfg is None: + return {} + if isinstance(cfg, BaseModel): + return cfg.model_dump(exclude_none=True) + return cfg + + +def is_opd_enabled(master_config: Any) -> bool: + """Whether on-policy distillation is enabled in the config.""" + return bool(_opd_cfg(master_config).get("enabled", False)) + + +def is_non_colocated_teachers_enabled(master_config: Any) -> bool: + """Whether OPD is enabled with non-colocated (separate-GPU) teachers.""" + if not is_opd_enabled(master_config): + return False + return bool( + _opd_cfg(master_config).get("non_colocated_teachers", {}).get("enabled", False) + ) + + +def _skip_prev_logprobs(master_config: Any) -> bool: + """Whether the training loop will zero ``prev_logprobs`` instead of computing it. + + Mirrors the predicate in ``grpo_train``: ``force_on_policy_ratio`` with no + ``seq_logprob_error_threshold`` skips the student logprob pass. + """ + force_on_policy_ratio = master_config.loss_fn.force_on_policy_ratio + seq_logprob_error_threshold = master_config.grpo.get( + "seq_logprob_error_threshold", None + ) + return bool(force_on_policy_ratio and seq_logprob_error_threshold is None) + + +def assert_prev_logprobs_available(master_config: Any) -> None: + """Raise if OPD is enabled but the config would zero ``prev_logprobs``. + + OPD's advantage is ``teacher_logprobs - prev_logprobs``, so it needs a real + student logprob. + """ + if is_opd_enabled(master_config) and _skip_prev_logprobs(master_config): + raise ValueError( + "adv_estimator='opd' requires real prev_logprobs, but the config zeros them " + "(loss_fn.force_on_policy_ratio=True with grpo.seq_logprob_error_threshold unset). " + "Set seq_logprob_error_threshold or disable force_on_policy_ratio." + ) + + +# --------------------------------------------------------------------------- +# Teacher routing +# --------------------------------------------------------------------------- + + +def resolve_reference_aliases( + agent_refs: list[dict], + teacher_model_by_agent_name: dict[str, str], + default_teacher_alias: Optional[str] = None, + strict_agent_name_match: bool = False, +) -> list[str]: + """Map each agent_ref to a teacher alias. + + Unmapped agents fall back to ``default_teacher_alias``; with + ``strict_agent_name_match`` an unmapped agent raises instead. + """ + aliases: list[str] = [] + for ref in agent_refs: + name = ref["name"] + if name in teacher_model_by_agent_name: + aliases.append(name) + elif strict_agent_name_match: + raise ValueError( + f"No teacher model mapping for agent '{name}'. " + f"Available: {sorted(teacher_model_by_agent_name.keys())}" + ) + elif default_teacher_alias: + print( + f"[OPD] Agent '{name}' not in teacher mapping, falling back to '{default_teacher_alias}'" + ) + aliases.append(default_teacher_alias) + else: + raise ValueError( + f"No teacher model mapping for agent '{name}' and no default_teacher_alias set." + ) + return aliases + + +def get_teacher_routing_metrics( + reference_aliases: list[str], + teacher_model_by_agent_name: dict[str, str], +) -> dict[str, float]: + """Compute teacher-routing diagnostics. + + Reports unique aliases, unique underlying models, and the alias→model + compression ratio (how many aliases share each underlying teacher model). + """ + alias_unique = len(set(reference_aliases)) + unique_models: set[str] = set() + for alias in reference_aliases: + if alias not in teacher_model_by_agent_name: + raise KeyError(f"Alias '{alias}' not found in teacher_model_by_agent_name") + unique_models.add(teacher_model_by_agent_name[alias]) + model_unique = len(unique_models) + return { + "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) + ), + } + + +# --------------------------------------------------------------------------- +# Setup helper — teacher worker group creation +# --------------------------------------------------------------------------- + + +def teacher_seq_pad_multiple( + teacher_worker_groups: dict[str, Any], policy_make_seq_div_by: int +) -> int: + """Sequence divisor to pre-pad teacher logprob inputs to. + + Packed teachers re-pad internally, so no pre-pad is needed (1). Non-packed + teachers need the ``[B, S]`` forward pre-padded to the policy divisor, which + must be a multiple of every teacher's ``sequence_length_pad_multiple``. All + teachers must share one packing mode. + """ + packing_modes = {twg.use_sequence_packing for twg in teacher_worker_groups.values()} + if len(packing_modes) > 1: + raise ValueError("All teachers must use the same sequence-packing mode.") + if packing_modes != {False}: + return 1 # no teachers, or all packed (they re-pad internally) + for alias, twg in teacher_worker_groups.items(): + if policy_make_seq_div_by % twg.sequence_length_pad_multiple: + raise ValueError( + f"policy.make_sequence_length_divisible_by ({policy_make_seq_div_by}) " + f"must be a multiple of teacher '{alias}'s pad requirement " + f"({twg.sequence_length_pad_multiple})." + ) + return policy_make_seq_div_by + + +def create_teacher_worker_groups( + master_config: Any, + policy_config: dict[str, Any], + tokenizer: Any, + *, + segment_size: Optional[int] = None, + teacher_segment_topology: Optional[dict[str, tuple[str, int]]] = None, +) -> tuple[dict[str, Any], dict[str, str]]: + """Create TeacherWorkerGroup instances for non-colocated teachers. + + Args: + segment_size: NVLink-domain segment size from the cluster config; when + set, each teacher is placed topology-aware so its TP/PP/CP stays + within an NVLink domain. + teacher_segment_topology: Topology of the nodes left after policy / + inference placement, used to pin teacher nodes (see + ``prepare_segment_topology``). + + Returns (teacher_worker_groups, alias_to_group_alias). + """ + from nemo_rl.distributed.virtual_cluster import ( + RayVirtualCluster, + prepare_segment_topology, + ) + from nemo_rl.models.policy.teacher_worker_group import ( + TeacherWorkerGroup, + create_teacher_configs_from_opd_config, + ) + + opd_cfg = _opd_cfg(master_config) + teacher_model_by_agent_name = dict(opd_cfg.get("teacher_model_by_agent_name", {})) + + # A non-strict run falls back to default_teacher_alias for unmapped agents, so + # it must itself be a mapped agent. + default_teacher_alias = opd_cfg.get("default_teacher_alias") + if ( + not opd_cfg.get("strict_agent_name_match", False) + and default_teacher_alias is not None + and default_teacher_alias not in teacher_model_by_agent_name + ): + raise ValueError( + f"default_teacher_alias '{default_teacher_alias}' is not a key in " + f"teacher_model_by_agent_name (available: " + f"{sorted(teacher_model_by_agent_name.keys())})." + ) + + teacher_configs = create_teacher_configs_from_opd_config(opd_cfg) + + # Running topology of still-free nodes; each teacher consumes a segment and + # passes the remainder to the next so teachers don't collide. + running_topology = ( + dict(teacher_segment_topology) if teacher_segment_topology else None + ) + + teacher_worker_groups: dict[str, Any] = {} + for tcfg in teacher_configs: + alias = tcfg.alias + num_nodes = tcfg.num_nodes + gpus_per_node = tcfg.gpus_per_node + + # Pin each teacher within one NVLink domain (its whole node span is one + # segment) so its TP/PP/CP collectives stay on NVLink, not InfiniBand. + teacher_segment_size = None + node_resource_constraints = None + if segment_size is not None: + teacher_segment_size = num_nodes + node_resource_constraints, remaining_ids, _ = prepare_segment_topology( + num_nodes, + num_nodes, + topology=running_topology, + role=f"teacher:{alias}", + ) + if running_topology is not None: + running_topology = {nid: running_topology[nid] for nid in remaining_ids} + + teacher_cluster = RayVirtualCluster( + name=f"teacher_{alias}", + bundle_ct_per_node_list=[gpus_per_node] * num_nodes, + use_gpus=True, + num_gpus_per_node=gpus_per_node, + max_colocated_worker_groups=1, + segment_size=teacher_segment_size, + node_resource_constraints=node_resource_constraints, + ) + # Eagerly claim domain-aligned nodes before the next teacher selects. + if node_resource_constraints is not None: + teacher_cluster.get_placement_groups() + twg = TeacherWorkerGroup( + teacher_cfg=tcfg, + cluster=teacher_cluster, + policy_config=policy_config, + tokenizer=tokenizer, + ) + teacher_worker_groups[alias] = twg + print( + f" ✓ Teacher '{alias}' cluster: {num_nodes} node(s), {gpus_per_node} GPUs/node", + flush=True, + ) + + # Verify all teacher workers are alive (actor __init__ runs async and + # failures are otherwise silent until the first remote call). + print(" Verifying teacher workers are healthy...", flush=True) + for alias, twg in teacher_worker_groups.items(): + try: + refs = [w.__ray_ready__.remote() for w in twg.worker_group.workers] + ray.get(refs, timeout=1800) + except Exception as e: + raise RuntimeError( + f"Teacher '{alias}' worker(s) failed during initialization. " + f"This often means a stale cached mcore checkpoint — try deleting " + f"the cached checkpoint under $HF_HOME/nemo_rl/ and rerunning.\n" + f"Original error: {e}" + ) from e + print(" ✓ All teacher workers healthy", flush=True) + + # Reject a mixed/incompatible teacher packing config (raises). + teacher_seq_pad_multiple( + teacher_worker_groups, policy_config["make_sequence_length_divisible_by"] + ) + + # Build alias -> group_alias mapping for deduplication + alias_to_group_alias: dict[str, str] = {} + model_to_primary: dict[str, str] = {} + for tcfg in teacher_configs: + model_to_primary[tcfg.model_name] = tcfg.alias + for alias, model_name in teacher_model_by_agent_name.items(): + alias_to_group_alias[alias] = model_to_primary.get(model_name, alias) + + return teacher_worker_groups, alias_to_group_alias diff --git a/nemo_rl/models/policy/teacher_worker_group.py b/nemo_rl/models/policy/teacher_worker_group.py new file mode 100644 index 0000000000..3a1068bd45 --- /dev/null +++ b/nemo_rl/models/policy/teacher_worker_group.py @@ -0,0 +1,305 @@ +# 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. + +"""Non-colocated teacher worker group for MOPD async distillation. + +Each TeacherWorkerGroup wraps a RayWorkerGroup running MegatronPolicyWorker +in inference-only mode for a single teacher model checkpoint. +""" + +from __future__ import annotations + +import warnings +from copy import deepcopy +from dataclasses import dataclass +from typing import Any, Optional + +import numpy as np +from transformers import PreTrainedTokenizerBase + +from nemo_rl.algorithms.opd import TeacherResourceConfig +from nemo_rl.distributed.batched_data_dict import ( + BatchedDataDict, + SequencePackingArgs, +) +from nemo_rl.distributed.named_sharding import NamedSharding +from nemo_rl.distributed.virtual_cluster import RayVirtualCluster +from nemo_rl.models.generation.interfaces import GenerationDatumSpec +from nemo_rl.models.policy.interfaces import ReferenceLogprobOutputSpec + + +@dataclass +class TeacherConfig: + """Resolved config for a single non-colocated teacher (built in-process).""" + + alias: str + model_name: str # checkpoint path + tensor_model_parallel_size: int + pipeline_model_parallel_size: int + context_parallel_size: int + expert_model_parallel_size: int + num_nodes: int + gpus_per_node: int + precision: str + micro_batch_size: int + megatron_cfg_overrides: dict[str, Any] + + +def create_teacher_configs_from_opd_config( + opd_cfg: dict[str, Any], +) -> list[TeacherConfig]: + """Build per-teacher configs from on_policy_distillation config. + + Handles deduplication (multiple aliases sharing one checkpoint produce + one TeacherConfig) and per-teacher overrides on top of defaults. + """ + teacher_model_by_agent_name: dict[str, str] = dict( + opd_cfg.get("teacher_model_by_agent_name", {}) + ) + non_coloc_cfg = dict(opd_cfg.get("non_colocated_teachers", {})) + default_cfg = dict(non_coloc_cfg.get("default_teacher_cfg", {})) + overrides = dict(non_coloc_cfg.get("teacher_overrides", {})) + deduplicate = bool(opd_cfg.get("deduplicate_shared_teacher_checkpoints", True)) + + configs: list[TeacherConfig] = [] + seen_models: set[str] = set() + + for alias, model_name in teacher_model_by_agent_name.items(): + if deduplicate and model_name in seen_models: + continue + seen_models.add(model_name) + + # defaults <- per-alias override, then validated/typed by the schema. + merged = {**default_cfg, **dict(overrides.get(alias, {}))} + res = TeacherResourceConfig(**merged) + + # Unknown top-level keys (extra="allow") fold into megatron_cfg_overrides; + # explicit megatron_cfg_overrides take precedence. + all_overrides = {**(res.model_extra or {}), **res.megatron_cfg_overrides} + + configs.append( + TeacherConfig( + alias=alias, + model_name=model_name, + tensor_model_parallel_size=res.tensor_model_parallel_size, + pipeline_model_parallel_size=res.pipeline_model_parallel_size, + context_parallel_size=res.context_parallel_size, + expert_model_parallel_size=res.expert_model_parallel_size, + num_nodes=res.num_nodes, + gpus_per_node=res.gpus_per_node, + precision=res.precision, + micro_batch_size=res.micro_batch_size, + megatron_cfg_overrides=all_overrides, + ) + ) + + return configs + + +class TeacherWorkerGroup: + """Inference-only mcore worker group for a single teacher model. + + Unlike the training policy, this group: + - Never initializes an optimizer + - Never initializes a reference model + - Loads the checkpoint once at startup + - Only exposes get_logprobs() + """ + + def __init__( + self, + teacher_cfg: TeacherConfig, + cluster: RayVirtualCluster, + policy_config: dict[str, Any], + tokenizer: PreTrainedTokenizerBase, + ): + self.alias = teacher_cfg.alias + self.model_name = teacher_cfg.model_name + self.teacher_cfg = teacher_cfg + + # Build a policy config for inference-only use. + cfg = deepcopy(policy_config) + cfg["model_name"] = self.model_name + # Override parallelism from teacher config. + if "megatron_cfg" not in cfg: + cfg["megatron_cfg"] = {} + cfg["megatron_cfg"]["enabled"] = True + cfg["megatron_cfg"]["tensor_model_parallel_size"] = ( + teacher_cfg.tensor_model_parallel_size + ) + cfg["megatron_cfg"]["pipeline_model_parallel_size"] = ( + teacher_cfg.pipeline_model_parallel_size + ) + cfg["megatron_cfg"]["context_parallel_size"] = teacher_cfg.context_parallel_size + cfg["megatron_cfg"]["expert_model_parallel_size"] = ( + teacher_cfg.expert_model_parallel_size + ) + + # Apply any additional megatron config overrides from teacher config. + for key, value in teacher_cfg.megatron_cfg_overrides.items(): + cfg["megatron_cfg"][key] = value + + # Teachers run Megatron inference-only. Don't let the student's other + # backend or parameter-adding features leak onto the frozen teacher. + if cfg.get("dtensor_cfg", {}).get("enabled", False): + raise ValueError( + f"Teacher '{self.alias}': only the Megatron backend is supported " + "for teachers, but the policy config has dtensor_cfg.enabled=True." + ) + if "dtensor_cfg" in cfg: + cfg["dtensor_cfg"]["enabled"] = False + if "peft" in cfg["megatron_cfg"]: + cfg["megatron_cfg"]["peft"]["enabled"] = False + if "draft" in cfg: + cfg["draft"]["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: + warnings.warn( + f"Teacher '{self.alias}': quantization is not supported for teachers; " + "running the teacher unquantized (ignoring the policy's quant_cfg)." + ) + cfg["quant_cfg"] = None + + tp = teacher_cfg.tensor_model_parallel_size + pp = teacher_cfg.pipeline_model_parallel_size + cp = teacher_cfg.context_parallel_size + + # Validate parallelism fits the cluster (matches lm_policy.py) + world_size = cluster.world_size() + model_parallel_size = tp * pp * cp + if world_size < model_parallel_size: + raise ValueError( + f"Teacher '{self.alias}': world_size ({world_size}) < TP({tp}) * PP({pp}) * CP({cp}) = {model_parallel_size}" + ) + if world_size % model_parallel_size != 0: + raise ValueError( + f"Teacher '{self.alias}': world_size ({world_size}) not divisible by TP({tp}) * PP({pp}) * CP({cp}) = {model_parallel_size}" + ) + + self.sharding_annotations = NamedSharding( + layout=np.arange(world_size).reshape(pp, -1, cp, tp), + names=[ + "pipeline_parallel", + "data_parallel", + "context_parallel", + "tensor_parallel", + ], + ) + + from nemo_rl.distributed.worker_groups import RayWorkerBuilder, RayWorkerGroup + + worker_builder = RayWorkerBuilder( + "nemo_rl.models.policy.workers.megatron_policy_worker.MegatronPolicyWorker", + cfg, + tokenizer=tokenizer, + processor=None, + init_optimizer=False, + weights_path=None, + optimizer_path=None, + init_reference_model=False, + worker_sharding_annotations=self.sharding_annotations, + ) + + env_vars = cfg["megatron_cfg"].get("env_vars", {}) + + self.worker_group = RayWorkerGroup( + cluster, + worker_builder, + name_prefix=f"teacher_{self.alias}", + sharding_annotations=self.sharding_annotations, + env_vars=env_vars or {}, + ) + + self.cfg = cfg + self._micro_batch_size = teacher_cfg.micro_batch_size + + # Set up sequence packing / dynamic batching (mirrors lm_policy.py) + self.use_sequence_packing = cfg["sequence_packing"]["enabled"] + self.use_dynamic_batches = cfg["dynamic_batching"]["enabled"] + # SP-forward divisor; the collector reads it to pre-pad non-packed inputs. + self.sequence_length_pad_multiple = cp * 2 * tp if cp > 1 else tp + if self.use_sequence_packing: + self.sequence_packing_args: SequencePackingArgs = { + "algorithm": cfg["sequence_packing"]["algorithm"], + "input_key": "input_ids", + "input_lengths_key": "input_lengths", + "sequence_length_pad_multiple": self.sequence_length_pad_multiple, + } + + def get_logprobs( + self, + data: BatchedDataDict[GenerationDatumSpec], + micro_batch_size: Optional[int] = None, + ) -> BatchedDataDict[ReferenceLogprobOutputSpec]: + """Run forward pass on teacher and return logprobs.""" + dp_size = self.sharding_annotations.get_axis_size("data_parallel") + mbs = micro_batch_size or self._micro_batch_size + + if self.use_sequence_packing: + self.sequence_packing_args["max_tokens_per_microbatch"] = self.cfg[ + "sequence_packing" + ]["logprob_mb_tokens"] + sharded_data, unsorted_data_indices = data.shard_by_batch_size( + dp_size, + batch_size=None, + sequence_packing_args=self.sequence_packing_args, + ) + else: + sharded_data = data.shard_by_batch_size(dp_size, batch_size=None) + unsorted_data_indices = None + + futures = self.worker_group.run_all_workers_sharded_data( + "get_logprobs", + data=sharded_data, + 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": mbs}, + ) + logprobs = BatchedDataDict.from_batches( + self.worker_group.get_all_worker_results(futures) + ) + + result = BatchedDataDict[ReferenceLogprobOutputSpec]( + reference_logprobs=logprobs["logprobs"].cpu() + ) + + # Undo packing reorder if needed — must use inverse permutation + # (argsort), matching lm_policy.py's reorder_data. + if unsorted_data_indices is not None: + result.reorder_data(unsorted_data_indices) + + return result + + def shutdown(self) -> bool: + """Shut down all workers and clean up resources.""" + try: + return self.worker_group.shutdown(cleanup_method="shutdown") + except Exception as e: + print(f"Error during teacher worker group shutdown: {e}") + return False + + def __del__(self) -> None: + """Safety net for cleanup.""" + if hasattr(self, "worker_group"): + self.shutdown() diff --git a/pyrefly.toml b/pyrefly.toml index bd1a5bae22..c8fd672555 100644 --- a/pyrefly.toml +++ b/pyrefly.toml @@ -49,6 +49,7 @@ project-includes = [ "nemo_rl/algorithms/loss/__init__.py", "nemo_rl/algorithms/loss/interfaces.py", "nemo_rl/algorithms/loss/utils.py", + "nemo_rl/algorithms/opd.py", "nemo_rl/algorithms/reward_functions.py", "nemo_rl/algorithms/utils.py", "nemo_rl/algorithms/x_token/__init__.py", diff --git a/tests/test_suites/llm/mopd-qwen3-1.7b-3n8g-megatron-pack.sh b/tests/test_suites/llm/mopd-qwen3-1.7b-3n8g-megatron-pack.sh new file mode 100755 index 0000000000..f845c75576 --- /dev/null +++ b/tests/test_suites/llm/mopd-qwen3-1.7b-3n8g-megatron-pack.sh @@ -0,0 +1,62 @@ +#!/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. +# +# MOPD: dense Qwen3-1.7B student distilled from a Qwen3-1.7B +# teacher (student == teacher -> OPD loss ~0), sequence packing ON, 3 nodes +# (1 policy + 1 vLLM + 1 teacher). MOPD is gym-only, so this drives the +# nemo_gym entrypoint. +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=15 +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +# Run the experiment +cd $PROJECT_ROOT +uv run examples/nemo_gym/run_grpo_nemo_gym.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 + +# Convert tensorboard logs to json +uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS + +# Student == teacher, so the OPD distillation signal is ~0: the policy loss +# should sit near 0 and the train-to-inference probability error near 1.0. +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' + + # Clean up checkpoint directory after a successful run to save space. + rm -rf "$CKPT_DIR" +fi diff --git a/tests/test_suites/nightly.txt b/tests/test_suites/nightly.txt index 57b48dca6c..6e278f9adb 100644 --- a/tests/test_suites/nightly.txt +++ b/tests/test_suites/nightly.txt @@ -225,6 +225,12 @@ tests/test_suites/llm/distillation-qwen3-32b-to-1.7b-base-1n8g-megatron-tp2pp2cp # healthy. Self-distill Qwen3-1.7B-Base, ~20 min on 1n8g. tests/test_suites/llm/distillation-qwen3-1.7b-1n8g-megatron-qa-nvfp4.sh +# MOPD (on-policy distillation via nemo_gym, non-colocated teacher). Dense +# Qwen3-1.7B self-distill: student == teacher so the OPD loss stays ~0. +# Sequence packing on, 3 nodes (1 policy + 1 vLLM + 1 teacher). Requires +# NRL_TRAIN_PATH / NRL_VAL_PATH (nemo_gym jsonl) in the environment. +tests/test_suites/llm/mopd-qwen3-1.7b-3n8g-megatron-pack.sh + # Nano3 hybrid MoE/Mamba ModelOpt layer-spec smoke. Keeps # policy.disable_modelopt_layer_spec=false to cover modelopt_mamba_stack_spec. tests/test_suites/llm/distillation-nano3-30ba3b-4n4g-megatron-qa-nvfp4-modelopt-spec.sh diff --git a/tests/unit/algorithms/test_advantage_estimator.py b/tests/unit/algorithms/test_advantage_estimator.py new file mode 100644 index 0000000000..e7cf031bbd --- /dev/null +++ b/tests/unit/algorithms/test_advantage_estimator.py @@ -0,0 +1,109 @@ +# Copyright (c) 2025, 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. + +import torch + +from nemo_rl.algorithms.advantage_estimator import OPDAdvantageEstimator + + +def _make_estimator(): + return OPDAdvantageEstimator({"name": "opd"}, {}) + + +def test_opd_basic_positive_distill_advantage(): + """teacher_lp > student_lp => positive advantages.""" + estimator = _make_estimator() + B, S = 2, 4 + teacher_lp = torch.zeros(B, S) # log(1) = 0 + student_lp = torch.full((B, S), -1.0) # lower logprob + mask = torch.ones(B, S) + prompt_ids = torch.arange(B) + rewards = torch.zeros(B) + + adv = estimator.compute_advantage( + prompt_ids, rewards, mask, teacher_logprobs=teacher_lp, prev_logprobs=student_lp + ) + + assert adv.shape == (B, S) + assert (adv > 0).all(), "teacher_lp > student_lp should yield positive advantages" + + +def test_opd_teacher_equals_student(): + """Same logprobs => zero advantages.""" + estimator = _make_estimator() + B, S = 2, 4 + logprobs = torch.randn(B, S) + mask = torch.ones(B, S) + prompt_ids = torch.arange(B) + rewards = torch.zeros(B) + + adv = estimator.compute_advantage( + prompt_ids, rewards, mask, teacher_logprobs=logprobs, prev_logprobs=logprobs + ) + + torch.testing.assert_close(adv, torch.zeros(B, S)) + + +def test_opd_mask_applied(): + """Masked tokens should have zero advantage.""" + estimator = _make_estimator() + B, S = 1, 6 + teacher_lp = torch.zeros(B, S) + student_lp = torch.full((B, S), -1.0) + mask = torch.tensor([[1, 1, 1, 0, 0, 0]], dtype=torch.float32) + prompt_ids = torch.arange(B) + rewards = torch.zeros(B) + + adv = estimator.compute_advantage( + prompt_ids, rewards, mask, teacher_logprobs=teacher_lp, prev_logprobs=student_lp + ) + + # Masked positions must be zero + assert (adv[:, 3:] == 0).all(), "Masked positions should be zero" + # Unmasked positions should be positive (teacher > student) + assert (adv[:, :3] > 0).all(), "Unmasked positions should be positive" + + +def test_opd_metrics_returned(): + """self.last_metrics should be populated after compute_advantage.""" + estimator = _make_estimator() + B, S = 2, 4 + teacher_lp = torch.zeros(B, S) + student_lp = torch.full((B, S), -1.0) + mask = torch.ones(B, S) + prompt_ids = torch.arange(B) + rewards = torch.zeros(B) + + estimator.compute_advantage( + prompt_ids, rewards, mask, teacher_logprobs=teacher_lp, prev_logprobs=student_lp + ) + + assert ( + "on_policy_distillation/teacher_student_logprob_gap_mean" + in estimator.last_metrics + ) + assert "on_policy_distillation/adv_mean" in estimator.last_metrics + assert "on_policy_distillation/adv_std" in estimator.last_metrics + # teacher - student = 0 - (-1) = 1.0 + assert ( + abs( + estimator.last_metrics[ + "on_policy_distillation/teacher_student_logprob_gap_mean" + ] + - 1.0 + ) + < 1e-5 + ) + assert abs(estimator.last_metrics["on_policy_distillation/adv_mean"] - 1.0) < 1e-5 + assert abs(estimator.last_metrics["on_policy_distillation/adv_std"]) < 1e-5 diff --git a/tests/unit/algorithms/test_async_utils.py b/tests/unit/algorithms/test_async_utils.py index cb33222a54..b7cff7eecd 100644 --- a/tests/unit/algorithms/test_async_utils.py +++ b/tests/unit/algorithms/test_async_utils.py @@ -1105,7 +1105,10 @@ def create_mock_config(self) -> MasterConfig: "max_rollout_turns": 1, "async_grpo": {"max_trajectory_age_steps": 2}, }, - "policy": {"max_total_sequence_length": 512}, + "policy": { + "max_total_sequence_length": 512, + "make_sequence_length_divisible_by": 1, + }, } return MasterConfig.model_construct(**config) @@ -1447,7 +1450,10 @@ def create_mock_config(self) -> MasterConfig: "max_rollout_turns": 1, "async_grpo": {"max_trajectory_age_steps": 1}, }, - "policy": {"max_total_sequence_length": 512}, + "policy": { + "max_total_sequence_length": 512, + "make_sequence_length_divisible_by": 1, + }, } return MasterConfig.model_construct(**config) diff --git a/tests/unit/algorithms/test_opd.py b/tests/unit/algorithms/test_opd.py new file mode 100644 index 0000000000..bb9d89f3e3 --- /dev/null +++ b/tests/unit/algorithms/test_opd.py @@ -0,0 +1,434 @@ +# Copyright (c) 2025, 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. + +import pytest +import torch + +from nemo_rl.distributed.batched_data_dict import BatchedDataDict + +# --------------------------------------------------------------------------- +# Mock teacher worker group for _compute_teacher_logprobs tests +# --------------------------------------------------------------------------- + + +class _MockShardingAnnotations: + def __init__(self, dp_size): + self._dp_size = dp_size + + def get_axis_size(self, name): + if name == "data_parallel": + return self._dp_size + return 1 + + +class _MockTeacherWorkerGroup: + """Returns logprobs filled with a constant; validates DP-divisible batch.""" + + def __init__(self, fill_value=1.0, dp_size=4): + self._fill_value = fill_value + self.sharding_annotations = _MockShardingAnnotations(dp_size) + + def get_logprobs(self, data): + input_ids = data["input_ids"] + B, S = input_ids.shape + # Verify the caller already padded to dp_size + dp_size = self.sharding_annotations.get_axis_size("data_parallel") + assert B % dp_size == 0, ( + f"get_logprobs received batch_size={B} not divisible by dp_size={dp_size}" + ) + return BatchedDataDict( + {"reference_logprobs": torch.full((B, S), self._fill_value)} + ) + + +def _make_collector(**overrides): + """Build a bare AsyncTrajectoryCollector (bypass Ray) for unit testing.""" + import threading + + from nemo_rl.algorithms.async_utils import AsyncTrajectoryCollector + + # AsyncTrajectoryCollector is @ray.remote-decorated; unwrap to the real class. + real_cls = AsyncTrajectoryCollector.__ray_metadata__.modified_class + defaults = { + "teacher_worker_groups": {}, + "alias_to_group_alias": {}, + "on_policy_distillation_cfg": {}, + "_has_distillation_teachers": False, + } + defaults.update(overrides) + obj = object.__new__(real_cls) + for k, v in defaults.items(): + setattr(obj, k, v) + obj._teacher_locks = {k: threading.Lock() for k in obj.teacher_worker_groups} + return obj + + +# --------------------------------------------------------------------------- +# DP padding tests +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "batch_size,dp_size", + [ + (1, 4), # the exact bug: 1 sample, dp=4 + (2, 4), # 2 samples, dp=4 + (3, 4), # 3 samples, dp=4 + (4, 4), # already aligned + (1, 8), # extreme: 1 sample, dp=8 + (5, 4), # 5 samples → pad to 8 + ], +) +def test_compute_teacher_logprobs_dp_padding(batch_size, dp_size): + """Teacher logprob computation must pad batch to dp_size multiple.""" + twg = _MockTeacherWorkerGroup(fill_value=2.0, dp_size=dp_size) + collector = _make_collector( + teacher_worker_groups={"math": twg}, + alias_to_group_alias={"math_agent": "math"}, + on_policy_distillation_cfg={ + "teacher_model_by_agent_name": {"math_agent": "/ckpt/math"}, + }, + _has_distillation_teachers=True, + ) + + S = 16 + input_ids = torch.randint(0, 100, (batch_size, S)) + agent_refs = [{"name": "math_agent"}] * batch_size + + result, _ = collector._compute_teacher_logprobs(input_ids, agent_refs) + + assert result.shape == (batch_size, S) + assert torch.allclose(result, torch.tensor(2.0)) + + +def test_compute_teacher_logprobs_routes_to_correct_teacher(): + """Samples are routed to the right teacher and results stitched back.""" + math_twg = _MockTeacherWorkerGroup(fill_value=1.0, dp_size=1) + code_twg = _MockTeacherWorkerGroup(fill_value=2.0, dp_size=1) + + collector = _make_collector( + teacher_worker_groups={"math": math_twg, "code": code_twg}, + alias_to_group_alias={"math_agent": "math", "code_agent": "code"}, + on_policy_distillation_cfg={ + "teacher_model_by_agent_name": { + "math_agent": "/ckpt/math", + "code_agent": "/ckpt/code", + }, + }, + _has_distillation_teachers=True, + ) + + B, S = 4, 8 + input_ids = torch.randint(0, 100, (B, S)) + agent_refs = [ + {"name": "math_agent"}, + {"name": "code_agent"}, + {"name": "math_agent"}, + {"name": "code_agent"}, + ] + + result, _ = collector._compute_teacher_logprobs(input_ids, agent_refs) + + assert result.shape == (B, S) + assert torch.allclose(result[0], torch.tensor(1.0)) + assert torch.allclose(result[1], torch.tensor(2.0)) + assert torch.allclose(result[2], torch.tensor(1.0)) + assert torch.allclose(result[3], torch.tensor(2.0)) + + +def test_compute_teacher_logprobs_deduplication(): + """alias_to_group_alias routes multiple aliases to one teacher group.""" + shared_twg = _MockTeacherWorkerGroup(fill_value=3.0, dp_size=1) + + collector = _make_collector( + teacher_worker_groups={"primary": shared_twg}, + alias_to_group_alias={"mcqa": "primary", "terminal": "primary"}, + on_policy_distillation_cfg={ + "teacher_model_by_agent_name": { + "mcqa": "/ckpt/shared", + "terminal": "/ckpt/shared", + }, + }, + _has_distillation_teachers=True, + ) + + B, S = 2, 4 + input_ids = torch.randint(0, 100, (B, S)) + agent_refs = [{"name": "mcqa"}, {"name": "terminal"}] + + result, _ = collector._compute_teacher_logprobs(input_ids, agent_refs) + assert result.shape == (B, S) + assert torch.allclose(result, torch.tensor(3.0)) + + +def test_compute_teacher_logprobs_default_alias_fallback_routes(): + """Unmapped agent_ref falls back to default_teacher_alias and routes to a valid group.""" + math_twg = _MockTeacherWorkerGroup(fill_value=7.0, dp_size=1) + collector = _make_collector( + teacher_worker_groups={"math": math_twg}, + alias_to_group_alias={"math_agent": "math"}, + on_policy_distillation_cfg={ + "teacher_model_by_agent_name": {"math_agent": "/ckpt/math"}, + "default_teacher_alias": "math_agent", + }, + _has_distillation_teachers=True, + ) + B, S = 2, 5 + input_ids = torch.randint(0, 100, (B, S)) + # second agent ("surprise_agent") is unmapped -> must fall back to math_agent + agent_refs = [{"name": "math_agent"}, {"name": "surprise_agent"}] + result, _ = collector._compute_teacher_logprobs(input_ids, agent_refs) + assert result.shape == (B, S) + assert torch.allclose(result, torch.tensor(7.0)) + + +# --------------------------------------------------------------------------- +# Unsort / reorder_data regression test +# --------------------------------------------------------------------------- + + +def test_reorder_data_vs_direct_gather(): + """Verify reorder_data inverts the permutation, while direct gather does not. + + This is the root cause of the num_gen>1 teacher logprob misalignment bug: + shard_by_batch_size returns a forward permutation (sorted_pos → orig_idx). + To restore original order we need the *inverse* (argsort), which + reorder_data computes. A direct gather ``result[indices]`` applies + the forward permutation and silently produces wrong results. + """ + # Simulate: 4 samples reordered by sequence packing as [3, 0, 2, 1] + forward_perm = [3, 0, 2, 1] + # After inference, results are in sorted order: + # position 0 = result for orig sample 3 + # position 1 = result for orig sample 0 etc. + sorted_results = BatchedDataDict( + {"logprobs": torch.tensor([[30.0], [0.0], [20.0], [10.0]])} + ) + # label: sorted_results[i] holds the value for original sample forward_perm[i] + # sorted_results[0]=30 → orig 3, sorted_results[1]=0 → orig 0, etc. + + # --- WRONG: direct gather (the old bug) --- + wrong = sorted_results["logprobs"][forward_perm] + # wrong[0] = sorted_results[3] = 10 (should be 0 for orig 0) + assert not torch.equal(wrong, torch.tensor([[0.0], [10.0], [20.0], [30.0]])), ( + "Direct gather should NOT produce the correct original order" + ) + + # --- CORRECT: reorder_data (inverse permutation) --- + correct = BatchedDataDict({"logprobs": sorted_results["logprobs"].clone()}) + correct.reorder_data(forward_perm) + assert torch.equal( + correct["logprobs"], torch.tensor([[0.0], [10.0], [20.0], [30.0]]) + ), "reorder_data should restore the original sample order" + + +# --------------------------------------------------------------------------- +# Teacher logprob alignment with variable-length sequences (num_gen > 1) +# --------------------------------------------------------------------------- + + +def test_reorder_data_inverse_permutation_various(): + """reorder_data correctly inverts arbitrary permutations, including identity.""" + # Identity permutation + bdd = BatchedDataDict({"x": torch.tensor([[0.0], [1.0], [2.0]])}) + bdd.reorder_data([0, 1, 2]) + assert torch.equal(bdd["x"], torch.tensor([[0.0], [1.0], [2.0]])) + + # Reversal + bdd = BatchedDataDict({"x": torch.tensor([[0.0], [1.0], [2.0]])}) + bdd.reorder_data([2, 1, 0]) + # batch_sorted_indices=[2,1,0] means sorted[0] came from orig 2, etc. + # Inverse: orig[2]=sorted[0]=0.0, orig[1]=sorted[1]=1.0, orig[0]=sorted[2]=2.0 + assert torch.equal(bdd["x"], torch.tensor([[2.0], [1.0], [0.0]])) + + # Non-trivial: simulate 4 samples reordered as [2, 3, 0, 1] + bdd = BatchedDataDict({"x": torch.tensor([[20.0], [30.0], [0.0], [10.0]])}) + bdd.reorder_data([2, 3, 0, 1]) + assert torch.equal(bdd["x"], torch.tensor([[0.0], [10.0], [20.0], [30.0]])), ( + "After reorder_data, row i should hold the result for original sample i" + ) + + +def test_is_opd_enabled(): + from nemo_rl.algorithms.opd import is_opd_enabled + + assert is_opd_enabled({"on_policy_distillation": {"enabled": True}}) + assert not is_opd_enabled({"on_policy_distillation": {"enabled": False}}) + assert not is_opd_enabled({}) + + +def test_is_opd_enabled_object_config(): + # _opd_cfg must also handle a config object (not just a dict): math recipes + # have no on_policy_distillation attribute at all. + from types import SimpleNamespace + + from nemo_rl.algorithms.opd import is_opd_enabled + + assert is_opd_enabled(SimpleNamespace(on_policy_distillation={"enabled": True})) + assert not is_opd_enabled( + SimpleNamespace(on_policy_distillation={"enabled": False}) + ) + assert not is_opd_enabled(SimpleNamespace()) + assert not is_opd_enabled(SimpleNamespace(on_policy_distillation=None)) + + +def test_is_non_colocated_teachers_enabled(): + from nemo_rl.algorithms.opd import is_non_colocated_teachers_enabled + + assert is_non_colocated_teachers_enabled( + { + "on_policy_distillation": { + "enabled": True, + "non_colocated_teachers": {"enabled": True}, + } + } + ) + assert not is_non_colocated_teachers_enabled( + { + "on_policy_distillation": { + "enabled": True, + "non_colocated_teachers": {"enabled": False}, + } + } + ) + + +def test_resolve_reference_aliases_bad_agent_ref(): + from nemo_rl.algorithms.opd import resolve_reference_aliases + + with pytest.raises(KeyError): + resolve_reference_aliases([{"not_name": "oops"}], {"math": "/ckpt/math"}) + + +def test_resolve_reference_aliases_fallback(): + from nemo_rl.algorithms.opd import resolve_reference_aliases + + aliases = resolve_reference_aliases( + [{"name": "math_agent"}, {"name": "unknown"}, {"name": "code_agent"}], + {"math_agent": "/ckpt/math", "code_agent": "/ckpt/code"}, + default_teacher_alias="math_agent", + ) + assert aliases == ["math_agent", "math_agent", "code_agent"] + + +def test_resolve_reference_aliases_strict_raises(): + from nemo_rl.algorithms.opd import resolve_reference_aliases + + with pytest.raises(ValueError, match="No teacher model mapping"): + resolve_reference_aliases( + [{"name": "unknown"}], {"math": "/ckpt/math"}, strict_agent_name_match=True + ) + + +def test_get_teacher_routing_metrics(): + from nemo_rl.algorithms.opd import get_teacher_routing_metrics + + metrics = get_teacher_routing_metrics( + ["math_a", "math_b", "if", "math_a"], + {"math_a": "t_math", "math_b": "t_math", "if": "t_if"}, + ) + assert metrics["on_policy_distillation/teacher_alias_unique"] == 3.0 + assert metrics["on_policy_distillation/teacher_model_unique"] == 2.0 + + +# --------------------------------------------------------------------------- +# teacher_seq_pad_multiple: teacher pre-pad multiple + packing-mode guard +# --------------------------------------------------------------------------- + + +def _twg(packed, pad_multiple=1): + """TeacherWorkerGroup stand-in with the two attrs the helper reads.""" + from types import SimpleNamespace + + return SimpleNamespace( + use_sequence_packing=packed, sequence_length_pad_multiple=pad_multiple + ) + + +def test_teacher_seq_pad_multiple_no_teachers_is_one(): + from nemo_rl.algorithms.opd import teacher_seq_pad_multiple + + assert teacher_seq_pad_multiple({}, 8) == 1 + + +def test_teacher_seq_pad_multiple_all_packed_is_one(): + from nemo_rl.algorithms.opd import teacher_seq_pad_multiple + + assert teacher_seq_pad_multiple({"a": _twg(True), "b": _twg(True)}, 8) == 1 + + +def test_teacher_seq_pad_multiple_mixed_packing_raises(): + from nemo_rl.algorithms.opd import teacher_seq_pad_multiple + + with pytest.raises(ValueError, match="same sequence-packing mode"): + teacher_seq_pad_multiple({"a": _twg(True), "b": _twg(False)}, 8) + + +def test_teacher_seq_pad_multiple_non_packed_uses_policy_divisor(): + # Non-packed teachers pre-pad to the policy divisor when it is a multiple of + # every teacher's requirement (here 2 and 4 both divide 8). + from nemo_rl.algorithms.opd import teacher_seq_pad_multiple + + teachers = {"a": _twg(False, pad_multiple=2), "b": _twg(False, pad_multiple=4)} + assert teacher_seq_pad_multiple(teachers, 8) == 8 + + +def test_teacher_seq_pad_multiple_non_packed_incompatible_divisor_raises(): + # policy divisor 8 is not a multiple of the teacher requirement 16. + from nemo_rl.algorithms.opd import teacher_seq_pad_multiple + + with pytest.raises(ValueError, match="make_sequence_length_divisible_by"): + teacher_seq_pad_multiple({"a": _twg(False, pad_multiple=16)}, 8) + + +# --------------------------------------------------------------------------- +# Teacher-logprob seq-length padding + the "opd" advantage-estimator branch +# --------------------------------------------------------------------------- + + +def test_pad_teacher_logprobs(): + from nemo_rl.algorithms.grpo import _pad_teacher_logprobs + + # teacher_S < train_S -> right zero-pad + padded = _pad_teacher_logprobs(torch.ones(2, 3), 5) + assert padded.shape == (2, 5) + assert (padded[:, :3] == 1).all() and (padded[:, 3:] == 0).all() + # teacher_S == train_S -> unchanged + assert _pad_teacher_logprobs(torch.ones(2, 4), 4).shape == (2, 4) + # teacher_S > train_S -> raises + with pytest.raises(ValueError, match="seq length"): + _pad_teacher_logprobs(torch.ones(2, 6), 4) + + +def test_create_advantage_estimator_opd_branch(): + import warnings + from types import SimpleNamespace + + from nemo_rl.algorithms.advantage_estimator import OPDAdvantageEstimator + from nemo_rl.algorithms.grpo import _create_advantage_estimator + + # loss_fn not MOPD-configured -> the 3 recommendation warnings fire. + loss_fn = SimpleNamespace( + disable_ppo_ratio=False, + use_importance_sampling_correction=False, + truncated_importance_sampling_type="none", + ) + master_config = SimpleNamespace( + grpo={"adv_estimator": {"name": "opd"}}, loss_fn=loss_fn + ) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + estimator = _create_advantage_estimator(master_config) + assert isinstance(estimator, OPDAdvantageEstimator) + assert len(caught) == 3 diff --git a/tests/unit/models/policy/test_teacher_worker_group.py b/tests/unit/models/policy/test_teacher_worker_group.py new file mode 100644 index 0000000000..94c25c5f8e --- /dev/null +++ b/tests/unit/models/policy/test_teacher_worker_group.py @@ -0,0 +1,79 @@ +# Copyright (c) 2025, 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. + + +def test_teacher_resource_config_defaults(): + from nemo_rl.algorithms.opd import TeacherResourceConfig + + res = TeacherResourceConfig(tensor_model_parallel_size=4) + assert res.tensor_model_parallel_size == 4 + assert res.pipeline_model_parallel_size == 1 + assert res.gpus_per_node == 8 + assert res.precision == "bf16" + + +def test_create_teacher_configs_homogeneous(): + from nemo_rl.models.policy.teacher_worker_group import ( + create_teacher_configs_from_opd_config, + ) + + configs = create_teacher_configs_from_opd_config( + { + "teacher_model_by_agent_name": {"math": "/ckpt/math", "code": "/ckpt/code"}, + "non_colocated_teachers": { + "default_teacher_cfg": {"tensor_model_parallel_size": 4} + }, + } + ) + assert len(configs) == 2 + assert all(c.tensor_model_parallel_size == 4 for c in configs) + + +def test_create_teacher_configs_heterogeneous_override(): + from nemo_rl.models.policy.teacher_worker_group import ( + create_teacher_configs_from_opd_config, + ) + + configs = create_teacher_configs_from_opd_config( + { + "teacher_model_by_agent_name": {"math": "/ckpt/math", "code": "/ckpt/code"}, + "non_colocated_teachers": { + "default_teacher_cfg": {"tensor_model_parallel_size": 4}, + "teacher_overrides": {"code": {"tensor_model_parallel_size": 8}}, + }, + } + ) + code_cfg = [c for c in configs if c.alias == "code"][0] + assert code_cfg.tensor_model_parallel_size == 8 + + +def test_create_teacher_configs_deduplicates(): + from nemo_rl.models.policy.teacher_worker_group import ( + create_teacher_configs_from_opd_config, + ) + + configs = create_teacher_configs_from_opd_config( + { + "teacher_model_by_agent_name": { + "math": "/shared", + "code": "/shared", + "rlhf": "/rlhf", + }, + "deduplicate_shared_teacher_checkpoints": True, + "non_colocated_teachers": { + "default_teacher_cfg": {"tensor_model_parallel_size": 2} + }, + } + ) + assert len(configs) == 2 diff --git a/tests/unit/reference_configs/grpo_math_1B.yaml b/tests/unit/reference_configs/grpo_math_1B.yaml index 1931ff578f..fffde35296 100644 --- a/tests/unit/reference_configs/grpo_math_1B.yaml +++ b/tests/unit/reference_configs/grpo_math_1B.yaml @@ -453,3 +453,7 @@ data_plane: local_buffer_size: 68719476736 # 64 GiB — used when backend == "mooncake_cpu" # observability: # NotRequired # enabled: false + +# Multi-Teacher On-Policy Distillation (MOPD); null = disabled (default). Mirrors +# the field on the GRPO MasterConfig added for MOPD support. +on_policy_distillation: null diff --git a/tests/unit/test_recipes_and_test_suites.py b/tests/unit/test_recipes_and_test_suites.py index c5833344a6..3583b7b19b 100644 --- a/tests/unit/test_recipes_and_test_suites.py +++ b/tests/unit/test_recipes_and_test_suites.py @@ -47,6 +47,7 @@ "dapo": "examples/configs/grpo_math_1B.yaml", "prorlv2": "examples/configs/prorlv2.v2.yaml", "ppo": "examples/configs/ppo_math_1B_megatron.yaml", + "mopd": "examples/configs/grpo_math_1B.yaml", } # Configuration keys that are allowed to be added to base configs during testing