`nemo_rl.algorithms.x_token.loss_utils`#

Shared utilities for cross-tokenizer distillation.

Used by both :mod:token_aligner and

mod:: nemo_rl.algorithms.loss.loss_functions:

class:

Fp32SparseMM — FP32 sparse-dense matmul that ignores BF16 autocast (no BF16 sparse-mm kernel exists).
Chunk aggregation: :func:chunk_log_prob_sums / :func:chunk_average_finalize / :func:chunk_average_log_probs / :func:valid_chunk_mask (the partial/finalize split lets callers insert a CP all-reduce between), plus

func:

nemo_rl.distributed.model_utils.group_all_reduce_sum for the global valid-chunk denominator.
Teacher-logit IPC: :func:rebuild_teacher_full_logits_from_ipc,

func:

assemble_teacher_logits_from_shards,

func:

collect_overlapping_teacher_shards reassemble full-vocab teacher logits from per-rank shards across heterogeneous TP/CP.
Projection: :func:parse_projection_file, the

func:

get_sparse_projection_matrix / :func:get_topk_projection process-local caches, :func:slice_sparse_projection_rows, and

func:

build_exact_token_map (cached common/uncommon partition).
func:

alignment_from_flat_batch rehydrates the flat alignment_* data-dict keys into an :class:AlignmentBatch.

Module Contents#

Classes#

`Fp32SparseMM`	FP32 `M.t() @ dense` (sparse-dense matmul) ignoring surrounding autocast.
`LocalizedAlignment`	CP-localized alignment tensors consumed by the loss reductions.

Functions#

`alignment_from_flat_batch`	Rebuild :class:`AlignmentBatch` from the flat `alignment_*` keys.
`chunk_log_prob_sums`	Local bmm + bucket count, no division.
`chunk_average_finalize`	Divide sums by sizes; `eps` guards empty buckets.
`chunk_average_log_probs`	Average `log_probs` over chunks defined by `chunk_id`.
`slice_sparse_projection_rows`	Row-slice a sparse-COO projection `[V_s, V_t]` to `[row_end-row_start, V_t]`.
`project_student_to_teacher_vocab`	Project student vocab probs `[B, T, V_s(/TP)]` to teacher vocab `[B, T, V_t]`.
`select_teacher_topk_indices`	Sorted global top-`k` teacher-vocab ids by max importance over the microbatch.
`localize_alignment`	Localize the chunk-alignment data-dict fields for the local CP shard.
`student_next_token_ce`	Per-token next-token cross-entropy `[B, T-1]` on the student.
`ce_label_mask`	Next-token label mask `[B, ce_seq_len]` = shifted token_mask * sample_mask.
`next_token_accuracy`	Masked next-token top-1 accuracy of the student (scalar, no gradient).
`collect_overlapping_teacher_shards`	Plan `(src_seq, src_vocab, dest_seq, dest_vocab)` slices per teacher shard.
`assemble_teacher_logits_from_shards`	P2P-IPC-read overlapping teacher shards into a `[T_t/CP_s, V_t]` dest.
`_try_zero_copy_teacher_logits`	Zero-copy `[B, T_t/CP_s, V_t]` view of the teacher logits, or None.
`rebuild_teacher_full_logits_from_ipc`	Rebuild `[B, T_t/CP_s, V_t]` teacher logits for this student rank.
`valid_chunk_mask`	Per-chunk validity gate: both sides non-empty and pair is valid.
`parse_projection_file`	Parse a projection-matrix file into COO components.
`get_sparse_projection_matrix`	Return the sparse-COO projection matrix on `device` (cached).
`get_topk_projection`	Return the dense top-k `(indices, likelihoods)` projection on `device` (cached).
`build_exact_token_map`	Build the common/uncommon vocab partition for the gold path (cached).
`prepare_xtoken_cross_tokenizer_loss_input`	Build the per-teacher cross-tokenizer distillation loss pieces from student logits + IPC teacher data.

Data#

`_SPARSE_PROJECTION_CACHE`
`_TOPK_PROJECTION_CACHE`
`_EXACT_TOKEN_MAP_CACHE`

API#

nemo_rl.algorithms.x_token.loss_utils.alignment_from_flat_batch( data: Mapping[str, Any], ) → nemo_rl.algorithms.x_token.token_aligner.AlignmentBatch#

Rebuild :class:AlignmentBatch from the flat alignment_* keys.

The field set is driven off :class:AlignmentBatch so the helper can’t drift from the schema.

class nemo_rl.algorithms.x_token.loss_utils.Fp32SparseMM#

