Adds the chat/instruct alignment path to TokenAligner (Part 2 of the offset-alignment PR): align_one_offset_per_asst aligns each assistant message independently (isolate content, rebase offsets to a shared coordinate system, run the offset-cluster kernel on the slice, map positions back to the full sequence), with native reasoning/close/answer regions, a synthetic EOT pair, drop_first_content_pair, and included_regions. align_one_offset adds the single-sample offset path with an NFC decode-fix is_correct mask, and the constructor gains alignment_method (offset_cluster_decode_fix | same_tokenizer_identity). Returns raw 7-tuples; the chat collator (next commit) converts to AlignmentPair.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Adithya Hanasoge <avenkateshha@nvidia.com>

Adds a chat mode to CrossTokenizerCollator (Part 2 of the offset-alignment PR). In chat mode the collator renders each side's chat template, tokenizes with char offsets, derives the per-token assistant mask + each assistant message's char span, and aligns each assistant message independently via TokenAligner.align_one_offset_per_asst (multi-teacher). Loss fires only on assistant content (token_mask = attention_mask * assistant_mask). find_rendered_message_content_span (chat_templates.py) locates message content in the rendered string. Config knobs (collator_mode, drop_first_assistant_chunk_kl, include_thinking_in_loss, native_thinking_alignment, kd_alignment_regions, num_packed_rows) are wired from data_config and default to the raw-text path.

Scope: whole-message chat alignment. native_thinking_alignment (native reasoning/close/answer parsing) and num_packed_rows>1 (lockstep packing) raise NotImplementedError as clearly-marked follow-ups. Validated by unit tests (char-level fake chat tokenizer); a GPU chat smoke on real instruct data is a follow-up.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Adithya Hanasoge <avenkateshha@nvidia.com>

Keep cross-tokenizer teachers on the single offset-cluster/decode-fix path while same-tokenizer teachers continue to bypass alignment for direct positional KL.

Make the single-region helper private, reuse it from per-assistant alignment, and update focused coverage.

Signed-off-by: Adithyakrishna Hanasoge <avenkateshha@nvidia.com>

Add a typed per-teacher aligner block with centralized defaults for projection and chat alignment.

Preserve the released projection path through an explicit migration, route drop-first behavior by original teacher index, and update configs, docs, launchers, and coverage.

Signed-off-by: Adithyakrishna Hanasoge <avenkateshha@nvidia.com>

Use each teacher's assistant mask for chat scoring so prompts, template
scaffolding, and padding do not affect teacher selection or weighting.
Add regression coverage for assistant masks and score sensitivity.

Reject the legacy root projection path with an error naming the nested
aligner setting, and remove the guard for the unreleased shared setting.
Update config validation tests for the explicit rejection behavior.

Remove unused alignment partition masks, chunk counts, and the unused
alignment reconstruction helper, including their plumbing and tests.
Document both collator modes and every constructor argument, and correct
the explanation for pairing padding tokens.

Addresses the six scoped review comments on PR #3286.

Signed-off-by: Adithyakrishna Hanasoge <avenkateshha@nvidia.com>

Include retained assistant terminators in CE, same-tokenizer KD, and teacher
scoring masks. Locate template suffixes across whitespace without extending
content spans or duplicating cross-tokenizer EOT alignment pairs.

Move the five collator options into a typed top-level config with centralized
defaults, update both exemplars, and reject obsolete data-level settings with
migration guidance. Exercise CLI overrides through schema loading and setup.

Add a public align_chat API returning AlignmentBatch and use AlignmentPair
records throughout chat alignment. Keep batch assembly and padding cleanup
inside the aligner rather than the collator.

Add regression coverage for loss gradients, unchanged raw cross-tokenizer
KD, truncation, multiple turns, template whitespace, and typed public APIs.
Document chat configuration and terminator supervision.

Validation: 163 focused tests passed, 2 skipped; Ruff lint/import/format and
git diff checks passed. Repository-wide pyrefly reports 11 errors in untouched
modules, including missing optional dependencies.

Addresses review comments 2, 3, and 4 on PR #3286.

Signed-off-by: Adithyakrishna Hanasoge <avenkateshha@nvidia.com>

Preserve historical reasoning and structured tool calls during chat
collation, with independent message loss selection and typed regions.

