feat(xtoken): auto-build the cross-tokenizer projection matrix on the fly (online)
- 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
Assessment
This issue has not been assessed yet.