Bases: torch.autograd.Function

FP32 M.t() @ dense (sparse-dense matmul) ignoring surrounding autocast.

addmm_sparse_cuda has no BF16 kernel on either forward or backward. The worker wraps forward + loss + backward in autocast(BF16), so a plain with autocast(enabled=False): around the forward call is not enough — loss.backward() runs inside the outer autocast and the sparse-mm backward kernel is still dispatched as BF16. The custom_fwd(cast_inputs=torch.float32) / custom_bwd decorators are PyTorch’s official escape: they force FP32 inputs on forward and run the backward as if autocast were disabled.

autograd’s builtin sparse-mm backward computes M @ grad_out. The gradient w.r.t. the sparse argument isn’t needed (the projection matrix is frozen), so it’s returned as None.

static forward( ctx: Any, sparse_M: torch.Tensor, dense: torch.Tensor, ) → torch.Tensor#

static backward( ctx: Any, grad_out: torch.Tensor, ) → tuple[None, torch.Tensor]#

nemo_rl.algorithms.x_token.loss_utils.chunk_log_prob_sums( log_probs: torch.Tensor, chunk_id: torch.Tensor, max_chunks: int, ) → tuple[torch.Tensor, torch.Tensor]#

Local bmm + bucket count, no division.

Output is summable across CP; callers that need cross-rank chunks to aggregate correctly should group_all_reduce_sum_with_grad both tensors before :func:chunk_average_finalize. chunk_id == -1 contributes to no bucket.

nemo_rl.algorithms.x_token.loss_utils.chunk_average_finalize( chunk_sums: torch.Tensor, chunk_sizes: torch.Tensor, ) → tuple[torch.Tensor, torch.Tensor]#: Divide sums by sizes; eps guards empty buckets.

nemo_rl.algorithms.x_token.loss_utils.chunk_average_log_probs( log_probs: torch.Tensor, chunk_id: torch.Tensor, max_chunks: int, *, cp_group: Optional[torch.distributed.ProcessGroup] = None, ) → tuple[torch.Tensor, torch.Tensor]#

Average log_probs over chunks defined by chunk_id.

Builds a one-hot chunk mask from chunk_id (-1 = no chunk), then bmm-aggregates and divides by chunk sizes. When cp_group has world

1, the per-chunk sums are group_all_reduce_sum_with_grad’d across CP ranks before the divide (mean is non-linear, so the reduce must precede it).

Parameters:

log_probs – [B, T, V] log-probabilities.
chunk_id – [B, T] long tensor, values in [-1, max_chunks).
max_chunks – number of chunk buckets.
cp_group – context-parallel group for cross-rank chunk aggregation.

Returns:

[B, max_chunks, V] averaged log-probs. chunk_sizes: [B, max_chunks] float tensor of bucket sizes.

Return type:

chunk_log_probs

nemo_rl.algorithms.x_token.loss_utils.slice_sparse_projection_rows( sparse_matrix: torch.Tensor, row_start: int, row_end: int, ) → torch.Tensor#

Row-slice a sparse-COO projection [V_s, V_t] to [row_end-row_start, V_t].

Filters COO indices in-place: keeps entries with row in [row_start, row_end) and shifts the row index by -row_start. Used by the TP-aware P-KL path where each rank owns a contiguous slab of the student vocab axis.

nemo_rl.algorithms.x_token.loss_utils.project_student_to_teacher_vocab( student_probs: torch.Tensor, sparse_projection: torch.Tensor, *, tp_group: Optional[torch.distributed.ProcessGroup] = None, ) → torch.Tensor#

Project student vocab probs [B, T, V_s(/TP)] to teacher vocab [B, T, V_t].

sparse_projection is the full [V_s, V_t] sparse-COO matrix. With tp_group world > 1 the student probs cover only this rank’s V_s/TP rows, so the matrix is row-sliced to that range, the sparse matmul produces a partial teacher-vocab sum, and a group_all_reduce_sum_with_grad over the TP group combines the partials into the full V_s contraction. Otherwise a single sparse matmul over the full matrix is used.

nemo_rl.algorithms.x_token.loss_utils.select_teacher_topk_indices( teacher_logits: torch.Tensor, k: int, *, cp_group: Optional[torch.distributed.ProcessGroup] = None, ) → torch.Tensor#

Sorted global top-k teacher-vocab ids by max importance over the microbatch.

Importance is the per-vocab max over flattened (B*T) teacher logits. With cp_group world > 1 the sequence is CP-sharded, so the local max only sees this rank’s slice; an all_reduce(MAX) makes every rank pick the same subset. No gradient.

