`nemo_rl.algorithms.opd`#

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).

Module Contents#

Classes#

`TeacherResourceConfig`	Per-teacher resourcing for a non-colocated teacher worker group.
`NonColocatedTeachersConfig`	Non-colocated (separate-GPU) teacher resourcing for on-policy distillation.
`OnPolicyDistillationConfig`	User-facing config for the top-level `on_policy_distillation` block.

Functions#

`_opd_cfg`	Return the on_policy_distillation sub-config as a plain dict.
`is_opd_enabled`	Whether on-policy distillation is enabled in the config.
`is_non_colocated_teachers_enabled`	Whether OPD is enabled with non-colocated (separate-GPU) teachers.
`_skip_prev_logprobs`	Whether the training loop will zero `prev_logprobs` instead of computing it.
`assert_prev_logprobs_available`	Raise if OPD is enabled but the config would zero `prev_logprobs`.
`resolve_reference_aliases`	Map each agent_ref to a teacher alias.
`get_teacher_routing_metrics`	Compute teacher-routing diagnostics.
`teacher_seq_pad_multiple`	Sequence divisor to pre-pad teacher logprob inputs to.
`create_teacher_worker_groups`	Create TeacherWorkerGroup instances for non-colocated teachers.

API#

class nemo_rl.algorithms.opd.TeacherResourceConfig#

Bases: pydantic.BaseModel

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(…)’

class nemo_rl.algorithms.opd.NonColocatedTeachersConfig#

Bases: pydantic.BaseModel

Non-colocated (separate-GPU) teacher resourcing for on-policy distillation.

enabled: bool#: False

default_teacher_cfg: nemo_rl.algorithms.opd.TeacherResourceConfig#: ‘Field(…)’

teacher_overrides: dict[str, nemo_rl.algorithms.opd.TeacherResourceConfig]#: ‘Field(…)’

class nemo_rl.algorithms.opd.OnPolicyDistillationConfig#

Bases: pydantic.BaseModel

User-facing config for the top-level on_policy_distillation block.

enabled: bool#: False

teacher_model_by_agent_name: dict[str, str]#: ‘Field(…)’

default_teacher_alias: Optional[str]#: None

strict_agent_name_match: bool#: False

deduplicate_shared_teacher_checkpoints: bool#: True

non_colocated_teachers: Optional[nemo_rl.algorithms.opd.NonColocatedTeachersConfig]#: None

nemo_rl.algorithms.opd._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.

nemo_rl.algorithms.opd.is_opd_enabled(master_config: Any) → bool#: Whether on-policy distillation is enabled in the config.

nemo_rl.algorithms.opd.is_non_colocated_teachers_enabled(master_config: Any) → bool#: Whether OPD is enabled with non-colocated (separate-GPU) teachers.

nemo_rl.algorithms.opd._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.

nemo_rl.algorithms.opd.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.

nemo_rl.algorithms.opd.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.

nemo_rl.algorithms.opd.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).

nemo_rl.algorithms.opd.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.

nemo_rl.algorithms.opd.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.

Parameters:

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).

nemo_rl.algorithms.opd#

Module Contents#

Classes#

Functions#

API#

`nemo_rl.algorithms.opd`#