[avenkateshha]

Superseded by #4680 after renaming the contribution branch to avenkateshha/csd. The implementation commit and stacked base are unchanged.

What does this PR do ?

Adds CSD as a separate cross-tokenizer distillation algorithm, with sub-token table preparation and a complete training workflow. Existing X-Token loss behavior, example configurations, and projection-matrix preparation remain available.

CSD provides common-event and prefix-support losses, dense or sparse teacher-logit transport, bounded-memory tensor/context-parallel student losses, multi-teacher weighting, and deterministic lockstep sequence packing. It reuses the existing alignment, colocated teacher execution, and training infrastructure, with integration for Automodel and Megatron, model-load precision controls, and router replay. CSD-specific collation keeps complete source samples and validates durable identities, semantic masks, and sequence limits.

The public guide covers building and validating both table directions, configuring their paths, and launching training. CPU preprocessing includes exact byte-preserving token chains, special-token roles, tokenizer fingerprints, and validation. The provided builder supports fast BPE tokenizers with ByteLevel decoders; unsupported decoder families fail explicitly.

Issues

Depends on and is stacked on #3286. The stack is main → #3286 → this PR. Both branches are in NVIDIA-NeMo/RL:

• Base: avenkateshha/xtoken-offset-alignment, at 247493fd04b8eae8af7bcfedf4470750ec11b380 when this contribution was prepared.
• Head: avenkateshha/csd.

The parent implementation is inherited through history; this diff contains the additional CSD work and integration changes. This PR does not target main.

Usage

Follow the cross-tokenizer guide for prerequisites, both table-building commands, validation, and a complete two-step GPU training example. The existing X-Token guide documents its separate projection workflow.

uv run python -m tools.csd.build_subtoks_table \
  --source-tokenizer meta-llama/Llama-3.2-1B \
  --destination-tokenizer Qwen/Qwen3-4B \
  --output cross_tokenizer_data/csd/student_to_teacher.pt
uv run python -m tools.csd.build_subtoks_table \
  --source-tokenizer Qwen/Qwen3-4B \
  --destination-tokenizer meta-llama/Llama-3.2-1B \
  --output cross_tokenizer_data/csd/teacher_to_student.pt
uv run python -m tools.csd.validate_subtoks_tables \
  --student-tokenizer meta-llama/Llama-3.2-1B \
  --teacher-tokenizer Qwen/Qwen3-4B \
  --student-to-teacher cross_tokenizer_data/csd/student_to_teacher.pt \
  --teacher-to-student cross_tokenizer_data/csd/teacher_to_student.pt

Training uses examples/run_csd.py with examples/configs/csd.yaml or csd_multiteacher.yaml. Each cross-tokenizer teacher consumes teachers[i].aligner.student_to_teacher_table_path and teachers[i].aligner.teacher_to_student_table_path. Prepare or filter the input corpus so complete rows fit both model context limits; see the guide for a self-contained smoke dataset and launch command.

Before your PR is "Ready for review"

• Read and followed the Contributor guidelines.
• Added and updated necessary tests.
• Ran the complete unit and functional GPU suites locally.
• Added and updated documentation.

Additional Information

Executed focused CPU checks include 162 combined X-Token/CSD controller tests; CSD loss values, gradients, dense/sparse parity and TP2/CP2 Gloo parity; 85 packing/distributed/IPC tests (three GPU skips); dataset/sample-identity, collator supervision, configuration, and table-builder/validator CLI regressions. Ruff lint/import ordering/formatting, recipe minimization, shell syntax, and whitespace checks were also run.

Full GPU training and CUDA IPC validation were unavailable in the local environment. Megatron runtime dependencies were unavailable, and real-model tokenizer checks require model access. Table preparation was exercised with local tokenizer fixtures. The local Python environment was 3.13.13 rather than the checkout's pinned 3.13.14. Full-project type checking is limited by unavailable or mismatched optional dependencies; focused checks cover the new CSD modules. The GPU recipes are registered as manual/disabled tests pending hardware qualification.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[avenkateshha]

/ok to test 22ebcf7