Unify decoded offset matching, repair byte and Unicode boundaries against
source text, and bind ordinary chat spans to their logical turns. Retain
EOT supervision, same-tokenizer reuse, and per-teacher alignment settings.

Add a Qwen history-preserving template, native-chat recipe and nightly
driver, documentation, and regression coverage.

Validation: 352 focused cases passed across scoped runs; 2 skipped. The
existing DP multiprocessing test timed out in combined runs and passed
alone. Ruff and scoped Pyrefly pass; global Pyrefly retains 11 existing
errors.

Signed-off-by: Adithyakrishna Hanasoge <avenkateshha@nvidia.com>

Record the pinned tokenizer revision and fixture checksums as exact secret-scan baseline findings. Add the missing NVIDIA headers to the new data transport and Qwen template tests.

Signed-off-by: Adithyakrishna Hanasoge <avenkateshha@nvidia.com>

Exercise native chat alignment with SmolLM2-135M and Qwen3-1.7B using
different tokenizers, and assert unequal token counts in a valid aligned
chunk. Fold the three-step chat smoke into the existing functional test
and remove its standalone nightly driver and registration.

Generate the arithmetic conversations and projection matrix at runtime.
Load the stock Qwen template from its pinned tokenizer instead of keeping
a copied Jinja fixture and source manifest. Update documentation and the
secret-scan baseline for the moved and removed test assets.

Validation: 109 focused CPU tests passed, followed by the relocated config
integration recheck. Ruff lint/format, shell syntax, secret scanning, suite
accounting, and diff checks passed. GPU training was not run locally.

Addresses review comments 4140208831, 4140223825, and 4140278823 on #3286.

Signed-off-by: Adithyakrishna Hanasoge <avenkateshha@nvidia.com>

The functional xtoken native-chat config references its chat template
relative to the repo root, but CI runs pytest from tests/, so
get_tokenizer raised FileNotFoundError. Change into the repo root for
this test, matching how the functional suite launches.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Adithyakrishna Hanasoge <avenkateshha@nvidia.com>

Main renamed the dtensor_cfg block to automodel_cfg and removed the _v2
key, and reject_outdated_automodel_block now fails on either. Update the
native-chat functional config and the multi-teacher unit-test fixture
added on this branch.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Adithyakrishna Hanasoge <avenkateshha@nvidia.com>

Provide the checkpoint block required during teacher worker initialization after the recipe replaces the inherited teachers list.

Signed-off-by: Adithyakrishna Hanasoge <avenkateshha@nvidia.com>

Delete the guard for never-released data collator keys, its migration-message test, and the obsolete upgrade instructions.

Signed-off-by: Adithyakrishna Hanasoge <avenkateshha@nvidia.com>

…op-k primitives

Add three primitives for evaluating a loss on context-parallel student
shards and for exporting sparse teacher rows, all without gathering the
vocabulary or sequence axes:

- DistributedSelectedLogprobs / distributed_selected_logprobs: exact
  log-probabilities of selected token ids at selected rows of a
  TP-sharded logit tensor. Forward reduces only per-row max/sum scalars
  and the selected values; backward recomputes the local softmax shard
  and scatters the gradient back into the selected rows. Columns past
  the real vocabulary size are excluded from log Z and the gradient.
- cp_load_balanced_segments: a rank's load-balanced CP rows as
  (local_start, global_start, length) intervals, matching
  _get_tokens_on_this_cp_rank.
- distributed_vocab_topk_logz_force: per-row natural top-k sorted by
  token id, exact temperature-scaled log Z, the score-ranked k-th token
  id, and a forced-label sidecar (raw logits plus natural top-k
  membership) resolved across TP. Ties are broken by the lowest global
  token id so the selection does not depend on the TP layout.

Tests cover values and gradients against dense references, padded
LM-head columns, segment/rank-token agreement, cutoff ties, sidecar
validation, and gloo TP1/TP2 runs.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Adithya Hanasoge <avenkateshha@nvidia.com>

Add a Megatron sender for sparse teacher logits that keeps rows in their
context-parallel shard:

