feat(csd): add cross-tokenizer distillation and sub-token preparation

[avenkateshha]

What does this PR do ?

Adds CSD (Common-Support Distillation), a second cross-tokenizer off-policy distillation method next to X-Token. CSD aligns the student and teacher tokenizations of the same text into chunks and compares the two models only on events both tokenizers can express — shared tokens and shared token chains, looked up in bidirectional sub-token tables — with a REST bucket for the remaining probability mass. Student next-token CE is added to the KD term, and only the student is trained.

	X-Token	CSD
Selector	algorithm: xtoken (default)	algorithm: csd
Vocabulary artifact	Projection matrix (tools/x_token/)	Student-to-teacher and teacher-to-student sub-token tables (tools/csd/)
Entrypoint / exemplars	run_xtoken_off_policy_distillation.py	run_csd.py, csd.yaml, csd_multiteacher.yaml
Backends	Automodel	Automodel or Megatron, per model
Teacher transport	Dense CUDA IPC	Sparse top-K (default K=128) or dense CUDA IPC
Sequence packing	Unpacked	Unpacked, or optional lockstep packing

Main changes

• Loss (nemo_rl/algorithms/loss/csd.py): CSDLossFn / CSDLossConfig. 1:1 chunks compare common tokens; 1:many and many:1 chunks compare prefix-support alternatives (or only the realized chains with csd_realized_only_mismatch); many:many chunks compare realized chains. Defaults: JSD (β=0.5), mismatch chunks weighted ×5, normalized by the DP-global valid chunk count × T², dynamic KD/CE scaling. Same-tokenizer teachers use the shared top-K KL from refactor(xtoken): move same-tokenizer KD math into shared helpers #4721. Multi-teacher sum / averaged_logits / select_teacher.
• Sub-token tables (tools/csd/): a CPU builder (one direction per call; complete chains, role-mapped special tokens, no truncation) and a validator (fingerprints, ID ranges, every mapped event). Fast BPE tokenizers with a ByteLevel decoder.
• Teacher transport: sparse top-K IPC (top-K logits, exact log-normalizer, forced realized labels) from Automodel and Megatron teachers, plus Megatron dense full-logit export. The default teacher_topk_ipc_k: 128 changes the objective (top-128 + REST instead of the full vocabulary); set 0 for dense.
• Distributed execution: TP/CP-sharded loss without gathering the vocabulary; eligible Megatron students (PP=1, unpacked, static sum) evaluate KD owner-local on each CP rank. New primitives in nemo_rl/distributed/model_utils.py.
• Data: durable sample_id on arrow_text, openai_format and ResponseDataset; CSDCrossTokenizerCollator is a thin subclass of the shared collator with feat(xtoken): cross-tokenizer collator options and pad text to longest row #4728's options on. Overlength text rows are clipped to a prefix that fits every model.
• Lockstep sequence packing (optional, CSD only): the student and every teacher pack the same samples into the same bins (nemo_rl/data/packing/lockstep.py, docs/design-docs/csd-lockstep-sequence-packing.md).
• Configs and docs: exemplars; three nightly recipes (Automodel TP4CP2 single- and multi-teacher, Megatron TP2CP2); docs/guides/xtoken-off-policy-distillation.md renamed to docs/guides/cross-tokenizer-distillation.md, with a new CSD section next to X-Token; README "Cross-Tokenizer Distillation" section covering X-Token and CSD.

Shared changes that also affect X-Token and other algorithms:

• Megatron CP > 1 no longer requires sequence packing, and packed Megatron runs pass raw cu_seqlens_q/kv and a CP-local MoE router padding mask.
• The arrow_text packed-cache schema version changes, so existing characters_per_sample caches are rebuilt on first use.
• rebuild_teacher_full_logits_from_ipc returns (logits, fallback_count); the strict exact-token map selects every row whose top weight is 1.0.
• Teacher table paths are accepted only under teachers[i].aligner.

Issues

Stacked on #4728: main ← #3286 ← #4721 ← #4728 ← this PR (base avenkateshha/xtoken-csd-shared-collator). Merge in stack order.

Usage

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
uv run python examples/run_csd.py

See the CSD section of the cross-tokenizer guide for table format, loss options, transports and lockstep packing.

Before your PR is "Ready for review"

• Read and followed the Contributor guidelines.
• Added and updated necessary tests.
• Added and updated documentation.
• Completed GPU qualification.

