`nemo_rl.algorithms.loss.loss_functions`#

Module Contents#

Classes#

`DraftCrossEntropyLossConfig`
`DraftCrossEntropyLossDataDict`
`DraftCrossEntropyLossFn`	Compute the auxiliary soft-target cross-entropy used for draft-model training.
`ClippedPGLossConfig`
`ClippedPGLossDataDict`	Required keys for the Clipped Policy Gradient loss function.
`ClippedPGLossFn`	Generalized Clipped Policy Gradient loss function w/ KL regularization.
`NLLLossFn`	Negative Log Likelihood Loss function.
`PreferenceLossDataDict`	Required keys for the preference loss function.
`PreferenceLossFn`	Preference Loss function.
`DPOLossConfig`
`DPOLossDataDict`	Required keys for the DPO loss function.
`DPOLossFn`	Direct Preference Optimization (DPO) loss function.
`DistillationLossConfig`
`DistillationLossDataDict`
`DistillationLossFn`	Distillation loss function.
`MseValueLossConfig`	Config for the MSE value loss used by PPO’s value model.
`MseValueLossFn`	Mean Squared Error value loss function with optional clipping (PPO-style).
`CrossTokenizerDistillationLossConfig`	Config for cross-tokenizer distillation loss.
`CrossTokenizerDistillationLossDataDict`	Student-side keys are fixed; teacher-side keys are teacher-indexed.
`CrossTokenizerDistillationLossFn`	Cross-tokenizer distillation loss.

Data#

Tensor

API#

nemo_rl.algorithms.loss.loss_functions.Tensor#: ‘TypeVar(…)’

class nemo_rl.algorithms.loss.loss_functions.DraftCrossEntropyLossConfig#

Bases: typing.TypedDict

vocab_parallel_group: Optional[torch.distributed.ProcessGroup]#: None

class nemo_rl.algorithms.loss.loss_functions.DraftCrossEntropyLossDataDict#

Bases: typing.TypedDict

teacher_logits: nemo_rl.algorithms.loss.loss_functions.Tensor#: None

student_logits: nemo_rl.algorithms.loss.loss_functions.Tensor#: None

token_mask: nemo_rl.algorithms.loss.loss_functions.Tensor#: None

sample_mask: nemo_rl.algorithms.loss.loss_functions.Tensor#: None

student_vocab_indices: NotRequired[nemo_rl.algorithms.loss.loss_functions.Tensor]#: None

class nemo_rl.algorithms.loss.loss_functions.DraftCrossEntropyLossFn( vocab_parallel_group: Optional[torch.distributed.ProcessGroup] = None, )#

Bases: nemo_rl.algorithms.loss.interfaces.LossFunction

Compute the auxiliary soft-target cross-entropy used for draft-model training.

Initialization

loss_type#: None

input_type#: None

__call__( teacher_logits: nemo_rl.algorithms.loss.loss_functions.Tensor, student_logits: nemo_rl.algorithms.loss.loss_functions.Tensor, token_mask: nemo_rl.algorithms.loss.loss_functions.Tensor, data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.algorithms.loss.loss_functions.DraftCrossEntropyLossDataDict], global_valid_seqs: torch.Tensor, global_valid_toks: torch.Tensor, ) → torch.Tensor#: Reduce the masked per-token draft loss to a scalar.

class nemo_rl.algorithms.loss.loss_functions.ClippedPGLossConfig#

Bases: pydantic.BaseModel

disable_ppo_ratio: bool#: False

token_level_loss: bool#: True

sequence_level_importance_ratios: bool#: False

ratio_clip_min: float#: 0.2

ratio_clip_max: float#: 0.2

ratio_clip_c: Optional[float]#: None

reference_policy_kl_penalty: float#: 0.01

reference_policy_kl_type: str#: ‘k3’

kl_input_clamp_value: Optional[float]#: 20.0

kl_output_clamp_value: Optional[float]#: 10.0

use_kl_in_reward: bool#: False

use_importance_sampling_correction: bool#: False

truncated_importance_sampling_type: Optional[str]#: None

truncated_importance_sampling_ratio: Optional[float]#: None

truncated_importance_sampling_ratio_min: Optional[float]#: None

use_on_policy_kl_approximation: bool#: False

force_on_policy_ratio: bool#: False

use_cispo: bool#: False

positive_example_nll_weight: float#: 0.0

class nemo_rl.algorithms.loss.loss_functions.ClippedPGLossDataDict#

Bases: typing.TypedDict

Required keys for the Clipped Policy Gradient loss function.

Initialization

Initialize self. See help(type(self)) for accurate signature.

input_ids: torch.Tensor#: None

advantages: torch.Tensor#: None

prev_logprobs: torch.Tensor#: None

generation_logprobs: torch.Tensor#: None

reference_policy_logprobs: torch.Tensor#: None

token_mask: torch.Tensor#: None

sample_mask: torch.Tensor#: None

__extra__: Any#: None

class nemo_rl.algorithms.loss.loss_functions.ClippedPGLossFn( cfg: nemo_rl.algorithms.loss.loss_functions.ClippedPGLossConfig, use_fused_linear_logprobs: bool = False, )#