class nemo_rl.algorithms.x_token.loss_utils.LocalizedAlignment#

CP-localized alignment tensors consumed by the loss reductions.

For a cross-tokenizer teacher every field is populated (chunk-averaged projection KL / gold path). For a same-tokenizer teacher (no projection, identity 1:1 token alignment) the chunk/pair fields stay None: its KD term reads only the shared student fields (student_input_ids / student_token_mask / sample_mask).

sample_mask: torch.Tensor#: None

student_chunk_id: Optional[torch.Tensor]#: None

teacher_chunk_id: Optional[torch.Tensor]#: None

pair_valid: Optional[torch.Tensor]#: None

pair_is_correct: Optional[torch.Tensor]#: None

student_input_ids: Optional[torch.Tensor]#: None

student_token_mask: Optional[torch.Tensor]#: None

nemo_rl.algorithms.x_token.loss_utils.localize_alignment( data: Mapping[str, Any], *, teacher_seq_len: int, alignment_prefix: str = 'alignment_', cp_group: Optional[torch.distributed.ProcessGroup] = None, ) → nemo_rl.algorithms.x_token.loss_utils.LocalizedAlignment#

Localize the chunk-alignment data-dict fields for the local CP shard.

Unwraps the {alignment_prefix}* / sample_mask entries from DTensor to their local tensors. Student-seq fields (student_chunk_id) come from cp_buffers in PyTorch’s load-balanced (2*cp interleaved) CP layout, not contiguous — the caller (:func:prepare_xtoken_cross_tokenizer_loss_input) must relayout them to this rank’s contiguous window via :func:cp_load_balanced_to_contiguous before use. The teacher-seq teacher_chunk_id is full, so it is sliced contiguously to this CP rank’s teacher_seq_len window to match the IPC consumer’s contiguous teacher-logit slice.

Parameters:: alignment_prefix – Data-dict key prefix for this teacher’s alignment tensors ("alignment_" single-teacher, "alignment_{i}_" per teacher in the multi-teacher trainer / collator). sample_mask is student-level and stays unprefixed.

nemo_rl.algorithms.x_token.loss_utils.student_next_token_ce( logits: torch.Tensor, *, input_ids: torch.Tensor, seq_index: Optional[torch.Tensor] = None, ) → torch.Tensor#

Per-token next-token cross-entropy [B, T-1] on the student.

DTensor (TP/CP) logits route through the vocab-parallel log-prob helper (which also handles the CP roll); plain logits use a local shifted cross_entropy. The next-token shift (drop the last predictor) matches the convention the KL terms use.

nemo_rl.algorithms.x_token.loss_utils.ce_label_mask( *, token_mask: torch.Tensor, sample_mask: torch.Tensor, ce_seq_len: int, dtype: torch.dtype, ) → torch.Tensor#

Next-token label mask [B, ce_seq_len] = shifted token_mask * sample_mask.

token_mask is gathered to the full sequence (CP) before the shift; both inputs are DTensor-unwrapped.

nemo_rl.algorithms.x_token.loss_utils.next_token_accuracy( logits: torch.Tensor, *, input_ids: torch.Tensor, token_mask: torch.Tensor, sample_mask: torch.Tensor, tp_group: Optional[torch.distributed.ProcessGroup] = None, cp_group: Optional[torch.distributed.ProcessGroup] = None, ) → torch.Tensor#

Masked next-token top-1 accuracy of the student (scalar, no gradient).

Uses :func:vocab_parallel_argmax for the (possibly TP-sharded) argmax. The next-token shift on labels/mask is CP-aware (:func:cp_shift_next) so the boundary token crosses CP ranks, and the correct/total counts are CP-reduced so every rank reports the same global accuracy.

nemo_rl.algorithms.x_token.loss_utils.collect_overlapping_teacher_shards( teacher_shards: list[dict[str, Any]], student_cp_rank: int, student_cp_size: int, full_seq_len: int, ) → list[tuple[dict[str, Any], slice, slice, slice, slice]]#

Plan (src_seq, src_vocab, dest_seq, dest_vocab) slices per teacher shard.

Dest is [T_t/CP_s, V_t] (vocab fully reassembled, seq is this student CP rank’s range). Shards with no seq overlap are skipped.

nemo_rl.algorithms.x_token.loss_utils.assemble_teacher_logits_from_shards( teacher_shards: list[dict[str, Any]], student_cp_rank: int, student_cp_size: int, device: int, ) → torch.Tensor#

