`nemo_rl.algorithms.xtoken_off_policy_distillation`#

Multi-teacher cross-tokenizer off-policy distillation.

Training-loop layout mirrors run_distillation.py / nemo_rl/algorithms/distillation.py minus the on-policy bits (no env, no rollout, no generation). Per step:

1. Pull a collated batch (student & teacher token ids + alignment).
2. Run each teacher forward via ``Policy.get_full_logits_ipc`` on TEACHER
   token ids — every teacher ships full-vocab logits over CUDA IPC (no
   driver round-trip), reassembled across the teacher's TP/CP shards on
   the student side.
3. Pack alignment payload + teacher IPC handles into a student-side
   ``train_data`` dict.
4. ``student_policy.train(train_data, loss_fn)`` — student forward +
   loss + backward + optimizer step happens inside the dtensor v2
   worker.

The collator and aligner do all the CPU-side cross-tokenizer work; the loss function does only loss math; this module is just plumbing.

Module Contents#

Classes#

`OffPolicyDistillationConfig`	Top-level distillation algo config.
`OffPolicyDistillationSaveState`
`TeacherConfig`	Per-teacher config for multi-teacher cross-tokenizer distillation.
`MasterConfig`

Functions#

`xtoken_non_student_seq_keys`	Build the set of `train_data` keys whose dim 1 is NOT the student seq axis.
`_default_off_policy_distillation_save_state`
`setup`	Construct cluster, dataloaders, policies, and loss fn for the run.
`export_teacher_logits_and_pack`	Serially run each teacher’s forward and pack the student `train_data`.
`xtoken_off_policy_distillation_train`	Off-policy CT distillation training loop.
`validate`	Held-out KL/CE on a validation dataloader.

API#

nemo_rl.algorithms.xtoken_off_policy_distillation.xtoken_non_student_seq_keys( loss_fn: nemo_rl.algorithms.loss.loss_functions.CrossTokenizerDistillationLossFn, ) → frozenset[str]#

Build the set of train_data keys whose dim 1 is NOT the student seq axis.

These ride on the student-side train_data so the loss fn can index them per microbatch, but the worker’s check_sequence_dim pre-flight (which assumes [B, student_seq, ...] for every 2+D tensor) must skip them. The set is teacher-dependent, so it is built from the loss fn’s per-teacher metadata rather than a static constant. Every teacher ships full logits, so each contributes teacher_{i}_full_logits_ipc (a list of CUDA IPC handle dicts, not a tensor); a cross-tokenizer teacher additionally rides its teacher-seq tokenization (teacher_{i}_input_ids / teacher_{i}_token_mask, [B, T_t]) and its teacher-seq / max_pairs alignment_{i}_* keys. The alignment_{i}_student_* ([B, T_s]) and alignment_{i}_num_chunks ([B]) keys follow the student-seq invariant and are NOT skipped.

class nemo_rl.algorithms.xtoken_off_policy_distillation.OffPolicyDistillationConfig#

Bases: typing.TypedDict

Top-level distillation algo config.

.. attribute:: num_prompts_per_step

Global batch size at the dataloader level.

.. attribute:: max_num_steps

Max training steps before early stop.

.. attribute:: max_num_epochs

Max passes over the training dataset.

.. attribute:: seed

RNG seed.

.. attribute:: val_period

Validation cadence in steps. 0 disables validation.

.. attribute:: val_at_start

Run validation before training begins.

.. attribute:: val_at_end

Run validation on the final step.

Initialization

Initialize self. See help(type(self)) for accurate signature.

num_prompts_per_step: int#: None

max_num_steps: int#: None

max_num_epochs: int#: None

seed: int#: None

val_period: int#: None

val_at_start: bool#: None

val_at_end: bool#: None

class nemo_rl.algorithms.xtoken_off_policy_distillation.OffPolicyDistillationSaveState#

Bases: typing.TypedDict

current_epoch: int#: None

current_step: int#: None

total_steps: int#: None

consumed_samples: int#: None

total_valid_tokens: int#: None

val_loss: NotRequired[float]#: None

nemo_rl.algorithms.xtoken_off_policy_distillation._default_off_policy_distillation_save_state() → nemo_rl.algorithms.xtoken_off_policy_distillation.OffPolicyDistillationSaveState#

class nemo_rl.algorithms.xtoken_off_policy_distillation.TeacherConfig#

Bases: pydantic.BaseModel

Per-teacher config for multi-teacher cross-tokenizer distillation.

Carries the full PolicyConfig content (model_name, tokenizer, dtensor_cfg, …) as permitted extras, plus the cross-tokenizer knobs declared below. Use :meth:policy_config to recover the plain PolicyConfig dict for Policy construction.

.. attribute:: projection_matrix_path

Path to this teacher’s student->teacher projection matrix. None marks a same-tokenizer teacher: projection and alignment are skipped and the loss uses a direct per-position KL on the shared vocab.

.. attribute:: weight

Static loss weight for this teacher when several teachers are aggregated (kd_loss_mode="sum" / the convex "averaged_logits" mix). Single-teacher runs leave it at 1.0.

.. attribute:: gold_loss

Optional per-teacher override of loss_fn.gold_loss. None falls back to the global value. Honored only in kd_loss_mode="sum" (other modes use the global).

.. attribute:: xtoken_loss

Optional per-teacher override of loss_fn.xtoken_loss, same semantics as gold_loss.

projection_matrix_path: Optional[str]#: None

weight: float#: 1.0

gold_loss: Optional[bool]#: None

