nemo_rl.data.cross_tokenizer_collate#

Collator that tokenizes the student once, then tokenizes+aligns each teacher.

The collator runs inside DataLoader worker processes. It does:

  1. Tokenizes the source text once with the student tokenizer (no chat template, no special handling); this tokenization is shared by all teachers.

  2. For each cross-tokenizer teacher, tokenizes with that teacher’s tokenizer and calls :class:TokenAligner.align to produce a dense-padded

    class:

    AlignmentBatch (P-KL, gold_loss, xtoken_loss), emitted under teacher-indexed keys teacher_{i}_* / alignment_{i}_*.

  3. Same-tokenizer teachers (aligners[i] is None) emit nothing extra — their forward reuses the student tokenization, so projection and alignment are skipped.

  4. Returns a :class:BatchedDataDict with the keys :class:Policy.train expects (input_ids, input_lengths, token_mask, sample_mask) plus per-teacher tensors and alignment tensors.

Loss-side projection-matrix work happens inside the loss fn; nothing related to KL/CE math runs here.

Module Contents#

Classes#

CrossTokenizerCollator

Tokenize the student once, tokenize+align each teacher, return a flat batch.

API#

class nemo_rl.data.cross_tokenizer_collate.CrossTokenizerCollator(
*,
student_tokenizer: transformers.tokenization_utils_base.PreTrainedTokenizerBase,
teacher_tokenizers: List[Optional[transformers.tokenization_utils_base.PreTrainedTokenizerBase]],
aligners: List[Optional[nemo_rl.algorithms.x_token.token_aligner.TokenAligner]],
ctx_length_student: int,
ctx_length_teachers: List[int],
make_seq_div_by_student: int = 1,
make_seq_div_by_teachers: Optional[List[int]] = None,
)#

Tokenize the student once, tokenize+align each teacher, return a flat batch.

Supports N teachers. The student text is tokenized once and shared; each cross-tokenizer teacher is tokenized with its own tokenizer and aligned with its own :class:TokenAligner, emitting teacher-indexed keys (teacher_{i}_* and alignment_{i}_*). A same-tokenizer teacher (aligners[i] is None) emits nothing extra — its forward reuses the student tokenization, so projection and alignment are skipped entirely.

Parameters:

  • student_tokenizer – HF tokenizer matching the student model.

  • teacher_tokenizers – Per-teacher HF tokenizers. May be None for a same-tokenizer teacher (its tokenization is the student’s).

  • aligners – Per-teacher :class:TokenAligner. None marks a same-tokenizer teacher (no projection / no alignment).

  • ctx_length_student – Hard tokenization length cap on the student side (also the padded sequence length of the student tensor).

  • ctx_length_teachers – Per-teacher tokenization length caps.

  • make_seq_div_by_student – Round student sequence length up to a multiple of this value (typically TP * CP * 2 for DTensor V2).

  • make_seq_div_by_teachers – Per-teacher sequence-length divisors.

Initialization

__call__(
batch: List[nemo_rl.data.interfaces.DatumSpec],
) → nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any]#
static _tokenize_batch(
texts: List[str],
tokenizer: transformers.tokenization_utils_base.PreTrainedTokenizerBase,
ctx_length: int,
make_seq_div_by: int,
) → tuple[torch.Tensor, torch.Tensor]#

Tokenize a batch and pad to a multiple of make_seq_div_by.