`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.
`OnPolicyDistillationFullConfig`	Full-vocabulary MOPD (`opd_full`).
`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.
`get_opd_full_config`	Return the validated `opd_full` sub-config when the feature is enabled.
`opd_full_payload_field`	Return the data-plane column name carrying this run’s teacher payload.
`opd_full_teacher_index_field`	Return the per-sample teacher-identity column, or `None` if unneeded.
`teacher_configs_by_index`	Deduplicated teacher configs in `teacher_index` order.
`opd_full_teacher_checkpoints_by_index`	Checkpoint path per `teacher_index`, or None when rows carry no tag.
`_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`.
`assert_trust_region_supported`	Raise if TROPD knobs are set on a path that never reads the advantage.
`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.OnPolicyDistillationFullConfig#

Bases: pydantic.BaseModel

Full-vocabulary MOPD (opd_full).

The student matches the teacher’s entire next-token distribution through an exact reverse KL instead of the sampled-token log-probability gap – the K=V limit of the top-k estimator, with no score-function tail term.

The policy-gradient knobs on ClippedPGLossConfig (ratio clipping, CISPO, importance-sampling correction, truncated IS) have no code path here and are rejected at construction. The OPD advantage estimator still runs: its output is unused, but advantages is a required training column.

enabled: bool#: False

teacher_payload: Literal[hidden_states, logits]#: ‘hidden_states’

divergence: Literal[reverse_kl]#: ‘reverse_kl’

payload_dtype: Literal[bfloat16, float16, float32]#: ‘bfloat16’

chunk_size: Optional[int]#: None

teacher_lm_head_lifecycle: Literal[none, offload, evict]#: ‘offload’

validate_decomposition: bool#: False

validate_positive_sizes() → nemo_rl.algorithms.opd.OnPolicyDistillationFullConfig#: Reject a chunk size that would make the divergence kernels useless.

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

full: Optional[nemo_rl.algorithms.opd.OnPolicyDistillationFullConfig]#: 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.get_opd_full_config( master_config: Any, ) → Optional[nemo_rl.algorithms.opd.OnPolicyDistillationFullConfig]#

Return the validated opd_full sub-config when the feature is enabled.

Parameters:: master_config – Full training configuration, a plain dict, or any object without the field (non-OPD recipes).
Returns:: The full-vocabulary config when both OPD and full.enabled are on, otherwise None.

nemo_rl.algorithms.opd.opd_full_payload_field( full_cfg: nemo_rl.algorithms.opd.OnPolicyDistillationFullConfig, ) → str#: Return the data-plane column name carrying this run’s teacher payload.

nemo_rl.algorithms.opd.opd_full_teacher_index_field( full_cfg: nemo_rl.algorithms.opd.OnPolicyDistillationFullConfig, ) → Optional[str]#

Return the per-sample teacher-identity column, or None if unneeded.

Only the hidden_states payload needs it: the student loads a teacher LM-head shard per unique checkpoint and must know which one projects each row. The logits payload ships an already-projected distribution, so no per-sample routing is needed at training time.

nemo_rl.algorithms.opd.teacher_configs_by_index( teacher_configs: collections.abc.Iterable[nemo_rl.models.policy.teacher_worker_group.TeacherConfig], ) → list[nemo_rl.models.policy.teacher_worker_group.TeacherConfig]#: Deduplicated teacher configs in teacher_index order.

nemo_rl.algorithms.opd.opd_full_teacher_checkpoints_by_index( master_config: Any, ) → Optional[list[str]]#

Checkpoint path per teacher_index, or None when rows carry no tag.

A list, not a map: the metadata round trips through JSON.

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. GRPO can also evaluate that threshold in the training loss without this 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.assert_trust_region_supported(master_config: Any) → None#

Raise if TROPD knobs are set on a path that never reads the advantage.

Full-vocabulary MOPD replaces the policy-gradient objective with an exact reverse KL and ignores advantages, so a proximal teacher or a global baseline would be silently dropped there.

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.

teacher_checkpoints_by_index() → Optional[list[str]]#

Checkpoint path per teacher_index, or None when rows carry no tag.

Read off the live worker groups, so it records what actually tagged the rows rather than what the config would produce now.

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`#