nemo_rl.algorithms.x_token.utils

Batch-grid helpers for cross-tokenizer off-policy distillation.

Teacher and student may run at different data-parallel degrees / microbatch sizes while consuming the same global batch. These helpers validate that the global batch tiles cleanly across both grids (:func:assert_teacher_student_batch_grid) and pad an uneven validation batch up to a tileable size (:func:pad_distillation_val_batch).

Module Contents

Functions

assert_teacher_student_batch_grid	Fail fast unless teacher and student share one global batch and both tile it.
assert_xtoken_ipc_node_local	Fail fast if the teacher->student logit transport would cross a node.
pad_distillation_val_batch	Pad every key of a validation batch up to target_size (batch axis).

API

nemo_rl.algorithms.x_token.utils.assert_teacher_student_batch_grid(

*,

global_batch_size: int,

student_gbs: int,

teacher_gbs: int,

student_dp: int,

teacher_dp: int,

student_mbs: int,

teacher_mbs: int,

) → None

Fail fast unless teacher and student share one global batch and both tile it.

A cross-tokenizer teacher and student may run at different data-parallel degrees and microbatch sizes, but they consume the SAME global batch in the SAME global order (the teacher exports a global-batch-ordered per-sample logits list and the student slices it by its own DP/MBS). That requires both sides’ train_global_batch_size to agree with the step’s global_batch_size, and each side to split that batch into whole per-DP-rank chunks and whole microbatches. Not specific to distillation — any xtoken teacher/student pairing needs it.

nemo_rl.algorithms.x_token.utils.assert_xtoken_ipc_node_local(

*,

num_nodes: int,

gpus_per_node: int,

student_tp: int,

student_cp: int,

teacher_tp: int,

teacher_cp: int,

student_dp: int,

teacher_dp: int,

) → None

Fail fast if the teacher->student logit transport would cross a node.

Teacher logits reach the student via CUDA IPC (cudaIpcOpenMemHandle), which can only map a buffer created by a process on the SAME physical node. A single-node job is therefore always safe; a multi-node job is only safe when every student rank’s required teacher shards live on its own node.

nemo_rl.algorithms.x_token.utils.pad_distillation_val_batch(

batch: nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any],

target_size: int,

) → nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any]

Pad every key of a validation batch up to target_size (batch axis).

Validation uses drop_last=False, so the final batch can be smaller than num_prompts_per_step and may not tile evenly across the teacher and student DP/MBS grids. Padding the whole batch symmetrically (student, teacher and alignment keys together) lets both sides run the existing even-split path with zero shared-code changes. The padded rows carry sample_mask == 0, so they are excluded from the valid-sample counts in process_global_batch and contribute nothing to the loss or gradients.