Bases: nemo_rl.algorithms.loss.interfaces.LossFunction

Generalized Clipped Policy Gradient loss function w/ KL regularization.

This implements:

PPO (Clipped) - https://arxiv.org/abs/1707.06347
GRPO - https://arxiv.org/abs/2402.03300
REINFORCE/RLOO (set disable_ppo_ratio = True and ignores ratio_clip_min/ratio_clip_max) - https://arxiv.org/abs/2402.14740
GSPO (set sequence_level_importance_ratios = True and token_level_loss = False) - https://arxiv.org/abs/2507.18071
CISPO (set use_cispo = True) - https://arxiv.org/abs/2506.13585
Truly on-policy (set force_on_policy_ratio = True to force ratio = 1.0, requires one update per rollout)

Formula: L(θ) = E_t [ min(r_t(θ) * A_t, clip(r_t(θ), 1-ε, 1+ε) * A_t) ] - β * KL(π_θ || π_ref)

where:

r_t(θ) = π_θ(a_t|s_t) / π_θ_old(a_t|s_t) is the probability ratio
A_t is the advantage estimate
ε is the clip parameter (ratio_clip_min/ratio_clip_max)
- As proposed in the DAPO paper (https://arxiv.org/pdf/2503.14476), we allow setting a distinct minimum and maximum value for the clip parameter (set to the same value for PPO/GRPO/etc.)
  - ratio_clip_min: minimum value for the clip parameter
  - ratio_clip_max: maximum value for the clip parameter
β is the KL penalty coefficient (reference_policy_kl_penalty)
KL(π_θ || π_ref) is the KL divergence between the current policy and reference policy (Schulman Approx.)

For REINFORCE/RLOO (when disable_ppo_ratio=True), the formula simplifies to: L(θ) = E_t [ π_θ(a_t|s_t) * A_t ] - β * KL(π_θ || π_ref)

Formula (CISPO): L(θ) = E_t [ sg(clip(r_t(θ), 1-ε_low, 1+ε_high)) * A_t * log π_θ(a_t|s_t) ]

Also supports “Dual-Clipping” from https://arxiv.org/pdf/1912.09729, which imposes an additional upper bound on the probability ratio when advantages are negative. This prevents excessive policy updates. \(rA << 0\) -> \(cA\)(clipped) The loss function is modified to the following when A_t < 0: L(θ) = E_t [ max(min(r_t(θ) * A_t, clip(r_t(θ), 1-ε, 1+ε) * A_t), c * A_t) ] - β * KL(π_θ || π_ref)

where:

c is the dual-clip parameter (ratio_clip_c), which must be greater than 1 and is usually set as 3 empirically.

Due to potential numerical instability, we cast the logits to float32 before computing the loss.

Initialization

input_type#: None

__call__( next_token_logprobs: nemo_rl.algorithms.loss.loss_functions.Tensor, data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.algorithms.loss.loss_functions.ClippedPGLossDataDict], global_valid_seqs: torch.Tensor, global_valid_toks: torch.Tensor, ) → tuple[torch.Tensor, dict]#: Clipped Policy Gradient RL loss function.

class nemo_rl.algorithms.loss.loss_functions.NLLLossFn(use_fused_linear_logprobs: bool = False)#

Bases: nemo_rl.algorithms.loss.interfaces.LossFunction

Negative Log Likelihood Loss function.

Initialization

loss_type#: None

input_type#: None

__call__( next_token_logprobs: nemo_rl.algorithms.loss.loss_functions.Tensor, data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any], global_valid_seqs: nemo_rl.algorithms.loss.loss_functions.Tensor | None, global_valid_toks: nemo_rl.algorithms.loss.loss_functions.Tensor, dpo_loss: bool = False, dpo_average_log_probs: bool = False, ) → tuple[torch.Tensor, dict[str, Any]]#

class nemo_rl.algorithms.loss.loss_functions.PreferenceLossDataDict#

Bases: typing.TypedDict

Required keys for the preference loss function.

Initialization

Initialize self. See help(type(self)) for accurate signature.

input_ids: torch.Tensor#: None

token_mask: torch.Tensor#: None

sample_mask: torch.Tensor#: None

class nemo_rl.algorithms.loss.loss_functions.PreferenceLossFn#

Bases: nemo_rl.algorithms.loss.interfaces.LossFunction

Preference Loss function.

Optimizes the model to prefer chosen responses over rejected ones

The preference loss is computed as: L_pref(θ) = -E[log(σ(β * (r_chosen - r_rejected)))]

where:

σ is the sigmoid function
β is a scaling factor (ex: reference_policy_kl_penalty in DPO)
r_chosen and r_rejected are the rewards for chosen and rejected responses

Returns:: A tuple containing: - The preference loss value - A dictionary with metrics including: - loss: Preference loss - accuracy: Fraction of examples where chosen response has higher reward
Return type:: tuple[torch.Tensor, dict]

loss_type#: None

input_type#: None

