[Feature] Add same-tokenizer off-policy distillation support
- Dominant language
- Python
- Stars
- 2k
- Forks
- 561
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 145
Description
## Background
The current [off-policy distillation](https://github.com/NVIDIA-NeMo/RL/pull/2508) work focuses on the cross-tokenizer path: raw text data is tokenized for both student and teacher, teacher logits are produced on the teacher-tokenized sequence, and the student is trained with cross-tokenizer alignment/projection logic.
Longer term, we should also support the same-tokenizer off-policy distillation path. In that case, student and teacher share the same tokenizer, so we should not require cross-tokenizer alignment or projection. The pipeline should consume raw text data, run teacher forward on the same tokenized sequence, and train the student directly from teacher logits.
## Goal
Add a same-tokenizer off-policy distillation mode that reuses the off-policy distillation pipeline without requiring x-token-specific components.
## Expected Behavior
When student and teacher use the same tokenizer:
- Use a standard single-tokenizer collator.
- Tokenize raw text once.
- Run teacher forward on the same `input_ids` / `input_lengths`.
- Store teacher logits in the normal same-tokenizer distillation format.
- Use the existing same-tokenizer `DistillationLossFn`.
- Do not require:
- `TokenAligner`
- cross-tokenizer alignment fields
- projection matrix
- `CrossTokenizerDistillationLossFn`
- `get_full_logits_ipc` unless full logits are explicitly needed by the chosen same-tokenizer loss path.
## Proposed Design
Introduce an off-policy distillation entrypoint/config that can route based on a config flag, for example:
```yaml
distillation:
cross_tokenizer: false
```
For cross_tokenizer: false:
```
loss_fn = DistillationLossFn(loss_config)
collator = standard_raw_text_or_kd_collator(...)
teacher_forward = get_topk_logits(...)
train_data = {
"input_ids": ...,
"input_lengths": ...,
"token_mask": ...,
"teacher_topk_logits": ...,
"teacher_topk_indices": ...,
}
```
For cross_tokenizer: true, keep the existing x-token path:
```
loss_fn = CrossTokenizerDistillationLossFn(loss_config)
collator = CrossTokenizerCollator(...)
teacher_forward = get_full_logits_ipc(...)
train_data = {
"input_ids": ...,
"teacher_input_ids": ...,
"teacher_full_logits_ipc": ...,
"alignment_*": ...,
}
```
## Scope
This issue should cover:
- Config updates for same-tokenizer off-policy distillation.
- Entrypoint/algorithm routing between same-tokenizer and cross-tokenizer modes.
- Teacher forward path for same-tokenizer off-policy distillation.
- Train-data assembly for DistillationLossFn.
- Example config for same-tokenizer off-policy distillation.
- Unit or smoke tests covering the same-tokenizer route.
Contributor guide
Assessment
This issue has not been assessed yet.