`nemo_rl.data.cross_tokenizer.cross_tokenizer_collate`#

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

The collator runs inside DataLoader worker processes. It does:

Tokenizes the student input once; this tokenization is shared by all teachers. In mode="text", tokenizes raw text without a chat template. In mode="chat", renders the full conversation with the student’s chat template and supervises only the final assistant turn (by default its content and end-of-turn token).
For each cross-tokenizer teacher, tokenizes with that teacher’s tokenizer and aligns with its :class:TokenAligner. Text mode aligns the full text; chat mode renders the teacher’s own chat template and aligns the final assistant message. Teacher scoring masks also select only the final assistant content and end-of-turn token. Dense-padded alignment and teacher inputs are emitted under alignment_{i}_* / teacher_{i}_*.
Same-tokenizer teachers (aligners[i] is None) emit nothing extra — their forward reuses the student tokenization, so projection and alignment are skipped.
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#

`CrossTokenizerCollatorConfig`	Shared batching options for cross-tokenizer distillation.
`CrossTokenizerCollator`	Tokenize the student once, tokenize+align each teacher, return a flat batch.

API#

class nemo_rl.data.cross_tokenizer.cross_tokenizer_collate.CrossTokenizerCollatorConfig#

Bases: pydantic.BaseModel

Shared batching options for cross-tokenizer distillation.

.. attribute:: mode

"text" tokenizes raw text; "chat" renders each model’s chat template and supervises assistant turns.

.. attribute:: include_thinking_in_loss

Include reasoning text and its closing marker in assistant supervision, excluding the opening scaffold. Separate reasoning_content fields require native thinking alignment; ordinary chat supports inline <think> content only when the template renders content verbatim. Only the final assistant turn is supervised; earlier turns remain context.

.. attribute:: native_thinking_alignment

Align reasoning, close, answer/tool and EOT independently on supported Qwen/Nano ChatML templates. Requires chat mode and thinking loss.

.. attribute:: kd_alignment_regions

Optional subset of reasoning, close, answer, and eot regions. Requires native thinking alignment.

.. attribute:: num_packed_rows

Positive number of examples packed into each row. Only one is currently supported.

mode: Literal[text, chat]#: ‘text’

include_thinking_in_loss: bool#: False

native_thinking_alignment: bool#: False

kd_alignment_regions: list[Literal[reasoning, close, answer, eot]] | None#: None

num_packed_rows: pydantic.PositiveInt#: 1

class nemo_rl.data.cross_tokenizer.cross_tokenizer_collate.CrossTokenizerCollator( *, config: nemo_rl.data.cross_tokenizer.cross_tokenizer_collate.CrossTokenizerCollatorConfig, 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], drop_first_assistant_chunk_kl_by_teacher: List[bool], make_seq_div_by_student: int = 1, make_seq_div_by_teachers: Optional[List[int]] = None, reject_overlength_text: bool = False, ce_covers_assistant_turn: bool = False, add_special_tokens_in_text: bool = False, )#

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

Supports N teachers in raw-text or chat mode. Student inputs are tokenized once and shared; each cross-tokenizer teacher uses its own tokenizer and

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. Chat mode applies each model’s chat template to the full conversation, masks teacher scoring (and, by default, loss) to the final assistant content and end-of-turn token, and aligns only that final message.

Parameters:

config – Typed batching options, including text/chat mode and the thinking-region and packing settings.
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. Batches pad to their longest tokenization rounded up to make_seq_div_by_student, so the padded width can exceed a cap that is not a multiple of the divisor.
ctx_length_teachers – Per-teacher tokenization length caps.
drop_first_assistant_chunk_kl_by_teacher – Per-teacher flags controlling whether chat alignment drops the first content pair in each assistant message. The list retains a slot for same-tokenizer teachers so its indices match aligners.
make_seq_div_by_student – Round student sequence length up to a multiple of this value (typically TP * CP * 2 for Automodel).
make_seq_div_by_teachers – Per-teacher sequence-length divisors.
reject_overlength_text – Text-mode policy for rows longer than a context length. False truncates the student and each cross-tokenizer teacher to its own context length; a same-tokenizer teacher shares the student’s truncation. True raises ValueError for a row whose student or teacher tokenization, or that length rounded up to the side’s sequence-length divisor, exceeds that side’s context length, so every side covers the same text and no padded batch is wider than its context; same-tokenizer teachers check the shared student tokenization against their own context length, rounded up to both the student’s divisor and their own. Chat mode always rejects rows longer than a context length; same-tokenizer teachers also check the width after padding to the student’s divisor.
ce_covers_assistant_turn – Chat-mode CE target. False supervises the final assistant content and end-of-turn token. True supervises every token of the final assistant turn (SFT-style), including chat-template markers and any rendered reasoning, and emits the content and end-of-turn mask as kd_token_mask. Alignment and teacher token masks are unchanged. The turn starts where the rendering of all earlier messages ends, or at the start of the rendering when there are none; that rendering must be a prefix of the full rendering. Requires mode="chat" without native_thinking_alignment.
add_special_tokens_in_text – Text-mode special tokens. False tokenizes the raw text only. True tokenizes the student and every cross-tokenizer teacher with each tokenizer’s own special tokens (for example a beginning-of-text marker at the start of each row); reject_overlength_text counts them. Added tokens carry (0, 0) offsets, so alignment pairs them by role and a marker without a counterpart forms its own chunk with an empty other side. Requires mode="text"; chat templates insert their own markers.

Initialization

__call__( batch: List[nemo_rl.data.interfaces.DatumSpec], ) → nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any]#

_call_text( batch: List[nemo_rl.data.interfaces.DatumSpec], ) → nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any]#

_render_document( tokenizer: transformers.tokenization_utils_base.PreTrainedTokenizerBase, datum: nemo_rl.data.interfaces.DatumSpec, ctx_length: int, side: str, ) → nemo_rl.data.cross_tokenizer.native_chat.RenderedChatDocument#

_final_turn_start( datum: nemo_rl.data.interfaces.DatumSpec, doc: nemo_rl.data.cross_tokenizer.native_chat.RenderedChatDocument, ) → int#

Return the index of the first student token of the final assistant turn.

The turn spans from the end of the rendering of the earlier messages (empty when there are none) to the end of the full rendering.

_call_chat( batch: List[nemo_rl.data.interfaces.DatumSpec], ) → nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any]#: Render complete conversations and align their final assistant regions.

static _padded_width(length: int, make_seq_div_by: int) → int#: length rounded up to a multiple of make_seq_div_by, and at least max(make_seq_div_by, 1).

static _pad_chat_batch( ids_list: List[List[int]], off_list: List[List[tuple[int, int]]], mask_list: List[List[int]], pad_token_id: int, make_seq_div_by: int, ) → tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]#

Right-pad per-sample chat tokenizations into dense [B, T] tensors.

Pads to the batch-max length rounded up to make_seq_div_by. Padding positions get attention_mask=0, offset=(0, 0), assistant_mask=0.

static _reject_overlength_texts( texts: List[str], idx: List[int], tokenizer: transformers.tokenization_utils_base.PreTrainedTokenizerBase, ctx_length: int, side: str, *, make_seq_div_by: int, add_special_tokens: bool, ) → None#

Raise ValueError if any text needs more than ctx_length positions.

A row needs its token count rounded up to make_seq_div_by.

static _tokenize_batch( texts: List[str], tokenizer: transformers.tokenization_utils_base.PreTrainedTokenizerBase, ctx_length: int, make_seq_div_by: int, *, add_special_tokens: bool, ) → tuple[torch.Tensor, torch.Tensor, torch.Tensor]#

Tokenize a batch, truncating to ctx_length, and pad to its longest row.

Rows are right-padded to the longest row rounded up to a multiple of make_seq_div_by, the same width rule as chat mode. Also returns per-token character offsets (offset_mapping), which

Meth:: TokenAligner.align needs to align the student and teacher tokenizations of the same source text. This requires a fast HF tokenizer; special and padding positions carry (0, 0).

nemo_rl.data.cross_tokenizer.cross_tokenizer_collate#

Module Contents#

Classes#

API#

`nemo_rl.data.cross_tokenizer.cross_tokenizer_collate`#