xtoken_loss: Optional[bool]#: None

policy_config() → nemo_rl.models.policy.PolicyConfig#: Recover the plain PolicyConfig dict (cross-tokenizer knobs stripped).

class nemo_rl.algorithms.xtoken_off_policy_distillation.MasterConfig#

Bases: pydantic.BaseModel

policy: nemo_rl.models.policy.PolicyConfig#: None

teachers: list[nemo_rl.algorithms.xtoken_off_policy_distillation.TeacherConfig]#: ‘Field(…)’

loss_fn: nemo_rl.algorithms.loss.loss_functions.CrossTokenizerDistillationLossConfig#: None

data: nemo_rl.data.DataConfig#: None

distillation: nemo_rl.algorithms.xtoken_off_policy_distillation.OffPolicyDistillationConfig#: None

logger: nemo_rl.utils.logger.LoggerConfig#: None

cluster: nemo_rl.distributed.virtual_cluster.ClusterConfig#: None

checkpointing: nemo_rl.utils.checkpoint.CheckpointingConfig#: None

nemo_rl.algorithms.xtoken_off_policy_distillation.setup( master_config: nemo_rl.algorithms.xtoken_off_policy_distillation.MasterConfig, student_tokenizer: transformers.tokenization_utils_base.PreTrainedTokenizerBase, teacher_tokenizers: list[transformers.tokenization_utils_base.PreTrainedTokenizerBase], train_dataset: nemo_rl.data.datasets.AllTaskProcessedDataset, val_dataset: Optional[nemo_rl.data.datasets.AllTaskProcessedDataset], ) → tuple[nemo_rl.models.policy.lm_policy.Policy, list[nemo_rl.models.policy.lm_policy.Policy], torchdata.stateful_dataloader.StatefulDataLoader, Optional[torchdata.stateful_dataloader.StatefulDataLoader], nemo_rl.algorithms.loss.loss_functions.CrossTokenizerDistillationLossFn, nemo_rl.utils.logger.Logger, nemo_rl.utils.checkpoint.CheckpointManager, nemo_rl.algorithms.xtoken_off_policy_distillation.OffPolicyDistillationSaveState, nemo_rl.algorithms.xtoken_off_policy_distillation.MasterConfig]#: Construct cluster, dataloaders, policies, and loss fn for the run.

nemo_rl.algorithms.xtoken_off_policy_distillation.export_teacher_logits_and_pack( teacher_policies: list[nemo_rl.models.policy.lm_policy.Policy], loss_fn: nemo_rl.algorithms.loss.loss_functions.CrossTokenizerDistillationLossFn, batch: nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any], teacher_mbs: list[int], *, timer: Optional[nemo_rl.utils.timer.Timer] = None, ) → nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any]#

Serially run each teacher’s forward and pack the student train_data.

Teachers run one at a time (collocated): each is onloaded for inference, forwarded, then offloaded. Every teacher ships full-vocab logits over CUDA IPC (teacher_{i}_full_logits_ipc) — the loss derives the microbatch-global top-k subset student-side, so there is no top-K transport. A same-vocab teacher (projection_matrix_paths[i] is None) reuses the student tokens and rides no extra alignment; a cross-tokenizer teacher’s own tokenization and alignment_{i}_* payload ride along (teacher-indexed). The persistent IPC buffers stay resident on the teacher GPUs until released by the caller after student.train. Shared by the train loop and validate so the forward+pack sequence can’t drift between them.

nemo_rl.algorithms.xtoken_off_policy_distillation.xtoken_off_policy_distillation_train( student_policy: nemo_rl.models.policy.lm_policy.Policy, teacher_policies: list[nemo_rl.models.policy.lm_policy.Policy], dataloader: torchdata.stateful_dataloader.StatefulDataLoader, val_dataloader: Optional[torchdata.stateful_dataloader.StatefulDataLoader], loss_fn: nemo_rl.algorithms.loss.loss_functions.CrossTokenizerDistillationLossFn, logger: nemo_rl.utils.logger.Logger, checkpointer: nemo_rl.utils.checkpoint.CheckpointManager, off_policy_distillation_state: nemo_rl.algorithms.xtoken_off_policy_distillation.OffPolicyDistillationSaveState, master_config: nemo_rl.algorithms.xtoken_off_policy_distillation.MasterConfig, ) → None#: Off-policy CT distillation training loop.

nemo_rl.algorithms.xtoken_off_policy_distillation.validate( student_policy: nemo_rl.models.policy.lm_policy.Policy, teacher_policies: list[nemo_rl.models.policy.lm_policy.Policy], val_dataloader: torchdata.stateful_dataloader.StatefulDataLoader, loss_fn: nemo_rl.algorithms.loss.loss_functions.CrossTokenizerDistillationLossFn, master_config: nemo_rl.algorithms.xtoken_off_policy_distillation.MasterConfig, skip_keys: frozenset[str], timer: Optional[nemo_rl.utils.timer.Timer] = None, ) → tuple[dict[str, Any], dict[str, Any]]#

Held-out KL/CE on a validation dataloader.

Reuses the same per-step path as training, but in eval mode so no backward / optimizer step runs. Returns mean train-style metrics.

skip_keys is the non-student-seq-axis key set for the worker’s check_sequence_dim pre-flight; built once in the train loop and threaded in so it can’t drift from a parallel rebuild here.

nemo_rl.algorithms.xtoken_off_policy_distillation#

Module Contents#

Classes#

Functions#

API#

`nemo_rl.algorithms.xtoken_off_policy_distillation`#