`nemo_rl.data.xtoken.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 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.xtoken.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.xtoken.cross_tokenizer_collate.CrossTokenizerCollator( *, config: nemo_rl.data.xtoken.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, )#

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 loss and teacher scoring 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. Text mode pads to this cap; chat mode pads to the batch’s longest tokenization. Sequence-length divisors may add padding.
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.

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.xtoken.native_chat.RenderedChatDocument#

_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 _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 _tokenize_batch( texts: List[str], tokenizer: transformers.tokenization_utils_base.PreTrainedTokenizerBase, ctx_length: int, make_seq_div_by: int, ) → tuple[torch.Tensor, torch.Tensor, torch.Tensor]#

Tokenize a batch and pad to a multiple of make_seq_div_by.

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.xtoken.cross_tokenizer_collate#

Module Contents#

Classes#

API#

`nemo_rl.data.xtoken.cross_tokenizer_collate`#