NVIDIA-NeMo / NVIDIA-NeMo/RL

feat(xtoken): auto-build the cross-tokenizer projection matrix on the fly (online)

Open
#2,792 0 comments 0 reactions 1 assignee Claimed by @RayenTian View on GitHub
Feature
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

## Summary

Optionally build (and cache) the cross-tokenizer (xtoken) distillation **projection matrix on the fly**, so users don't have to run a separate offline step before `run_xtoken_off_policy_distillation.py`.

## Current behavior

The cross-tokenizer distillation loss requires a pre-built projection matrix:

- `examples/run_xtoken_off_policy_distillation.py` errors when `loss_fn.projection_matrix_path` is null (it does not auto-build).
- The matrix is produced by a separate offline tool, e.g.:
```bash
uv run python -m tools.x_token.minimal_projection_via_multitoken --student-model meta-llama/Llama-3.2-1B --teacher-model Qwen/Qwen3-4B --top-k 4 --enable-special-token-mapping --enable-exact-match --disable-reverse-pass --disable-scale-trick --output-filename xtoken_proj --output-dir /tmp/xtoken_proj
```
- The nightly recipe `xtoken-off-policy-distillation-qwen3-4b-to-llama3.2-1b-1n8g-dtensor-tp4cp2.sh` wraps this as a shell "build-if-missing" step.

## Proposal

When `loss_fn.projection_matrix_path` is null, build the projection during `setup` from an optional config block and cache it to a deterministic path, then reuse it. Sketch:

```yaml
loss_fn:
projection_matrix_path: null # if null and projection_build is set -> auto-build + cache
projection_build:
student_model: ...
teacher_model: ...
top_k: 4
special_token_mapping: true
exact_match: true
# ... remaining minimal_projection_via_multitoken flags
```

- Cache key derived from (student tokenizer, teacher tokenizer, build flags) so repeated runs reuse the same artifact and the result is reproducible.
- Explicit `projection_matrix_path` continues to work and takes precedence.
- Log the resolved build flags + cache path so the (convergence-affecting) choices stay visible.

## Why deferred (not in the originating PR)

The build flags (`top_k`, `special_token_mapping`, `exact_match`, `reverse_pass`, `scale_trick`, ...) are convergence-affecting experiment choices; an explicit pre-built artifact is more reproducible and keeps those choices visible. The originating PR therefore keeps the explicit two-step usage and tracks the convenience feature here.

## Acceptance criteria

- [ ] Null `projection_matrix_path` + `projection_build` config → projection is built, cached, and reused on subsequent runs.
- [ ] Explicit `projection_matrix_path` still works and takes precedence.
- [ ] Build flags and cache path are logged.
- [ ] Cache is keyed so a flag/tokenizer change rebuilds rather than silently reusing a stale matrix.

## References

- Follow-up to the TP/CP/diff-DP sharded cross-tokenizer distillation PR (branch `ruit/xtoken-tp-cp`).
- Builder: `tools/x_token/minimal_projection_via_multitoken.py`
- Run entry: `examples/run_xtoken_off_policy_distillation.py` (null-path check)

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.