- SparseTeacherLogitsPostProcessor computes, per teacher row and across
  TP only, the natural top-k sorted by token id, the exact
  temperature-scaled log Z over the real vocabulary, the score-ranked
  k-th token id, and a sidecar for up to F forced labels (raw logits and
  natural top-k membership). Forced ids come from the teacher data's
  force_include_token_ids [B, T_t, F] and are sharded with the same
  load-balanced CP layout as the logits. It is dispatched without
  sampling-temperature scaling.
- MegatronPolicyWorker.get_sparse_logits_ipc writes the seven fields
  into persistent [N_mb, B, T_local, width] slabs whose IPC handles are
  cached across calls, sized once for max_total_sequence_length.
  Only TP rank 0 on the last pipeline stage exports; each record lists
  the global intervals its rows cover in sequence_segments. Models that
  pack and context-shard inside their own forward are rejected at
  CP > 1. release_ipc_buffer also frees the sparse slabs.
- Policy.get_sparse_logits_ipc routes the call to Megatron workers and
  keeps one record per teacher CP shard and sample.
- ensure_teacher_ipc_buffer takes an optional capacity_seq_len.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Adithya Hanasoge <avenkateshha@nvidia.com>

Add readers that copy only the teacher rows a loss asks for, instead of
reassembling a teacher's full sequence on every GPU:

- SparseTeacherRowReader / gather_sparse_teacher_rows_from_ipc route
  each (sample, global position) to the teacher CP shard that holds it
  through the records' sequence_segments, open each backing slab once,
  and return one exact K-wide support per request: the natural top-k
  when the requested teacher label is in it, otherwise the natural
  top-(k-1) plus that label's forced logit (the recorded k-th id is
  evicted). Each request also gets the exact log Z and the label's
  natural top-k membership.
- DenseTeacherRowReader reads full-vocabulary rows from the existing
  dense handles (compact or rectangular), stitching the teacher's TP
  vocabulary shards; rows in a compact handle's omitted suffix read as
  zero. _rebuild_teacher_ipc_row now shares its row slicing with it.
- Both readers cache opened storage until close().
- _build_teacher_force_include_token_id_slots builds the
  [B, T_t, 2] forced labels the Megatron sender needs: every predictor
  row of every valid chunk keeps its realized teacher label, and a row
  shared by an unshifted chunk at the sequence start and the next
  shifted chunk keeps both labels.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Adithya Hanasoge <avenkateshha@nvidia.com>

Add an owner-local CSD evaluation for Megatron students, selected by the
runtime-injected CSDLossConfig.owner_local flag (set by setup in a
later commit). The student logits stay [B, S_pad/CP, V/TP] in their
load-balanced CP layout and are never gathered across CP:

- prepare_csd_owner_local_loss_input hands each cross-tokenizer teacher
  to the loss as a TeacherRowPayload (its sparse or dense IPC records)
  with global spans, full ids and the global position of every local
  logit row. A same-tokenizer teacher keeps today's contiguous-window
  reconstruction, built only when one is configured.
- _compute_csd_loss_owner_local classifies every chunk identically on
  all ranks; the rank holding a chunk's final student predictor row
  evaluates it, reading only that chunk's teacher rows through the
  sparse or dense row reader and the student values through
  distributed_selected_logprobs. A mismatch chunk's student prefix
  rows may sit on other CP ranks; each rank adds its share to one
  fixed-length vector summed across CP with a SUM backward. Dense
  teachers follow the replicated support rules; sparse teachers use
  each request's top-k with its realized label kept. Before the first
  collective, a MIN-reduced ok flag over TP then CP makes every rank
  raise together when one rank's setup fails.
- The loss is CP-partial: KD is the sum over owned chunks divided by
  the step-global valid-chunk count (required on this path), CE sums
  -log p over the rank's own rows divided by global_valid_toks, and the
  same-tokenizer KD is already a per-window partial. Dynamic loss
  scaling and the reported metrics use detached CP sums, so every rank
  applies the same ratio and reports complete values. Accuracy uses the
  rank's rows and scalar CP reductions.
- The Megatron worker builds LossPostProcessor with cp_normalize=False
  for such a loss (loss_returns_context_parallel_partial); the
  num_microbatches / cp_size schedule compensation stays. Megatron sums
  gradients over DP x CP, so the partials give gradient factor 1.
- owner_local_config_violations lists unsupported settings
  (kd_loss_mode other than sum, dynamic weights, pure ALM, BCE on
  non-binary partitions, sparse IPC without kept realized labels);
  CSDLossFn rejects them when owner_local is set.