Additional Information

Validation so far

Check	Result
CPU unit suites (CSD, X-Token, data, distributed, tools)	Pass; remaining failures are offline Hugging Face downloads, identical on the base
GPU unit tests, default markers	499 passed, 0 failed
GPU smoke, csd-qwen3-4b-to-llama3.2-1b-1n8g-dtensor-tp4cp2	8 steps + validation, finite loss
Megatron owner-local vs. replicated KD (TP2×CP2, K=0, 20 steps)	Step 1: logits bitwise equal, loss/CE identical, KD rel. diff 9e-8; grad-norm rel. diff 2.2e-4. K=128: finite for 20 steps
Same-tokenizer KD and collator vs. pre-refactor CSD	Bitwise identical (28 KD cases; 21 of 24 collator cases, the rest are intended behavior changes)

Open items

• test_nightly_compute_stays_below_4947_hours fails: the CSD nightlies bring the nightly total to 4955 GPU-hours.
• Megatron-marked unit tests need an image newer than nemo_rl_main_20260803 on the current base.
• The multi-teacher recipe and lockstep packing have no GPU run on the current head; no convergence claims yet (results will be added to the guide).

Suggested review order: docs/guides/cross-tokenizer-distillation.md (CSD section) → tools/csd/ → nemo_rl/algorithms/loss/csd.py → nemo_rl/algorithms/x_token/loss_utils.py and nemo_rl/distributed/model_utils.py → nemo_rl/algorithms/xtoken_off_policy_distillation.py → nemo_rl/data/cross_tokenizer/csd_collate.py and nemo_rl/data/packing/lockstep.py → backend/worker changes under nemo_rl/models/ → tests.

🤖 Generated with Claude Code

[avenkateshha]

/ok to test 22ebcf7

[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]

Self-review of the CSD PR at c516f035b. It was run with an agent team: RL guidelines, three bug-finders, an upstream-API expert, design and tests. Three devil's-advocate passes then re-checked every finding at this head; findings below 80% confidence were dropped. Everything actionable is inline; this body holds context and follow-ups.

Fix in this PR (details inline):

1. megatron/data.py:2032: TE infers pad_between_seqs=True for every Megatron packed run, so the grpo-cispo-…-megatron-cispo nightly fails. This affects non-CSD algorithms.
2. megatron_policy_worker.py:1607: empty mb_losses on an all-padding DP rank → torch.stack([]) crash (Megatron SFT/DPO validation).
3. csd.py:654: Automodel CP>1 dynamic loss scaling makes KD CP× too strong (both Automodel CP=2 nightlies).
4. nightly.txt:324: nightly GPU-hours 4955 > 4947 cap (L0 unit failure).
5. test_subtoks_tables.py:1: copyright header (CI red).
6. test_train.py:335: the new Megatron test fails on the pinned MCore (a test bug).

PR description: remaining gaps (the latest rewrite fixed the rest):