split_output_tensor( tensor: nemo_rl.algorithms.loss.loss_functions.Tensor, ) → tuple[nemo_rl.algorithms.loss.loss_functions.Tensor, nemo_rl.algorithms.loss.loss_functions.Tensor]#

_preference_loss( rewards: nemo_rl.algorithms.loss.loss_functions.Tensor, sample_mask: nemo_rl.algorithms.loss.loss_functions.Tensor, global_valid_seqs: nemo_rl.algorithms.loss.loss_functions.Tensor, beta: float = 1.0, ) → tuple[nemo_rl.algorithms.loss.loss_functions.Tensor, nemo_rl.algorithms.loss.loss_functions.Tensor, nemo_rl.algorithms.loss.loss_functions.Tensor, nemo_rl.algorithms.loss.loss_functions.Tensor]#

__call__( logits: nemo_rl.algorithms.loss.loss_functions.Tensor, data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.algorithms.loss.loss_functions.PreferenceLossDataDict], global_valid_seqs: nemo_rl.algorithms.loss.loss_functions.Tensor, global_valid_toks: nemo_rl.algorithms.loss.loss_functions.Tensor | None, ) → tuple[torch.Tensor, dict[str, Any]]#

class nemo_rl.algorithms.loss.loss_functions.DPOLossConfig#

Bases: pydantic.BaseModel

reference_policy_kl_penalty: float#: 0.05

preference_loss_weight: float#: 1.0

sft_loss_weight: float#: 0.0

preference_average_log_probs: bool#: False

sft_average_log_probs: bool#: False

class nemo_rl.algorithms.loss.loss_functions.DPOLossDataDict#

Bases: typing.TypedDict

Required keys for the DPO loss function.

Initialization

Initialize self. See help(type(self)) for accurate signature.

input_ids: torch.Tensor#: None

reference_policy_logprobs: torch.Tensor#: None

token_mask: torch.Tensor#: None

sample_mask: torch.Tensor#: None

class nemo_rl.algorithms.loss.loss_functions.DPOLossFn( cfg: nemo_rl.algorithms.loss.loss_functions.DPOLossConfig, use_fused_linear_logprobs: bool = False, )#

Bases: nemo_rl.algorithms.loss.loss_functions.PreferenceLossFn

Direct Preference Optimization (DPO) loss function.

This loss function implements the DPO algorithm as described in: “Direct Preference Optimization: Your Language Model is Secretly a Reward Model” (https://arxiv.org/abs/2305.18290)

The loss combines two main components:

Preference Loss: Optimizes the model to prefer chosen responses over rejected ones
SFT Loss (optional): Auxiliary supervised fine-tuning loss on chosen responses

The total loss is computed as: L(θ) = w_p * L_pref(θ) + w_s * L_sft(θ)

where:

w_p is the preference_loss_weight
w_s is the sft_loss_weight
L_pref(θ) is the preference loss term
L_sft(θ) is the supervised fine-tuning loss term

The preference loss term is computed as: L_pref(θ) = -E[log(σ(β * (r_chosen - r_rejected)))]

where:

σ is the sigmoid function
β is the reference_policy_kl_penalty
r_chosen and r_rejected are the rewards for chosen and rejected responses
The rewards are computed as the sum of log probability differences between the current policy and reference policy

If preference_average_log_probs is True, the rewards are averaged over tokens: r = (1/n) * Σ_t (log π_θ(a_t|s_t) - log π_ref(a_t|s_t))

Otherwise, the rewards are summed over tokens.

The SFT loss term is a standard negative log likelihood loss on the chosen responses. If sft_average_log_probs is True, the loss is averaged over tokens.

Parameters:

cfg (DPOLossConfig) –

Configuration dictionary containing:

reference_policy_kl_penalty (float): Strength of the KL penalty term (β)
preference_loss_weight (float): Weight for the preference loss term (w_p)
sft_loss_weight (float): Weight for the SFT loss term (w_s)
preference_average_log_probs (bool): Whether to average log probs across tokens in preference loss
sft_average_log_probs (bool): Whether to average log probs across tokens in SFT loss

Returns:

A tuple containing: - The total loss value - A dictionary with metrics including: - loss: Total loss value - sft_loss: SFT loss component - preference_loss: Preference loss component - accuracy: Fraction of examples where chosen response has higher reward

Return type:

tuple[torch.Tensor, dict]

Initialization

loss_type#: None

input_type#: None

_dpo_loss( next_token_logprobs: nemo_rl.algorithms.loss.loss_functions.Tensor, data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.algorithms.loss.loss_functions.DPOLossDataDict], global_valid_seqs: nemo_rl.algorithms.loss.loss_functions.Tensor, ) → tuple[nemo_rl.algorithms.loss.loss_functions.Tensor, nemo_rl.algorithms.loss.loss_functions.Tensor, nemo_rl.algorithms.loss.loss_functions.Tensor, nemo_rl.algorithms.loss.loss_functions.Tensor]#