- New knobs csd_selected_row_chunk_size (128) and
  csd_support_chunk_batch_size (16) bound memory only.

Gloo tests over TP/CP in {1, 2}^2 check that the CP-summed loss and the
reassembled gradient equal the replicated reference (dense) or the
TP1/CP1 owner-local result (sparse, itself checked against the
replicated sparse path), in both mismatch modes, with a noise filter,
fixed weights, an added same-tokenizer teacher, and a rank that owns no
chunk; that no CP gather runs without a same-tokenizer teacher; and
that a setup error on one rank raises on every rank.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Adithya Hanasoge <avenkateshha@nvidia.com>

Setup now decides, per CSD run, whether cross-tokenizer KD runs
owner-local and injects the result as loss_fn.owner_local:

- csd_owner_local_ineligibility lists why a run cannot: a student that
  is not Megatron, student PP > 1, sequence packing, dynamic batching,
  the loss settings from owner_local_config_violations, and, with
  sparse teacher IPC, cross-tokenizer teachers that are not Megatron or
  whose make_sequence_length_divisible_by is not a multiple of 2 x CP
  (and TP with sequence parallelism).
- A Megatron student with no reasons runs owner-local, with dense or
  sparse teacher rows. With reasons it keeps the replicated loss, which
  only dense IPC can feed, so teacher_topk_ipc_k > 0 is rejected with
  the reasons and a hint to set it to 0. An Automodel student keeps the
  replicated path and still requires Automodel sparse teachers.
- Owner-local runs export cross-tokenizer sparse teachers through
  get_sparse_logits_ipc with the two-slot forced labels; dense export
  is unchanged.
- Setup prints which evaluation the run uses and why.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Adithya Hanasoge <avenkateshha@nvidia.com>

On the replicated Megatron path at CP > 1, a same-tokenizer teacher's
KD term covers only this rank's contiguous window (a per-window partial
over the global KD-token count), but the loss wrapper then divides the
whole loss by the CP size. The term was therefore under-weighted by
1/CP, and every rank reported only its own window's share.

prepare_loss_input now marks the Megatron replicated path
(cp_normalized_by_caller, set when there is no model-owned CP sharder),
and CSDLossFn sums the same-tokenizer KD term across CP with a SUM
backward there: every rank holds the full term, and after the /CP
division the gradient factor is 1. The Automodel path and the
owner-local path, which do not divide by CP, are unchanged.

A gloo test checks that loss, metrics and reassembled gradient at CP2
match CP1.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Adithya Hanasoge <avenkateshha@nvidia.com>

Set CSDLossConfig.teacher_topk_ipc_k to 128 and the CSD exemplars
(csd.yaml, csd_multiteacher.yaml) to match, and update the guide.

This changes the objective, not just the transport: a cross-tokenizer
teacher row now contributes its top-128 logits plus the exact REST mass
instead of the full vocabulary. Megatron students run it owner-local
with the Megatron sparse sender; Automodel students keep the replicated
sparse path. Set loss_fn.teacher_topk_ipc_k=0 for the previous dense
transport; CSD lockstep packing still requires it and says so at setup.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Adithya Hanasoge <avenkateshha@nvidia.com>

Add a "Teacher transport and context parallelism" section to the
cross-tokenizer distillation guide: sparse and dense teacher transport,
how owner-local evaluation splits the loss across context-parallel
ranks, the student x teacher x teacher_topk_ipc_k matrix, the
owner-local requirements, and memory notes (no student gathers, row
reads, the two tiling knobs, and the persistent sparse buffers whose
reused IPC handles may not be released after an error or at shutdown).
The lockstep packing design no longer says Megatron sparse IPC is
missing.

Remove the unused CSDLossFn._rebuild_teacher_full_logits and
_rebuild_teacher_sparse_logits helpers and their test; the loss-input
preparation rebuilds or reads teacher logits.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Adithya Hanasoge <avenkateshha@nvidia.com>

The setup check that each sequence-length divisor is a multiple of the
context-parallel size read every entity's CP from automodel_cfg and
detected same-vocab teachers by a null projection path. In the shared
trainer a Megatron student or teacher keeps its CP in megatron_cfg and
pads its own inputs to lcm(divisor, 2 * CP), so a Megatron entity with
divisor 1 and CP=2 is valid, and CSD marks same-vocab teachers by a
null sub-token table path instead.