• "Shared changes" doesn't mention that X-Token now offloads the student optimizer after every step and every validation when policy.offload_optimizer_for_logprob: true (controller#L2348-L2356). Both X-Token exemplars set it; the nightly recipes override it to false.
• Megatron-FSDP (use_megatron_fsdp / use_custom_fsdp: true) is now honored for every algorithm, with HF weights loaded directly (setup.py#L669-L677, #L983-L991); the base ignored the flag.
• New shared policy keys aren't listed: freeze_parameter_patterns (Automodel-only, see inline), automodel_cfg.load_precision / shard_before_load, the Megatron optimizer dtype pass-throughs, and reduce_scatter_with_fp32_accumulation.
• "CP > 1 no longer requires sequence packing" should note that logprob, top-k and full-payload export still need it (see setup.py:1344).
• The two Megatron-marked failures attributed to the image are test-side (test_train.py:335, test_megatron_setup.py:205).
• The validation section doesn't mention the red copyright-check.

Docs: docs/design-docs/csd-lockstep-sequence-packing.md L51-53, L602-604 and L608 still describe limitations this PR removed. For example, "Both teacher IPC entry points in Policy explicitly reject sequence packing" is no longer true: get_full_logits_ipc now takes a packing plan (lm_policy.py#L921-L1005). Reword those passages as history, or delete them.

Follow-ups (not blocking):

• About 330 lines of multi-teacher aggregation in CSDLossFn are near-copies of X-Token's CrossTokenizerDistillationLossFn (similarity: _teacher_weight_score 1.00, _compute_dynamic_weights 0.99, _teacher_score_inputs 0.94, …). The select_teacher / CP crash (inline at controller L1557) already lives in both copies. The CSD copy also reduces over a per-call side-channel group (_dp_cp_group / _sum_window_kd_over_cp, csd.py#L331-L334, set at L557-L558). A shared helper that takes the group explicitly would fix both.
• CSDSequencePackingLossWrapper._accumulate_metrics adds an 8th copy of the min/max metric sets (csd_packing.py#L455-L456), the duplication the TODO in wrapper.py#L29-L46 asks to remove. CSD emits no _min/_max metrics, so import _SEQ_METRIC_MIN/_MAX or drop the branches.

Worth keeping: the algorithm-free CP/TP primitives in distributed/model_utils.py; owner_local_config_violations / csd_owner_local_ineligibility as pure reason-list functions reused by __init__ and setup; and the CPU-testable CSD math with its gloo TP2/CP2 parity tests.

Generated by Claude Code

[avenkateshha]

1 action item, please fix in this PR. nemo_rl/models/megatron/data.py:2032 · PR-introduced; affects every Megatron sequence-packed run, not only CSD.

TL;DR: TE 2.18 now infers pad_between_seqs=True for every packed Megatron forward, which disables FlashAttention 2/4; the grpo-cispo-…-megatron-cispo nightly (packing + attention_backend: flash) will fail with "No dot product attention backend is available".

PackedSeqParams carries raw and padded boundaries; TE uses them to pick the attention backend.

1. The base passed the same cu_seqlens_padded object to all four fields (base L1832-L1836). This PR passes raw cu_seqlens for q/kv; with no padding cu_seqlens_padded is cu_seqlens.clone() (L2027), so it is always a distinct tensor.
2. MCore forwards pad_between_seqs=None (transformer_engine.py#L2314-L2340), so TE infers it by object identity: if cu_seqlens_q_padded is cu_seqlens_q and cu_seqlens_kv_padded is cu_seqlens_kv: pad_between_seqs = False (TE dot_product_attention.py#L1549-L1561).
3. With it True, TE disables FA2/FA4 and unfused attention for THD (TE utils.py#L976-L991).
4. grpo-cispo-…-megatron-cispo.yaml (nightly.txt#L101) sets packing + attention_backend: flash, which turns fused/unfused off (MLM utils.py#L63-L82); the default flash_attention_version=2 turns FA3 off (setup.py#L1906-L1909). Nothing is left: ValueError: No dot product attention backend is available (TE#L1687-L1693) on the first packed forward. On sm<90 with attention_backend: auto, packed runs silently move from FA2 to cuDNN.

Probe (real function body + TE's rule): base → False; head → True, even for [0,5,12,16] vs [0,5,12,16].

The raw boundaries look intentional (FullLogitsPostProcessor and _packed_row_geometry read them), so a revert doesn't work, and neither does pad_between_seqs=not torch.equal(...): TP2+SP (divisor 2) pads odd rows, so it is still True for cispo.

Action: keep the raw PackedSeqParams for NeMo-RL consumers, but hand the model a copy whose q/kv boundaries are the padded tensors, at both call sites: model_forward (train.py#L262-L264) and the draft/Eagle call (train.py#L475-L480). One way to do it:

def _te_packed_seq_params(params: Optional[PackedSeqParams]) -> Optional[PackedSeqParams]:
    """Packed boundaries as TE expects them: q/kv point at the padded tensors."""
    if params is None:
        return None
    return dataclasses.replace(
        params,
        cu_seqlens_q=params.cu_seqlens_q_padded,
        cu_seqlens_kv=params.cu_seqlens_kv_padded,
    )

Add a model_forward test asserting kwargs["packed_seq_params"].cu_seqlens_q is kwargs["packed_seq_params"].cu_seqlens_q_padded; the existing pins at test_megatron_data.py L2912-2913 / L3053 stay valid. Not GPU-verified; GPU numbers taken at this head went through cuDNN.

[avenkateshha]

1 action item, please fix in this PR. megatron_policy_worker.py:1607 · PR-introduced; hits every algorithm on Megatron.

The new skip of num_valid_samples <= 0 records can leave mb_losses == [] on a DP rank whose shard is all padding. all(...) over [] is True, so L1641-L1642 runs torch.stack([]) → RuntimeError: stack expects a non-empty TensorList. The base appended every record, and the Automodel worker already handles this case.

Breaking case (reproduced on CPU with the real maybe_pad_last_batch + shard_by_batch_size): Megatron SFT, DP=8, val_micro_batch_size=1, last validation batch of 3 samples (drop_last=False, sft.py#L574-L585) → padded to 8 rows with sample_mask=0 (utils.py#L558-L577) → ranks 3-7 hold only padding → every record skipped → crash. Same for DPO validation, and for GRPO when a DP shard is fully zeroed (grpo.py#L3697).

Action: at L1641 (outside the diff), so the else branch yields tensor(0.):

                if mb_losses and all(isinstance(loss, torch.Tensor) for loss in mb_losses):

plus a Megatron worker test with an all-sample_mask=0 DP shard.

[avenkateshha]

2 action items, please fix in this PR. csd.py:654-657 · PR-introduced.

TL;DR: on Automodel with CP>1 and dynamic_loss_scaling: true (the csd.yaml default, inherited by both Automodel CP=2 nightlies), the dynamic KD scale is built from a CP-partial KD, so the KD gradient is CP× too large relative to CE and kl_loss is logged at 1/CP.

1. On Automodel CP the worker multiplies the chunk count by cp_size (automodel_policy_worker.py#L517-L525), so each rank's cross-tokenizer KD is KD_true / CP (same-tokenizer KD is a per-window share); cp_normalized_by_caller is False there (loss_input.py#L395).
2. CE is forward-summed over CP, so every rank holds the full CE (csd.py#L1266-L1272).
3. This call gives kl_scale = |CE| / |KD_rank| = CP·|CE| / |KD_true|; after the DP×CP gradient sum (automodel/train.py#L604-L614) the KD gradient is CP× the intended one.
4. The owner-local path already uses the CP sum (_detached_cp_sum, L581-L582).

Probe (Automodel CP=2 vs a CP=1 reference): dynamic=True: ref scale=14.3198 | cp2 scale=28.6397 | max|grad diff|=5.065e-01. With only a same-tokenizer teacher the per-rank scales even diverge (8.76 / 12.20 vs 5.10). With the fix: same scale, grad diff ≤ 9e-8. Megatron is unchanged (_sum_window_kd_over_cp is True there). The nightlies only check isfinite, so they can't catch this.

AI-1

Use the CP-summed KD for the scale and the reported value (metric wiring in the next comment):

Suggested change

	# Combine the aggregated KD term with the single student CE term.
	loss, _, kl_scale = self._combine_kd_and_ce(
	total_kd, ce_loss, total_kd.detach(), ce_loss.detach()
	)
	# Combine the aggregated KD term with the single student CE term.
	# Without caller-side CP division (Automodel CP) each rank's KD is its
	# CP share while CE is complete, so the dynamic scale and the reported
	# KD use the CP sum, as on the owner-local path.
	reported_kd = (
	total_kd.detach()
	if self._sum_window_kd_over_cp
	else self._detached_cp_sum(total_kd, cp_group)
	)
	loss, reported_loss, kl_scale = self._combine_kd_and_ce(
	total_kd, ce_loss, reported_kd, ce_loss.detach()
	)

AI-2

Add a CP=2 gloo test with the Automodel contract (cp_normalized_by_caller=False) and dynamic scaling next to _run_same_vocab_megatron_cp in tests/unit/algorithms/csd/test_csd_loss_functions.py; the existing dynamic-scaling CP test only covers cp_normalized_by_caller=True.

Follow-up

Per-teacher kl_loss_t{i} / teacher_{i}/weighted_kl are also CP-partial on Automodel CP (logged at 1/CP); this fix doesn't cover them.

[avenkateshha]

Part of the csd.py:654 comment (same action). csd.py:677-678

Report the same CP-summed values the scale uses, so loss / kl_loss aren't logged at 1/CP on Automodel CP (only CP rank 0's metrics are kept, lm_policy.py#L1186-L1190):

Suggested change

	loss=loss,
	kd=total_kd,
	loss=reported_loss,
	kd=reported_kd,

[avenkateshha]

1 action item, please fix in this PR (L0 unit failure). tests/test_suites/nightly.txt:324 · PR-introduced.

The three CSD suites add 8 GPU-hours (2 + 2 + 4 under tools/launch's integer formula), taking the nightly total to 4955, while test_nightly_compute_stays_below_4947_hours asserts <= 4947. DRYRUN=1 ./tools/launch tests/test_suites/nightly.txt gives 4955 at head and 4947 for the base list, so the base passes exactly at the cap and this PR tips it over.

Action: feature PRs that add nightlies bump the cap in the same PR (#4458: 4830→4947; also #4526, #4341). Once the 100-step GPU runs (see the suite comment) fix NUM_MINUTES, update the constant, test name and message in tests/unit/test_recipes_and_test_suites.py, or keep one CSD suite out of nightly.txt until then.

[avenkateshha]

1 action item (low priority). examples/configs/csd.yaml:28 · PR-introduced.

With teacher_topk_ipc_keep_realized: true (the exemplar default), sparse KD adds each chunk's realized label to its support and lowers the mismatch minimum support from 2 to 1 (csd.py#L2618-L2623). So K=128 and K=0 optimize different losses even when K covers the whole vocabulary. Probe at K=V: dense 0.2297 vs sparse 0.3085, max |Δgrad| 0.037; with keep_realized: false they agree to 1e-9. The guide says the two "optimize different objectives" (guide#L543) but blames only the top-K cut, and this comment still reads as a pure transport choice:

Suggested change

	teacher_topk_ipc_k: 128 # 0 sends dense teacher logits
	teacher_topk_ipc_k: 128 # 0 sends dense teacher logits (a different KD objective; see the guide)

Do the same for csd_multiteacher.yaml L28. In the guide at L543, consider adding: "With teacher_topk_ipc_keep_realized: true the sparse KD also adds each chunk's realized label to its support, so the two differ even when K covers the vocabulary."

[avenkateshha]

1 action item. xtoken_off_policy_distillation.py:1444-1450 · PR-introduced.

The only caller is the CSD-only validate_csd_tokenizer_reuse (L1466). In CSD, a same-tokenizer teacher is one without a table pair; the projection path is null for CSD cross-tokenizer teachers too (L821-L825), so this message points users at the wrong key. The version below keeps the safe student-token reuse substring the test matches (test_distillation.py#L1872):

Suggested change

	raise ValueError(
	f"teachers[{teacher_idx}].aligner.projection_matrix_path is null, "
	"but safe "
	"student-token reuse requires identical tokenizer mapping, backend, "
	"special-token configuration, chat template, and template kwargs; "
	f"differing fields: {differing_fields}."
	)
	raise ValueError(
	f"teachers[{teacher_idx}].aligner has no CSD table pair, so it is "
	"treated as a same-tokenizer teacher, but safe "
	"student-token reuse requires identical tokenizer mapping, backend, "
	"special-token configuration, chat template, and template kwargs; "
	f"differing fields: {differing_fields}."
	)

[avenkateshha]

1 action item. csd.py:401-404 · PR-introduced. "Every teacher ships full-vocab logits" contradicts the teacher_topk_ipc_k: 128 default:

Suggested change

	# Per-teacher metadata: parallel lists (one entry per ``teachers[i]``),
	# injected by ``xtoken_off_policy_distillation.setup``. Every teacher ships
	# full-vocab logits (the loss derives the top-k subset student-side), so
	# there is no per-teacher ``send_full_logits`` flag.
	# Per-teacher metadata: parallel lists (one entry per ``teachers[i]``),
	# injected by ``xtoken_off_policy_distillation.setup``.

[avenkateshha]

1 action item. …-dtensor-tp4cp2.sh:12 · PR-introduced. This teacher uses CSD sub-token tables, not P-KL:

Suggested change

	TEACHER0_MODEL=Qwen/Qwen3-4B # cross-tokenizer (P-KL via projection)
	TEACHER0_MODEL=Qwen/Qwen3-4B # cross-tokenizer (CSD sub-token tables)

[avenkateshha]

1 action item. xtoken_off_policy_distillation.py:506 · PR-introduced.

TeacherAlignerConfig.projection_matrix_path's docstring (L500-L502, just above the diff) still says None marks a same-tokenizer teacher. For algorithm: csd, a null projection with a table pair is a cross-tokenizer teacher (L821-L825). Suggested wording: "X-Token: None marks a same-tokenizer teacher. CSD: the table pair below selects cross-tokenizer." The validate_csd_tokenizer_reuse docstring (L1458, "every null-projection teacher") needs the same fix.