__call__( next_token_logprobs: nemo_rl.algorithms.loss.loss_functions.Tensor, data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.algorithms.loss.loss_functions.DPOLossDataDict], global_valid_seqs: nemo_rl.algorithms.loss.loss_functions.Tensor, global_valid_toks: nemo_rl.algorithms.loss.loss_functions.Tensor | None, ) → tuple[torch.Tensor, dict[str, Any]]#

class nemo_rl.algorithms.loss.loss_functions.DistillationLossConfig#

Bases: typing.TypedDict

kl_type: str#: None

mixed_kl_weight: float#: None

zero_outside_topk: bool#: None

class nemo_rl.algorithms.loss.loss_functions.DistillationLossDataDict#

Bases: typing.TypedDict

input_ids: torch.Tensor#: None

input_lengths: torch.Tensor#: None

token_mask: torch.Tensor#: None

sample_mask: torch.Tensor#: None

teacher_topk_logits: torch.Tensor#: None

teacher_topk_indices: torch.Tensor#: None

class nemo_rl.algorithms.loss.loss_functions.DistillationLossFn( cfg: nemo_rl.algorithms.loss.loss_functions.DistillationLossConfig, )#

Bases: nemo_rl.algorithms.loss.interfaces.LossFunction

Distillation loss function.

Initialization

loss_type#: None

input_type#: None

__call__( student_topk_logprobs: torch.Tensor, teacher_topk_logprobs: torch.Tensor, H_all: torch.Tensor | None, data: nemo_rl.algorithms.loss.loss_functions.DistillationLossDataDict, global_valid_seqs: torch.Tensor, global_valid_toks: torch.Tensor, ) → tuple[torch.Tensor, dict[str, Any]]#: Compute distillation loss between teacher and student logits.

class nemo_rl.algorithms.loss.loss_functions.MseValueLossConfig#

Bases: pydantic.BaseModel

Config for the MSE value loss used by PPO’s value model.

scale: float#: 1.0

cliprange: Optional[float]#: None

class nemo_rl.algorithms.loss.loss_functions.MseValueLossFn( cfg: nemo_rl.algorithms.loss.loss_functions.MseValueLossConfig, )#

Bases: nemo_rl.algorithms.loss.interfaces.LossFunction

Mean Squared Error value loss function with optional clipping (PPO-style).

When cliprange is set, value predictions are clipped to [old_values - cliprange, old_values + cliprange] and the loss is 0.5 * max(mse(vpred, returns), mse(vpred_clipped, returns)). This prevents the value function from changing too drastically in a single update, mirroring the policy ratio clipping in PPO.

Initialization

input_type#: None

__call__( logits: torch.Tensor, data: nemo_rl.distributed.batched_data_dict.BatchedDataDict, global_valid_seqs: torch.Tensor, global_valid_toks: torch.Tensor, ) → tuple[torch.Tensor, dict[str, Any]]#: Compute Mean Squared Error value loss, optionally with clipping.

class nemo_rl.algorithms.loss.loss_functions.CrossTokenizerDistillationLossConfig#

Bases: typing.TypedDict

Config for cross-tokenizer distillation loss.

.. attribute:: projection_matrix_paths

Per-teacher list of filesystem paths to the .pt projection file (None marks a same-tokenizer teacher: direct KL, no projection). Each .pt holds either the dense top-k projection (dict with ‘indices’ and ‘likelihoods’ tensors of shape [V_student, top_k]) or the sparse multi-token format (dict[(student_id, teacher_id)] -> count), loaded lazily on first call by each worker process. Runtime-injected by xtoken_off_policy_distillation.setup from teachers[i]; not a user loss_fn key in YAML.

.. attribute:: gold_loss

If True, switch to the gold-loss formulation: split the vocab into an exact-token-mapped common set (KL) and an uncommon set (sorted L1).

.. attribute:: xtoken_loss

Modifier inside the gold-loss path. If True, relaxes the exact-map threshold to >= 0.6 (vs == 1.0) and adds a collision-replacement rule so multi-token projections can still contribute exact maps. Requires gold_loss=True.

.. attribute:: temperature

Softmax temperature applied symmetrically to student and teacher logits before KL.

.. attribute:: vocab_topk

Microbatch-global top-k size used by the P-KL path (gold_loss=False). Computed inside the loss fn from full teacher logits. Inert when gold_loss=True.

.. attribute:: uncommon_topk

Cap on the L1 uncommon-tail sort in the gold path. Defaults to 8192. Inert when gold_loss=False.

.. attribute:: reverse_kl

If True, compute KL(student || teacher) instead of KL(teacher || student).

.. attribute:: exact_token_match_only

If True, only aligned pairs flagged as ‘is_correct’ contribute to KL; mismatched pairs are masked out. Used by the P-KL path only.

.. attribute:: kl_loss_weight

Scalar multiplier on the distillation (KD) term in fixed-weight mode (dynamic_loss_scaling=False). Applies to both the P-KL and gold-loss paths.

.. attribute:: ce_loss_scale

Scalar multiplier on the next-token CE term in fixed-weight mode. Applies to both the P-KL and gold-loss paths.

