NVIDIA-NeMo / NVIDIA-NeMo/RL

[Feature] Add same-tokenizer off-policy distillation support

Open
#2,545 0 comments 0 reactions 1 assignee Claimed by @terrykong View on GitHub
enhancement Feature
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.