[2/5] feat: cross-tokenizer collator, Arrow dataset, and eval datasets

[avenkateshha]

Data-layer plumbing for cross-tokenizer off-policy distillation, plus in-training eval datasets. Builds on the TokenAligner package from PR 1.

• nemo_rl/data/cross_tokenizer_collate.py: CrossTokenizerCollator and TeacherCTSpec. Runs in StatefulDataLoader worker processes — does per-teacher tokenize + DP alignment up front so the train loop only consumes pre-built per_teacher_ct_data. Lazy-imports TokenAligner so workers that don't need cross-tokenizer never touch x_token.
• nemo_rl/data/__init__.py: add NotRequired prefetch_factor to DataConfig.
• nemo_rl/data/datasets/response_datasets/arrow_text_dataset.py: ArrowTextDataset with lazy packing, registered as "arrow_text" in DATASET_REGISTRY.
• nemo_rl/data/datasets/eval_datasets/{humaneval_plus, mbpp_plus, mmlu}.py and registry entries: in-training eval datasets. mmlu.py adds an optional num_few_shot argument with a static _build_few_shot_prefixes helper; default of 0 preserves existing behavior.

What does this PR do?

Adds the data-layer plumbing (collator + Arrow dataset) consumed by the cross-tokenizer distillation training loop, plus HumanEval+/MBPP+/MMLU eval datasets used for in-training evaluation.

Issues

None linked yet.

Usage

from torchdata.stateful_dataloader import StatefulDataLoader
from nemo_rl.data.cross_tokenizer_collate import CrossTokenizerCollator, TeacherCTSpec

teacher_ct_specs = [TeacherCTSpec(
    teacher_tokenizer_name="microsoft/Phi-4-mini-instruct",
    student_tokenizer_name="meta-llama/Llama-3.2-1B",
    projection_matrix_path="cross_tokenizer_data/llama_phi-mini_proj.pt",
    use_sparse_format=False, learnable=False, max_comb_len=4,
    projection_matrix_multiplier=1.0, project_teacher_to_student=False,
    max_teacher_len=4096, dp_chunk_size=128, use_align_fast=True,
    exact_token_match_only=False,
)]

collator = CrossTokenizerCollator(
    pad_token_id=tokenizer.pad_token_id,
    make_sequence_length_divisible_by=1,
    teacher_ct_specs=teacher_ct_specs,
    fallback_student_tokenizer_name="meta-llama/Llama-3.2-1B",
)

loader = StatefulDataLoader(
    dataset, batch_size=768, collate_fn=collator,
    num_workers=8, prefetch_factor=4, persistent_workers=True,
)

Before your PR is "Ready for review"

• Read Contributor guidelines
• No new tests in this PR. Collator behavior is exercised by the off-policy distillation recipe in PR 5.
• Static py_compile confirmed clean. Full functional run lands with PR 5.
• No docs entry yet — added alongside PR 5.

Additional Information

Stacked on PR 1 (TokenAligner) — #2347. Imports TokenAligner lazily so workers without CT teachers don't pay the cost.

Full chain:

1. TokenAligner + projection utilities — [1/5] feat: add TokenAligner and cross-tokenizer projection utilities #2347
2. (this PR) Collator + Arrow dataset + eval datasets
3. CT distillation loss + multi-teacher aggregator — [3/5] feat: cross-tokenizer distillation loss and multi-teacher aggregator #2349
4. CUDA IPC for teacher logits — [4/5] feat: CUDA IPC for teacher logits transfer #2350
5. Algorithm + worker integration — [5/5] feat: off-policy distillation algorithm and worker integration #2351

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.