P2P-IPC-read overlapping teacher shards into a [T_t/CP_s, V_t] dest.

device is a CUDA device index (matches

Func:: rebuild_cuda_tensor_from_ipc’s device_id signature).

nemo_rl.algorithms.x_token.loss_utils._try_zero_copy_teacher_logits( per_sample_entries: list[dict[str, Any]], *, student_cp_rank: int, student_cp_size: int, device: int, ) → Optional[torch.Tensor]#

Zero-copy [B, T_t/CP_s, V_t] view of the teacher logits, or None.

Returns a view into the producer’s IPC storage only when reassembly is unnecessary: every sample’s seq range is covered by a single full-vocab teacher shard (i.e. teacher tp_size == 1 and teacher_cp == student_cp or teacher_cp == 1), and the microbatch’s samples are a contiguous slab (same payload + buf_idx, sample rows 0..B-1) in one storage slot. Otherwise returns None and the caller falls back to assemble + stack.

nemo_rl.algorithms.x_token.loss_utils.rebuild_teacher_full_logits_from_ipc( per_sample_entries: list[dict[str, Any]], cp_group: Optional[torch.distributed.ProcessGroup], device: int, ) → torch.Tensor#

Rebuild [B, T_t/CP_s, V_t] teacher logits for this student rank.

Fast path (zero-copy view via :func:_try_zero_copy_teacher_logits): when the teacher is not vocab-sharded and each sample’s seq range is covered by a single shard, return a view into the IPC storage. Otherwise reassemble each sample from its overlapping shards and stack.

nemo_rl.algorithms.x_token.loss_utils.valid_chunk_mask( s_sizes: torch.Tensor, t_sizes: torch.Tensor, pair_valid: torch.Tensor, ) → torch.Tensor#: Per-chunk validity gate: both sides non-empty and pair is valid.

nemo_rl.algorithms.x_token.loss_utils.parse_projection_file( path: Union[str, os.PathLike], ) → Tuple[torch.Tensor, torch.Tensor, int, int]#

Parse a projection-matrix file into COO components.

Detects either the dense top-k format (dict["indices"] / dict["likelihoods"]) or the sparse multi-token format (dict[(student_id, teacher_id)] -> count) and converts both to a uniform COO representation.

The function does not apply any sizing or validity policy: the -1 sentinel used by _exact_map_remapped projection files is preserved in the returned indices, and the inferred vocab sizes are derived from the file alone (caller may override them upward against tokenizer / config knowledge). This keeps a single parser while letting :mod:token_aligner and the loss fn keep their own clipping rules.

Parameters:

path – Path to a torch.saved projection-matrix file.

Returns:

LongTensor[2, nnz] — (student_idx, teacher_idx). values: FloatTensor[nnz]. v_student_inferred: int — dense format: row count; sparse format: max(student_idx) + 1. v_teacher_inferred: int — max(positive teacher_idx) + 1 (0 if no positive entries exist).

Return type:

indices

Raises:

FileNotFoundError – path does not exist.
ValueError – the file is not in a recognized format.

nemo_rl.algorithms.x_token.loss_utils._SPARSE_PROJECTION_CACHE: dict[Tuple[str, torch.device, int, int], torch.Tensor]#: None

nemo_rl.algorithms.x_token.loss_utils._TOPK_PROJECTION_CACHE: dict[Tuple[str, torch.device], Tuple[torch.Tensor, torch.Tensor]]#: None

nemo_rl.algorithms.x_token.loss_utils.get_sparse_projection_matrix( path: Union[str, os.PathLike], device: torch.device, *, student_vocab_size: int, teacher_vocab_size: int, ) → torch.Tensor#

Return the sparse-COO projection matrix on device (cached).

On a cache miss, parses the file via :func:parse_projection_file, drops -1 teacher sentinels (illegal in sparse-COO), sizes V_s = max(student_vocab_size, max_observed_student_idx + 1) and V_t = max(teacher_vocab_size, max_observed_teacher_idx + 1), and builds a coalesced torch.sparse_coo_tensor on device. Subsequent calls with the same (path, device, student_vocab_size, teacher_vocab_size) return the cached tensor — no disk I/O, no re-materialization.

Both vocab sizes are keyword-only to prevent a positional swap (two same-magnitude ints, no error if confused).

Parameters:

path – Path to a torch.saved projection-matrix file.
device – Device the sparse tensor must live on.
student_vocab_size – Minimum width of the student-side axis.
teacher_vocab_size – Minimum width of the teacher-side axis.