.. attribute:: dynamic_loss_scaling

If True, rescale the KD term each step so its detached magnitude matches CE, then add CE; kl_loss_weight / ce_loss_scale are ignored in this mode. Applies to both the P-KL and gold-loss paths.

.. attribute:: student_vocab_size

Full student tokenizer vocab size, used to size the projection matrix’s student-side (V_s) axis. Runtime-injected by xtoken_off_policy_distillation.setup from len(student_tokenizer); not a user knob in YAML. Sizing V_s from the configured tokenizer vocab (rather than max(observed student_id) + 1 from the sparse projection file) keeps V_s in lockstep with logits.shape[-1] when the file’s highest student ids happen to be absent.

.. attribute:: teacher_vocab_sizes

Per-teacher list of full teacher tokenizer vocab sizes, used to size each projection matrix’s teacher-side (V_t) axis. Runtime-injected symmetrically to student_vocab_size from each len(teacher_tokenizer); not a user loss_fn key in YAML.

Initialization

Initialize self. See help(type(self)) for accurate signature.

gold_loss: bool#: None

xtoken_loss: bool#: None

temperature: float#: None

vocab_topk: int#: None

uncommon_topk: int#: None

reverse_kl: bool#: None

exact_token_match_only: bool#: None

kl_loss_weight: float#: None

ce_loss_scale: float#: None

dynamic_loss_scaling: bool#: None

kd_loss_mode: str#: None

normalize_teacher_by_vocab: bool#: None

alpha: float#: None

sum_weights_metric: NotRequired[Optional[str]]#: None

student_vocab_size: NotRequired[int]#: None

teacher_vocab_sizes: NotRequired[list[int]]#: None

projection_matrix_paths: NotRequired[list[Optional[str]]]#: None

teacher_weights: NotRequired[list[float]]#: None

teacher_gold_loss: NotRequired[list[Optional[bool]]]#: None

teacher_xtoken_loss: NotRequired[list[Optional[bool]]]#: None

class nemo_rl.algorithms.loss.loss_functions.CrossTokenizerDistillationLossDataDict#

Bases: typing.TypedDict

Student-side keys are fixed; teacher-side keys are teacher-indexed.

Only the student keys below are static. Each teacher i contributes a dynamic set of keys produced by CrossTokenizerCollator / the trainer and so cannot be enumerated here:

Every teacher: teacher_{i}_full_logits_ipc — List[B] of CUDA IPC handle dicts (payload_ipc + buf_idx/sample_index_in_buf + TP/CP shard metadata) from Policy.get_full_logits_ipc. rebuild_teacher_full_logits_from_ipc (in prepare_loss_input) P2P-reads and reassembles full-vocab teacher logits, routing across heterogeneous teacher/student TP/CP.
Cross-tokenizer teacher only: teacher_{i}_input_ids / teacher_{i}_token_mask [B, T_t] and alignment_{i}_* (pair_valid / pair_is_correct [B, max_pairs]; student_chunk_id [B, T_s]; teacher_chunk_id [B, T_t]; partition masks; num_chunks).
Same-tokenizer teacher: no teacher_{i}_input_ids / alignment_{i}_*; it reuses the student tokenization (identity 1:1 aligned).

Initialization

Initialize self. See help(type(self)) for accurate signature.

input_ids: torch.Tensor#: None

input_lengths: torch.Tensor#: None

token_mask: torch.Tensor#: None

sample_mask: torch.Tensor#: None

class nemo_rl.algorithms.loss.loss_functions.CrossTokenizerDistillationLossFn( cfg: nemo_rl.algorithms.loss.loss_functions.CrossTokenizerDistillationLossConfig, )#

Bases: nemo_rl.algorithms.loss.interfaces.LossFunction

Cross-tokenizer distillation loss.

Mode is selected by the (gold_loss, xtoken_loss) flags:

(False, False) -> P-KL: full-vocab projection KL (student logits mapped through the projection matrix M) plus a standard next-token student CE term, combined as kl_loss_weight * kl + ce_loss_scale * ce — or, when dynamic_loss_scaling is set, with the KL term rescaled each step to match the detached CE magnitude.
(True, False) -> gold-loss: KL on the exact-mapped common partition plus a sorted-L1 term on the uncommon tail (kd = (kl_common + l1_uncommon) * T**2), combined with a next-token student CE term the same way as the P-KL path — kl_loss_weight * kd + ce_loss_scale * ce, or, when dynamic_loss_scaling is set, with the KD term rescaled each step to match the detached CE magnitude.
(True, True) -> gold-loss with the xtoken modifier: same objective, but the exact-map threshold is relaxed (>= 0.6 instead of == 1.0) and a collision-replacement rule lets multi-token projections still contribute exact maps.

(False, True) is rejected in __init__: xtoken_loss is a modifier inside the gold path and is undefined for P-KL.

