`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.
`TQTeacherLogprobCoordinator`	Enrich SingleController rollout rows with teacher logprobs through TQ.

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.
`_validate_default_teacher_alias`	Validate the fallback teacher alias before reserving resources.
`reserve_teacher_clusters`	Create and reserve topology-aware clusters for non-colocated teachers.
`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).

class nemo_rl.algorithms.opd.TQTeacherLogprobCoordinator( *, dp_client: nemo_rl.data_plane.interfaces.DataPlaneClient, teacher_worker_groups: dict[str, Any], alias_to_group_alias: dict[str, str], on_policy_distillation_cfg: dict[str, Any], )#

Enrich SingleController rollout rows with teacher logprobs through TQ.

A prompt group is already present in the data plane when :meth:enrich is called. The coordinator sends its metadata to the inference-only teacher; teacher workers fetch their own DP shards and write the token column back under the same sample IDs. The controller materializes only the final row when unique temporary rows are needed for DP divisibility. Calls targeting the same deduplicated teacher group are serialized to keep NCCL collective ordering identical across that group’s workers; distinct teachers remain independent and can run concurrently.

Initialization

teacher_logprobs_field#: ‘teacher_reference_logprobs’

_resolve_teacher( record: nemo_rl.experience.interfaces.PromptGroupRecord, ) → tuple[str, str]#

_enrich_sync( meta: nemo_rl.data_plane.interfaces.KVBatchMeta, group_alias: str, ) → float#: Run the blocking TQ-read, teacher inference, and TQ-write sequence.

async enrich( meta: nemo_rl.data_plane.interfaces.KVBatchMeta, record: nemo_rl.experience.interfaces.PromptGroupRecord, ) → nemo_rl.data_plane.interfaces.KVBatchMeta#: Write teacher logprobs before the replay-buffer slot becomes ready.

drain_metrics() → dict[str, float]#: Return and reset teacher activity accumulated since the last drain.

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._validate_default_teacher_alias( opd_cfg: dict[str, Any], ) → None#: Validate the fallback teacher alias before reserving resources.

nemo_rl.algorithms.opd.reserve_teacher_clusters( master_config: Any, *, segment_size: Optional[int] = None, teacher_segment_topology: Optional[dict[str, tuple[str, int]]] = None, ) → dict[str, nemo_rl.distributed.virtual_cluster.RayVirtualCluster]#

Create and reserve topology-aware clusters for non-colocated teachers.

This reserves the teachers’ Ray placement groups without starting teacher workers or loading model checkpoints. Call it before starting other opportunistically placed GPU services, then pass the result to

Func:

create_teacher_worker_groups after policy initialization.

Parameters:

master_config – Full training configuration containing the OPD settings.
segment_size – NVLink-domain segment size from the cluster config. When set, every teacher is constrained to one NVLink domain.
teacher_segment_topology – Topology remaining after policy and inference placement.

Returns:

A mapping from each deduplicated teacher alias to its reserved cluster.

Raises:

ValueError – If the configured fallback teacher alias is invalid.
ResourceInsufficientError – If the requested topology segments cannot be formed.
TimeoutError – If Ray cannot reserve a teacher placement group.

nemo_rl.algorithms.opd.create_teacher_worker_groups( master_config: Any, policy_config: dict[str, Any], tokenizer: Any, *, teacher_clusters: dict[str, nemo_rl.distributed.virtual_cluster.RayVirtualCluster], ) → tuple[dict[str, Any], dict[str, str]]#

Create TeacherWorkerGroup instances for non-colocated teachers.

Parameters:

master_config – Full training configuration containing the OPD settings.
policy_config – Student policy configuration used as the teacher worker configuration template.
tokenizer – Tokenizer passed to every teacher worker.
teacher_clusters –
Clusters already reserved by

func:

reserve_teacher_clusters, keyed by teacher alias.

Returns:

A tuple containing the worker groups by primary teacher alias and the mapping from every configured alias to its primary group alias.

Raises:

ValueError – If teacher routing or the supplied cluster aliases are invalid, or teacher sequence-packing settings are incompatible.
RuntimeError – If any teacher worker fails during initialization.

nemo_rl.algorithms.opd#

Module Contents#

Classes#

Functions#

API#

`nemo_rl.algorithms.opd`#