Returns:

torch.sparse_coo_tensor of shape (V_s, V_t), coalesced, dtype=float32.

nemo_rl.algorithms.x_token.loss_utils.get_topk_projection( path: Union[str, os.PathLike], device: torch.device, ) → Tuple[torch.Tensor, torch.Tensor]#

Return the dense top-k (indices, likelihoods) projection on device (cached).

Used by the gold-loss exact-map builder, which needs the per-row top-k weights — the sparse dict[(s, t)] -> count projection format doesn’t carry those, so this loader rejects it.

Parameters:

path – Path to a torch.saved projection-matrix file.
device – Device the returned tensors must live on.

Returns:

(indices, likelihoods) — LongTensor[V_s, top_k] and FloatTensor[V_s, top_k] on device.

Raises:

FileNotFoundError – path does not exist.
ValueError – the file is not in the dense top-k format.

nemo_rl.algorithms.x_token.loss_utils._EXACT_TOKEN_MAP_CACHE: dict[Tuple[str, torch.device, bool, int], Dict[str, torch.Tensor]]#: None

nemo_rl.algorithms.x_token.loss_utils.build_exact_token_map( path: Union[str, os.PathLike], device: torch.device, *, xtoken_loss: bool, teacher_vocab_size: int, ) → Dict[str, torch.Tensor]#

Build the common/uncommon vocab partition for the gold path (cached).

Reads the dense projection arrays via :func:get_topk_projection, sorts each student row’s projection weights descending, then picks an exact-token map per the xtoken_loss flag:

xtoken_loss=False (strict): has_exact_map = (sorted_values[:, 0] == 1.0) & (projection_indices[:, 1] == -1). On collision (multiple students mapping to the same teacher id), the earliest (lowest) student index wins.
xtoken_loss=True (relaxed): has_exact_map = sorted_values[:, 0] >= 0.6. On collision, the student with the highest first-projection weight wins; ties are broken by lowest student index.

Both branches are vectorized via scatter_reduce so the build is O(V_s) and happens once per (path, device, xtoken_loss, teacher_vocab_size) for the run.

Parameters:

path – Path to a torch.saved projection-matrix file (dense top-k format).
device – Device the returned tensors must live on.
xtoken_loss – Selects strict vs relaxed exact-map rule (see above).
teacher_vocab_size – Width of the teacher-side vocab axis. The partition is bounded by this — teacher ids outside the range are dropped.

Returns:

Dict with keys common_student, common_teacher (paired), uncommon_student, uncommon_teacher (each independently sorted). All [long] tensors on device.

nemo_rl.algorithms.x_token.loss_utils.prepare_xtoken_cross_tokenizer_loss_input( logits: torch.Tensor, data: Mapping[str, Any], *, projection_matrix_paths: list[Optional[str]], vocab_parallel_group: Optional[torch.distributed.ProcessGroup] = None, context_parallel_group: Optional[torch.distributed.ProcessGroup] = None, ) → tuple[torch.Tensor, Dict[int, torch.Tensor], Dict[int, nemo_rl.algorithms.x_token.loss_utils.LocalizedAlignment], Optional[torch.distributed.ProcessGroup], Optional[torch.distributed.ProcessGroup]]#

Build the per-teacher cross-tokenizer distillation loss pieces from student logits + IPC teacher data.

Rebuilds each teacher’s full-vocab logits from its per-rank CUDA IPC handles (teacher_{i}_full_logits_ipc) and does the shared CP-resolution the loss needs. The contiguous student logits / input_ids / token_mask are relaid once and shared across teachers. Per teacher, a localized alignment is built: a cross-tokenizer teacher (projection_matrix_paths[i] set) gets the localized, next-token-shifted chunk alignment from its alignment_{i}_* keys; a same-tokenizer teacher (None path) gets a thin alignment carrying only the shared student fields (identity 1:1 token alignment, no chunks). TP/CP groups come from the student logits’ device mesh, falling back to the passed groups for non-DTensor logits.

Parameters:: projection_matrix_paths – Per-teacher projection paths. Its length is the teacher count and drives the teacher_{i}_* / alignment_{i}_* keys read here; a None entry marks a same-tokenizer teacher.
Returns:: (student_logits_contig, teacher_full_logits_by_idx, aligns_by_idx, tp_group, cp_group).

nemo_rl.algorithms.x_token.loss_utils#

Module Contents#

Classes#

Functions#

Data#

API#

`nemo_rl.algorithms.x_token.loss_utils`#