Multi-teacher: setup injects per-teacher metadata (projection paths, weights, vocab sizes, per-teacher gold/xtoken overrides). The per-teacher KD terms are aggregated by kd_loss_mode (sum / averaged_logits / select_teacher) and combined with a single student CE term. A teacher with a None projection path is a same-tokenizer teacher: projection and alignment are skipped and its KD term is a direct top-k per-position KL on the shared vocab (top-k selected student-side over the reassembled full-vocab teacher logits). The single-teacher path is just num_teachers == 1.

Inputs (via LossInputType.DISTILLATION_CROSS_TOKENIZER): logits: [B, T_s, V_s] raw student logits from the worker forward. student_logits_contig: CP-relaid contiguous student logits shared by every teacher’s KD term. teacher_full_logits_by_idx: dict[int, [B, T, V_t]] full-vocab teacher logits per teacher, rebuilt from the CUDA IPC handles by prepare_loss_input (see :func:nemo_rl.algorithms.x_token.loss_utils.rebuild_teacher_full_logits_from_ipc). aligns_by_idx: dict[int, LocalizedAlignment] per teacher (cross-tok: localized chunk alignment; same-tok: thin, student fields only).

Inputs (via data: BatchedDataDict): See :class:CrossTokenizerDistillationLossDataDict.

Returns:: loss, kl_loss (the aggregated KD term), ce_loss, kl_loss_scale, accuracy, num_valid_samples. Per-teacher metrics are suffixed _t{i} (e.g. kl_loss_t0, proj_accuracy_t0, weight_t0); select_teacher additionally reports selected_teacher.
Return type:: (loss, metrics). Aggregate metrics

Initialization

loss_type#: None

input_type#: None

_teacher_is_same_vocab(i: int) → bool#: A teacher is same-vocab (direct KL, no projection) iff its path is None.

__call__( data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.algorithms.loss.loss_functions.CrossTokenizerDistillationLossDataDict], global_valid_seqs: torch.Tensor, global_valid_toks: torch.Tensor, logits: torch.Tensor, student_logits_contig: torch.Tensor, teacher_full_logits_by_idx: dict[int, torch.Tensor], aligns_by_idx: dict[int, nemo_rl.algorithms.x_token.loss_utils.LocalizedAlignment], *, tp_group: Optional[torch.distributed.ProcessGroup] = None, cp_group: Optional[torch.distributed.ProcessGroup] = None, ) → tuple[torch.Tensor, dict[str, Any]]#

Compute the (multi-teacher) cross-tokenizer distillation loss.

Per-teacher KD terms are aggregated per kd_loss_mode and combined with a single student next-token CE term — dynamic-scaled when dynamic_loss_scaling is set (KD rescaled to match the detached CE magnitude, kl_loss_weight / ce_loss_scale ignored), else fixed-weighted. The single-teacher path is just num_teachers == 1.

student_logits_contig (CP-relaid) and the per-teacher aligns_by_idx / teacher_full_logits_by_idx are precomputed in prepare_loss_input; the raw logits is kept for the CE term.

_resolve_gold_xtoken( i: int, use_per_teacher: bool, ) → tuple[bool, bool]#

Effective (gold_loss, xtoken_loss) for teacher i.

Per-teacher overrides are honored only when use_per_teacher is set (sum mode); select_teacher / averaged_logits use the global flags. A None override falls back to the global value.

_compute_teacher_kd( i: int, student_logits_contig: torch.Tensor, data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.algorithms.loss.loss_functions.CrossTokenizerDistillationLossDataDict], teacher_full_logits_by_idx: dict[int, torch.Tensor], aligns_by_idx: dict[int, nemo_rl.algorithms.x_token.loss_utils.LocalizedAlignment], global_valid_toks: torch.Tensor, *, use_per_teacher_flags: bool, tp_group: Optional[torch.distributed.ProcessGroup], cp_group: Optional[torch.distributed.ProcessGroup], ) → tuple[torch.Tensor, dict[str, Any]]#

KD term for teacher i plus its (unsuffixed) metrics.

Dispatches on tokenizer kind: same-vocab -> direct top-k per-position KL; cross-tokenizer -> P-KL or gold path using teacher i’s projection and its localized alignment. Both consume the shared CP-relaid student logits and route TP/CP through main’s parameterized loss-mode helpers.

_compute_same_vocab_kl( i: int, student_logits_contig: torch.Tensor, teacher_full_logits: torch.Tensor, align: nemo_rl.algorithms.x_token.loss_utils.LocalizedAlignment, global_valid_toks: torch.Tensor, *, tp_group: Optional[torch.distributed.ProcessGroup], cp_group: Optional[torch.distributed.ProcessGroup], ) → tuple[torch.Tensor, dict[str, Any]]#

Direct top-k per-position KL for a same-tokenizer teacher.

Identical tokenizer => teacher tokens == student tokens (identity position alignment), so no projection / no chunk-averaging. The reduction matches CE: masked next-token mean normalized by global_valid_toks, scaled by T**2.

_direct_topk_kl( student_logits: torch.Tensor, teacher_full_logits: torch.Tensor, align: nemo_rl.algorithms.x_token.loss_utils.LocalizedAlignment, global_valid_toks: torch.Tensor, *, tp_group: Optional[torch.distributed.ProcessGroup], cp_group: Optional[torch.distributed.ProcessGroup], ) → torch.Tensor#