Skip Megatron entities, which pad themselves, and detect same-vocab
teachers with _teacher_is_cross_tokenizer, so the check covers both
algorithms.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Adithya Hanasoge <avenkateshha@nvidia.com>

Rename docs/guides/cross-tokenizer-distillation.md to
csd-cross-tokenizer-distillation.md and rewrite it as a CSD guide that
mirrors the X-Token guide: how CSD compares the models per chunk type,
backends and teacher transports, step-by-step sub-token table building
and validation, launch, loss knobs, multiple teachers and lockstep
packing.

Add a README "Cross-Tokenizer Distillation" section that introduces both
methods, keeps the X-Token subsection (and its anchor), and adds a CSD
subsection with table preparation and single- and multi-node launches.
Add CSD to the distillation recipe table, and give X-Token and CSD
separate cards in docs/index.md.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Adithya Hanasoge <avenkateshha@nvidia.com>

Keep a single cross-tokenizer distillation guide at
docs/guides/xtoken-off-policy-distillation.md: a shared introduction
with the tokenizer-overlap and X-Token/CSD comparison tables, then an
X-Token section and a CSD section, and combined "Where files live" and
"Related" lists. Remove docs/guides/csd-cross-tokenizer-distillation.md
and point the README, docs index and CSD exemplar comments at the CSD
section.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Adithya Hanasoge <avenkateshha@nvidia.com>

The guide now covers both X-Token and CSD, so rename
docs/guides/xtoken-off-policy-distillation.md to
docs/guides/cross-tokenizer-distillation.md. Point the README, docs
index, exemplar config comments and build_projection_matrix.sh at the
new path, and link the guide from the README's cross-tokenizer section.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Adithya Hanasoge <avenkateshha@nvidia.com>

CI's copyright-check job reports tests/unit/tools/csd/test_subtoks_tables.py
as the only Python file without the NVIDIA Apache-2.0 header. Add it.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Adithya Hanasoge <avenkateshha@nvidia.com>

…ests

- test_train.py: give test_model_forward_passes_cp_local_router_padding_mask
  a config with sequence_parallel=False. A bare MagicMock config is
  truthy, so model_forward sent the mask through the sequence-parallel
  scatter, which needs an initialized TP group.
- test_megatron_setup.py: unset HF_HUB_OFFLINE in
  test_megatron_fsdp_uses_local_hf_snapshot. validate_model_paths
  derives local_files_only from it, so the test failed wherever
  HF_HUB_OFFLINE=1 is exported.
- test_automodel_train.py: mark TestFullLogitsPostProcessorSeqPacking
  with @pytest.mark.automodel. It was the only unmarked class in the
  module, so no L0 lane ran it.
- test_megatron_data.py: run the unpacked-CP process_microbatch test
  through the real _get_tokens_on_this_cp_rank and
  get_ltor_masks_and_position_ids, and check CP rank 1's shard values.
  The old version mocked the shard helper with made-up shards and still
  passed when (cp_rank, cp_size) were swapped.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Adithya Hanasoge <avenkateshha@nvidia.com>

The unpacked context-parallel branch of process_microbatch sharded
input_ids and position_ids but left the nonpacked router padding mask at
full sequence length. With expert-bias MoE
(create_nonpacked_router_padding_mask=True), CP>1 and packing off,
GPTModel's padding_mask.shape == input_ids.shape assert fails (CP rank 1
of 2: (1, 4) vs (1, 8)). Shard the mask with the same load-balanced
layout.

Add a test that runs process_microbatch on CP rank 0 of 2 with the real
shard helper and checks the mask's shape and values. It fails without
the fix.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Adithya Hanasoge <avenkateshha@nvidia.com>

csd_multiteacher.yaml pointed the Phi-4-mini teacher at
cross_tokenizer_data/csd/{student_to_teacher,teacher_to_student}.pt, the
Llama-3.2-1B / Qwen3-4B pair that csd.yaml and the guide's main flow
build. Every Qwen3 id is also a valid Phi-4-mini id, so running the
exemplar after the main flow trained on the wrong mapping with no error.
Default the paths to cross_tokenizer_data/csd/phi4mini/, the location the
guide already tells users to build the Phi-4-mini tables at.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Adithya Hanasoge <avenkateshha@nvidia.com>

