[3/5] feat: cross-tokenizer distillation loss and multi-teacher aggregator

[avenkateshha]

Adds the loss-fn layer for cross-tokenizer distillation. Builds on the TokenAligner package (PR 1).

• CrossTokenizerDistillationLossFn: per-token KL/CE loss over 1:1 aligned positions, with optional gold-loss path. Holds a reference to a TokenAligner; teacher data (input ids, aligned pairs, optional chunked COO masks) is set per-step via set_cross_tokenizer_data.
• CrossTokenizerDistillationLossConfig and CrossTokenizerDistillationLossDataDict TypedDicts.
• MultiTeacherLossAggregator: wraps a list of optional CrossTokenizerDistillationLossFn instances with per-teacher weights. N=1 is a degenerate case used by the unified single-/multi-teacher worker path; the algorithm-layer multi-teacher orchestration comes in PR 5.
• _scatter_chunk_mask_from_coo helper for the chunked-CE path.

What does this PR do?

Adds CrossTokenizerDistillationLossFn and MultiTeacherLossAggregator — the per-token KL loss and weighted aggregator used by cross-tokenizer off-policy distillation.

Issues

None linked yet.

Usage

from nemo_rl.algorithms.x_token.tokenalign import TokenAligner
from nemo_rl.algorithms.loss.loss_functions import (
    CrossTokenizerDistillationLossFn,
    MultiTeacherLossAggregator,
)

aligner = TokenAligner(...)  # see PR 1
loss_cfg = {"loss_type": "KL", "temperature": 1.0, "vocab_topk": 8192,
            "exact_token_match_only": False, "reverse_kl": False}
loss_fn = CrossTokenizerDistillationLossFn(loss_cfg, aligner)

# Per training step (after collator produces per-teacher CT data):
loss_fn.set_cross_tokenizer_data(teacher_input_ids, aligned_pairs)

# Or for multi-teacher:
agg = MultiTeacherLossAggregator(
    [loss_fn_t0, loss_fn_t1], weights=[0.5, 0.5], cfg=loss_cfg,
)
agg.set_cross_tokenizer_data(teacher_input_ids_t0, aligned_pairs_t0, teacher_idx=0)
agg.set_cross_tokenizer_data(teacher_input_ids_t1, aligned_pairs_t1, teacher_idx=1)

Before your PR is "Ready for review"

• Read Contributor guidelines
• No new unit tests in this PR (next-step follow-up). Existing tests/unit/algorithms/test_loss_functions.py continues to pass for stock loss fns.
• Static py_compile confirmed clean. Functional smoke covered by PR 5.
• No docs entry — added alongside PR 5.

Additional Information

Draft. Stacked on PR 2 (collator) — #2348. The aggregator class is included here as a building block; no algorithm-level multi-teacher orchestration ships in this PR.

Full chain:

1. TokenAligner + projection utilities — [1/5] feat: add TokenAligner and cross-tokenizer projection utilities #2347
2. Collator + Arrow dataset + eval datasets — [2/5] feat: cross-tokenizer collator, Arrow dataset, and eval datasets #2348
3. (this PR) CT distillation loss + multi-teacher aggregator
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.