Top-K per-position KL on a shared vocab (same tokenizer), TP/CP-aware.

Top-K columns are selected at the student from the reassembled full-vocab teacher logits (select_teacher_topk_indices MAX-reduces across CP so every CP rank agrees on the same columns). The student is gathered to full vocab across TP (vocab_parallel_full_log_softmax) before slicing — slicing a TP-local shard would pick the wrong columns. Both sides are renormalized within the K-subset (the teacher’s subset-softmax and the student’s full-then-renorm are mathematically identical, the full-vocab partition function cancels). The masked next-token mean is normalized by the CP/DP-global valid-token count, exactly like the CE term; at CP=1 the cp_shift_next mask reduces to the token_mask[:, 1:] next-token shift.

_direct_full_vocab_kl( student_logits: torch.Tensor, teacher_logits: torch.Tensor, align: nemo_rl.algorithms.x_token.loss_utils.LocalizedAlignment, global_valid_toks: torch.Tensor, *, tp_group: Optional[torch.distributed.ProcessGroup], cp_group: Optional[torch.distributed.ProcessGroup], ) → torch.Tensor#

Full-vocab per-position KL on a shared vocab (same tokenizer), TP/CP-aware.

Used by averaged_logits over the convex-averaged teacher logits. The student is gathered to full vocab across TP; the teacher (already full vocab) is sliced to the student width to drop HF lm_head padding.

_same_vocab_masked_kl( per_pos: torch.Tensor, align: nemo_rl.algorithms.x_token.loss_utils.LocalizedAlignment, global_valid_toks: torch.Tensor, cp_group: Optional[torch.distributed.ProcessGroup], ) → torch.Tensor#

Masked next-token mean of a per-position KL (same-tokenizer reduction).

per_pos is the per-position KL on this CP rank’s contiguous window (student and teacher are position-aligned). The CP-aware next-token shift (:func:cp_shift_next) selects positions whose target token (p+1) is a valid label — at CP=1 this is the plain token_mask[:, 1:] shift, the global-last position dropped via fill=0. Reduction matches CE: masked_mean over global_valid_toks, scaled by T**2.

_sum_kd( student_logits_contig: torch.Tensor, data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.algorithms.loss.loss_functions.CrossTokenizerDistillationLossDataDict], teacher_full_logits_by_idx: dict[int, torch.Tensor], aligns_by_idx: dict[int, nemo_rl.algorithms.x_token.loss_utils.LocalizedAlignment], global_valid_toks: torch.Tensor, *, tp_group: Optional[torch.distributed.ProcessGroup], cp_group: Optional[torch.distributed.ProcessGroup], ) → tuple[torch.Tensor, dict[str, Any]]#

Weighted sum: total_kd = Σ_i weight_i · KD_i.

Weights are static (config weight) or dynamic (sum_weights_metric). When normalize_teacher_by_vocab is set, each teacher’s KD is additionally scaled by log(V_t_i) / log(min_j V_t_j).

_averaged_logits_kd( student_logits_contig: torch.Tensor, data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.algorithms.loss.loss_functions.CrossTokenizerDistillationLossDataDict], teacher_full_logits_by_idx: dict[int, torch.Tensor], aligns_by_idx: dict[int, nemo_rl.algorithms.x_token.loss_utils.LocalizedAlignment], global_valid_toks: torch.Tensor, *, tp_group: Optional[torch.distributed.ProcessGroup], cp_group: Optional[torch.distributed.ProcessGroup], ) → tuple[torch.Tensor, dict[str, Any]]#

Convex-weighted average of teacher logits, then one direct KL.

Valid only when all teachers are same-tokenizer (no projection) and ship full logits of identical shape. Otherwise falls back to a plain static-weight sum (no dynamic weights, no normalize_teacher_by_vocab).

_dp_global_masked_mean( values: torch.Tensor, mask: torch.Tensor, ) → torch.Tensor#

Masked mean of values over the process-global valid count.

Teacher selection / dynamic weighting must be identical on every rank: a rank-local mean lets ranks pick a different teacher / different weights, and the per-teacher KD’s collectives then see divergent participation (deadlock when one rank’s choice fires a collective another’s does not). All-reduce the masked sum and the mask count over the full group so every rank gets the same score (mirrors _compute_p_kl’s WORLD-reduced denominator). The result is detached (it gates selection / weighting and is not back-propagated).

_select_teacher_kd( student_logits_contig: torch.Tensor, data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.algorithms.loss.loss_functions.CrossTokenizerDistillationLossDataDict], teacher_full_logits_by_idx: dict[int, torch.Tensor], aligns_by_idx: dict[int, nemo_rl.algorithms.x_token.loss_utils.LocalizedAlignment], global_valid_toks: torch.Tensor, *, tp_group: Optional[torch.distributed.ProcessGroup], cp_group: Optional[torch.distributed.ProcessGroup], ) → tuple[torch.Tensor, dict[str, Any]]#: Use only the teacher with the lowest next-token CE on its own tokens.

