nemo_rl.models.policy.teacher_worker_group

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.

Module Contents

Classes

TeacherConfig	Resolved config for a single non-colocated teacher (built in-process).
TeacherWorkerGroup	Inference-only mcore worker group for a single teacher model.

Functions

create_teacher_configs_from_opd_config

Build per-teacher configs from on_policy_distillation config.

API

class nemo_rl.models.policy.teacher_worker_group.TeacherConfig

Resolved config for a single non-colocated teacher (built in-process).

alias: str

None

model_name: str

None

tensor_model_parallel_size: int

None

pipeline_model_parallel_size: int

None

context_parallel_size: int

None

expert_model_parallel_size: int

None

num_nodes: int

None

gpus_per_node: int

None

precision: str

None

micro_batch_size: int

None

megatron_cfg_overrides: dict[str, Any]

None

nemo_rl.models.policy.teacher_worker_group.create_teacher_configs_from_opd_config(

opd_cfg: dict[str, Any],

) → list[nemo_rl.models.policy.teacher_worker_group.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.

class nemo_rl.models.policy.teacher_worker_group.TeacherWorkerGroup(

teacher_cfg: nemo_rl.models.policy.teacher_worker_group.TeacherConfig,

cluster: nemo_rl.distributed.virtual_cluster.RayVirtualCluster,

policy_config: dict[str, Any],

tokenizer: transformers.PreTrainedTokenizerBase,

)

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

Initialization

get_logprobs(

data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationDatumSpec],

micro_batch_size: Optional[int] = None,

) → nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.policy.interfaces.ReferenceLogprobOutputSpec]

Run forward pass on teacher and return logprobs.

shutdown() → bool

Shut down all workers and clean up resources.

__del__() → None

Safety net for cleanup.