[5/5] feat: off-policy distillation algorithm and worker integration - #2351

Closed

avenkateshha wants to merge 11 commits into

NVIDIA-NeMo:mainfrom

avenkateshha:avenkateshha/xtoken-off-policy-distillation/05-off-policy-distillation

Commits on Apr 27, 2026

feat: add TokenAligner and cross-tokenizer projection utilities

avenkateshha
and
claude
committed
c492ac1
feat: cross-tokenizer collator, Arrow dataset, and eval datasets

avenkateshha
and
claude
committed
5934699
feat: cross-tokenizer distillation loss and multi-teacher aggregator

avenkateshha
and
claude
committed
70a14eb
feat: CUDA IPC for teacher logits transfer

avenkateshha
and
claude
committed
a146953
feat: off-policy distillation algorithm and worker integration

avenkateshha
and
claude
committed
f9d931b

Commits on Apr 28, 2026

feat(policy): add cross-tokenizer init/update wrappers on Policy

avenkateshha
and
claude
committed
2fb373f
fix(policy): repair Phi RoPE meta buffers + correct teacher_logits dispatch

avenkateshha
and
claude
committed
a14fa99
fix(distillation): wire token_aligners through setup→runner + val cached loss fn fallback

avenkateshha
and
claude
committed
a590172
fix: restore optimizer state to CUDA when keep_models_resident=False

avenkateshha
and
claude
committed
0cc0c28
fix(policy): load inference-only Phi models in compute dtype

avenkateshha
and
claude
committed
214e570
refactor(x_token): drop use_align_fast in off-policy distillation setup

avenkateshha
and
claude
committed
6d9246d