_teacher_score_inputs( i: int, data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.algorithms.loss.loss_functions.CrossTokenizerDistillationLossDataDict], teacher_full_logits_by_idx: dict[int, torch.Tensor], aligns_by_idx: dict[int, nemo_rl.algorithms.x_token.loss_utils.LocalizedAlignment], ) → tuple[torch.Tensor, torch.Tensor, torch.Tensor]#

Return (logits, input_ids, token_mask) for teacher i’s CE / weight-metric scores.

The token mask is over the tokenization the score is computed on: the shared student tokens (CP-relaid) for a same-vocab teacher, teacher i’s own otherwise. Every teacher ships full logits, so the full distribution is always available.

_compute_dynamic_weights( data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.algorithms.loss.loss_functions.CrossTokenizerDistillationLossDataDict], teacher_full_logits_by_idx: dict[int, torch.Tensor], aligns_by_idx: dict[int, nemo_rl.algorithms.x_token.loss_utils.LocalizedAlignment], ) → list[torch.Tensor]#

Sequence-level dynamic teacher weights via sum_weights_metric.

Per teacher computes a scalar score (ce -> -CE, entropy -> -entropy, max_prob -> max prob; higher = more trusted), optionally rescaled by log(V_t_i)/log(min_j V_t_j), then softmax(alpha * scores) across teachers.

_teacher_weight_score( t_logits: torch.Tensor, t_ids: torch.Tensor, t_mask: torch.Tensor, sample_mask: torch.Tensor, ) → torch.Tensor#

Scalar weight-metric score for one teacher (higher = more trusted).

Padded positions and masked-out samples are excluded, so long-padded batches don’t let near-uniform padding logits dominate the score.

_compute_p_kl( student_logits: torch.Tensor, teacher_full_logits: torch.Tensor, align: nemo_rl.algorithms.x_token.loss_utils.LocalizedAlignment, *, projection_matrix_path: Optional[str], teacher_vocab_size: int, tp_group: Optional[torch.distributed.ProcessGroup] = None, cp_group: Optional[torch.distributed.ProcessGroup] = None, ) → tuple[torch.Tensor, torch.Tensor, torch.Tensor]#

P-KL: chunk-averaged KL over a microbatch-global top-k teacher subset.

student_logits (CP-relaid to contiguous) and align (localized, with next-token-shifted chunk ids) are precomputed in prepare_loss_input.

Steps:

Project full-vocab student probs through M to teacher vocab.
Use the full teacher logits materialized by prepare_loss_input.
Compute one global_top_indices [k] per microbatch from the teacher’s importance: max over flat (B*T_t), topk over V_t. Same vocab subset across every sample/position — keeps chunk-averaged KL well-defined.
Slice both the projected student probs and the teacher logits to those k columns.
Build per-token chunk masks from alignment_*_chunk_id and chunk-average via bmm (shared helper).
Renormalize student chunk distributions inside the top-k subset (avg-then-renormalize, log).
Forward (or reverse) KL between chunk distributions.

_compute_gold( student_logits: torch.Tensor, teacher_full_logits: torch.Tensor, align: nemo_rl.algorithms.x_token.loss_utils.LocalizedAlignment, *, projection_matrix_path: Optional[str], teacher_vocab_size: int, xtoken_loss: bool, tp_group: Optional[torch.distributed.ProcessGroup] = None, cp_group: Optional[torch.distributed.ProcessGroup] = None, ) → tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]#

Gold-loss path: KL on common (exact-mapped) vocab + L1 on uncommon.

student_logits (CP-relaid to contiguous) and align (localized, with next-token-shifted chunk ids) are precomputed in prepare_loss_input.

Lazy-build the exact-token map (cached per device).
Use the full teacher logits materialized by prepare_loss_input.
log_softmax on full vocab both sides; chunk-average via the shared helper using the precomputed next-token-shifted chunk ids.
Slice each chunk-averaged tensor to common_* indices and compute (forward or reverse) KL, reduced as sum / valid_chunk.sum() where valid_chunk is the geometric chunk mask AND’d with sample_mask (mirrors the P-KL path).
Slice to uncommon_* indices, .exp() to probs, sort/topk descending (capped at self.uncommon_topk), truncate to min(student_len, teacher_len), L1 with reduction="none" summed over vocab and meaned across valid chunks.
Combine: loss = (kl_common + l1_uncommon) * T**2.
Top-1 accuracy on the common slice over valid chunks.

Returns (loss, kl_common, l1_uncommon, num_valid_chunks, top1_acc). Components other than loss are detached.

_compute_ce( logits: torch.Tensor, data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.algorithms.loss.loss_functions.CrossTokenizerDistillationLossDataDict], global_valid_toks: torch.Tensor, ) → torch.Tensor#: Next-token CE on the student side (TP/CP handled by the helpers).

nemo_rl.algorithms.loss.loss_functions#

Module Contents#

Classes#

Data#

API#

`nemo_rl.algorithms.loss.loss_functions`#