With CP>1 and packing off, process_microbatch built the dense
[1, 1, S, S] causal mask through get_ltor_masks_and_position_ids and
then dropped it, since TE's context-parallel attention applies the
causal mask itself. Megatron allocates it as float32 plus a bool copy,
about 5 * S^2 bytes per microbatch: ~5 GiB at 32k tokens and ~80 GiB at
128k, the lengths CP is for. Pass create_attention_mask=cp_size == 1 so
the mask is only built when the model uses it; CP=1 is unchanged.

Add a test that checks the mask is requested for CP=1 and not for CP=2.
The CP=2 case fails without the fix.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Adithya Hanasoge <avenkateshha@nvidia.com>

policy.freeze_parameter_patterns was only read by the Automodel setup,
so on a Megatron policy (as csd.yaml configures) a pattern such as
["mtp"] was silently ignored and the matching parameters kept training.
Port the MCore hook from Mingyu Yang's Cascade SFT-KD branch (70bb51230):
setup_model_and_optimizer appends a pre-wrap hook that sets
requires_grad=False on every parameter whose name contains a pattern,
using the Automodel backend's substring rule. It runs after the built-in
and caller-supplied hooks, so parameters added by PEFT or draft hooks are
covered, and before DDP and the optimizer are built. freeze_moe_router
now returns the model it receives, like the other hooks, so the list can
be chained.

Document that both backends support the key, and port the tests: the
hook freezes literal substring matches across model chunks, leaves a
model without matches unchanged, and setup_model_and_optimizer runs it
last with every hook returning the model.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Adithya Hanasoge <avenkateshha@nvidia.com>

reduce_mb_metric averages only the keys in MEAN_REDUCED_MB_METRICS and
sums everything else over microbatches and DP ranks. weight_t{i} and
selected_teacher are per-microbatch values, not shares of a global
total, so they were summed: a static teacher weight of 1.0 over 64
microbatches logged as weight_t0 = 64.0. Add both to the mean-reduced
set. The X-Token multi-teacher loss emits the same keys through the same
trainer, so its logged values are corrected too. Logging only; the loss
and gradients are unchanged.

Extend test_reduce_mb_metric_means_per_chunk_diagnostics_and_sums_shares
with weight_t{i} and selected_teacher; it fails (64.0 != 1.0) without
the fix.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Adithya Hanasoge <avenkateshha@nvidia.com>

…trainer

config-conventions keeps a config value's default in one place, never at
the call site. In xtoken_off_policy_distillation.py:
- offload_optimizer_for_logprob is a required PolicyConfig bool, so read
  it with policy["offload_optimizer_for_logprob"] instead of
  .get(..., False).
- offload_student_after_step is NotRequired, so read it with .get(key)
  instead of .get(key, False); an absent key still means off.
- validate_csd_router_replay_setup and the lockstep packing check used
  .get("megatron_cfg") or {} plus .get("enabled", False); use the
  existing _megatron_enabled helper.

Behaviour is unchanged for valid configs; a policy config that omits
offload_optimizer_for_logprob now fails with a KeyError instead of being
treated as False. Add the key to the two test fixtures that omitted it
(_make_master_config in csd/test_distillation.py and the master config in
x_token/test_xtoken_off_policy_distillation.py).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Adithya Hanasoge <avenkateshha@nvidia.com>

With clip_overlength_text, the CSD collator clips each text row that
_fits_every_side says does not fit, and the shared collator then rejects
any row that is still overlength. A same-tokenizer teacher reuses the
student batch, which is padded to the student's divisor (packed rows
pad to the teacher's), so the shared reject path checks its context
length against both divisors. _fits_every_side checked only the
teacher's divisor, so a row could pass the clip check and then raise in
the DataLoader worker: with a student divisor of 4 and a teacher context
of 6, a 5-token row pads to 8 and is rejected instead of clipped.

Check same-tokenizer teachers against both divisors, as the reject path
does. Add the student-rounding case to
test_text_overlength_rows_are_clipped_to_a_shared_prefix; it fails
without the fix with the shared collator's "padded to 8 ... exceeding
context length 6" error.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Adithya Hanasoge <avenkateshha@nvidia.com>