loss_functions.py

# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import warnings
from typing import (
    TYPE_CHECKING,
    Any,
    Literal,
    NotRequired,
    Optional,
    TypedDict,
    TypeVar,
)

import torch
from pydantic import BaseModel, Field

from nemo_rl.algorithms.loss.clipped_pg import (
    clipped_pg_actor_objective,
    clipped_pg_diagnostic_metrics,
    materialize_scalar_metrics,
)
from nemo_rl.algorithms.loss.draft import (
    DEFAULT_DRAFT_TOKEN_CHUNK_SIZE,
    DraftLossStats,
    streaming_vocab_parallel_soft_ce,
)
from nemo_rl.algorithms.loss.interfaces import (
    LossFunction,
    LossInputType,
    LossType,
    MetricNormalizer,
)
from nemo_rl.algorithms.utils import (
    calculate_kl,
    compute_seq_logprob_errors,
    masked_mean,
)
from nemo_rl.algorithms.x_token.loss_utils import (
    LocalizedAlignment,
    build_exact_token_map,
    ce_label_mask,
    chunk_average_log_probs,
    get_sparse_projection_matrix,
    next_token_accuracy,
    project_student_to_teacher_vocab,
    select_teacher_topk_indices,
    slice_sparse_projection_cols,
    student_next_token_ce,
    valid_chunk_mask,
)
from nemo_rl.distributed.batched_data_dict import BatchedDataDict
from nemo_rl.distributed.model_utils import (
    cp_shift_next,
    group_all_reduce_sum,
    group_all_reduce_sum_with_grad,
    to_local_if_dtensor,
    vocab_parallel_full_log_softmax,
    vocab_parallel_gather_columns,
    vocab_parallel_log_softmax,
)
from nemo_rl.telemetry.vocabulary import TeedMetric, register_teed_metrics

if TYPE_CHECKING:
    # Import-time only: nemo_rl.algorithms.opd imports the data plane, which
    # would pull a heavy dependency chain into every loss-function consumer.
    from nemo_rl.algorithms.opd import OnPolicyDistillationFullConfig

Tensor = TypeVar("Tensor", bound=torch.Tensor)

KL_PENALTY_KEY = "kl_penalty"
APPROX_ENTROPY_KEY = "approx_entropy"

#: Teed rows for the two metrics the losses below compute. The kl_penalty key
#: is named for the penalty, but the value has the coefficient divided back
#: out, so the series it feeds really is the divergence.
LOSS_TEED_METRICS = (
    TeedMetric(
        KL_PENALTY_KEY,
        "rl.kl.divergence",
        description="KL divergence from the reference policy.",
    ),
    TeedMetric(
        APPROX_ENTROPY_KEY,
        "rl.entropy",
        description="Approximate policy entropy.",
    ),
)

register_teed_metrics(LOSS_TEED_METRICS)


class DraftCrossEntropyLossConfig(TypedDict):
    vocab_parallel_group: Optional[torch.distributed.ProcessGroup]
    # Tokens per FP32 vocab tile in the streaming soft-CE reduction. Defaults to
    # DEFAULT_DRAFT_TOKEN_CHUNK_SIZE; see policy.draft.token_chunk_size.
    token_chunk_size: NotRequired[int]


class DraftCrossEntropyLossDataDict(TypedDict):
    teacher_logits: Tensor
    student_logits: Tensor
    token_mask: Tensor
    sample_mask: Tensor
    student_vocab_indices: NotRequired[Tensor]


class DraftCrossEntropyLossFn(LossFunction):
    """Compute the auxiliary soft-target cross-entropy used for draft-model training."""

    loss_type = LossType.TOKEN_LEVEL
    input_type = LossInputType.DRAFT

    def __init__(
        self,
        vocab_parallel_group: Optional[torch.distributed.ProcessGroup] = None,
        token_chunk_size: int = DEFAULT_DRAFT_TOKEN_CHUNK_SIZE,
    ):
        self.vocab_parallel_group = vocab_parallel_group
        self.token_chunk_size = token_chunk_size

    def __call__(
        self,
        teacher_logits: Tensor,
        student_logits: Tensor,
        token_mask: Tensor,
        data: BatchedDataDict[DraftCrossEntropyLossDataDict],
        global_valid_seqs: torch.Tensor,
        global_valid_toks: torch.Tensor,
    ) -> torch.Tensor:
        """Reduce the masked per-token draft loss to a scalar."""
        stats = self.loss_stats(
            teacher_logits=teacher_logits,
            student_logits=student_logits,
            token_mask=token_mask,
            data=data,
        )
        return stats.normalized(
            normalization_counts=global_valid_toks.reshape(1),
        )

    def loss_stats(
        self,
        teacher_logits: Tensor,
        student_logits: Tensor,
        token_mask: Tensor,
        data: BatchedDataDict[DraftCrossEntropyLossDataDict],
    ) -> DraftLossStats:
        """Return raw one-bin EAGLE statistics for deferred normalization."""
        mask = token_mask * data["sample_mask"].unsqueeze(-1)
        return streaming_vocab_parallel_soft_ce(
            student_logits=student_logits,
            teacher_logits=teacher_logits,
            mask=mask,
            token_chunk_size=self.token_chunk_size,
            tp_group=self.vocab_parallel_group,
        )


class ClippedPGLossConfig(BaseModel, extra="allow"):
    # --- Loss type ---
    disable_ppo_ratio: bool = False
    token_level_loss: bool = True
    # If True, apply the off-policy importance-sampling correction at the
    # sequence level (one weight per generated sample), as in GSPO.
    # If False (default), correction is applied at the token level as in the
    # original GRPO paper.
    sequence_level_importance_ratios: bool = False

    # --- Clipping ---
    ratio_clip_min: float = 0.2
    ratio_clip_max: float = 0.2
    # Dual-clipping value (should be >1 if enabled; usually set to 3 empirically). None to disable.
    ratio_clip_c: Optional[float] = None

    # --- KL regularization ---
    reference_policy_kl_penalty: float = Field(default=0.01, ge=0, allow_inf_nan=False)
    # Can be set to k1, k2, k3
    # For more details, see http://joschu.net/blog/kl-approx.html
    reference_policy_kl_type: str = "k3"
    kl_input_clamp_value: Optional[float] = 20.0
    kl_output_clamp_value: Optional[float] = 10.0
    # If True, add KL penalty to reward instead of loss (used by Reinforce++)
    use_kl_in_reward: bool = False

    # --- Importance sampling correction ---
    # Async GRPO requires importance sampling correction enabled
    # Set to true when async_grpo.enabled is true
    use_importance_sampling_correction: bool = False
    # --- Truncated importance sampling ---
    # Type of truncated importance sampling:
    #   "tis"          – clamp IS weights to [min, max], where min defaults to 0
    #   "icepop"       – zero out tokens with IS weight outside [min, max]
    #   "seq-mask-tis" – zero out sequences by geometric-mean IS ratio, non-truncated token IS correction
    truncated_importance_sampling_type: Optional[str] = None
    truncated_importance_sampling_ratio: Optional[float] = None
    # Lower bound for TIS clipping, ICE-POP filtering, or seq-mask-tis filtering
    truncated_importance_sampling_ratio_min: Optional[float] = None

    # --- On-policy ---
    # (default off) loss formulation improvements (docs/guides/grpo.md#loss)
    use_on_policy_kl_approximation: bool = False
    # If True, force the PPO ratio to 1.0 while preserving its gradient.
    # Generation-to-trainer importance sampling remains independently enabled.
    # NOTE: This should only be used when doing exactly one update per rollout
    # (i.e., num_prompts_per_step * num_generations_per_prompt == train_global_batch_size)
    force_on_policy_ratio: bool = False
    # Evaluate grpo.seq_logprob_error_threshold in the training loss, then
    # normalize accumulated gradients over survivors before the optimizer step.
    # Opt-in; supported only for token-level, force-on-policy Megatron GRPO
    # with the non-streaming trainer.
    seq_logprob_error_in_loss: bool = False
    # If True, use CISPO (Clipped IS-weight Policy Optimization) from MiniMax-M1.
    use_cispo: bool = False
    # VAPO: weight μ for positive-example NLL loss on correct samples.
    # L = L_PPO + μ·L_NLL(correct)   (arXiv:2504.05118, Eq. 10)
    # Set to 0 to disable.
    positive_example_nll_weight: float = 0.0
    # Diagnostic metrics are useful for monitoring, but are not needed for
    # the objective. ``full`` preserves the historical metric surface;
    # ``minimal`` keeps only objective/importance-sampling bookkeeping.
    metrics_level: Literal["full", "minimal"] = "full"
    # Compile the tensor-only actor objective when supported by the active
    # backend. The eager implementation remains the default path.
    enable_torch_compile: bool = False


class ClippedPGLossDataDict(TypedDict):
    """Required keys for the Clipped Policy Gradient loss function."""

    input_ids: torch.Tensor
    advantages: torch.Tensor
    prev_logprobs: torch.Tensor
    generation_logprobs: torch.Tensor
    reference_policy_logprobs: torch.Tensor
    token_mask: torch.Tensor
    sample_mask: torch.Tensor
    __extra__: Any


class ClippedPGLossFn(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 [ log π_θ(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.
    """

    input_type = LossInputType.LOGPROB

    def __init__(
        self,
        cfg: ClippedPGLossConfig,
        use_fused_linear_logprobs: bool = False,
        opd_full: Optional["OnPolicyDistillationFullConfig"] = None,
        *,
        seq_logprob_error_threshold: float | None = None,
    ):
        """Initialize the loss and its worker normalization requirements.

        Args:
            cfg: Policy-gradient loss configuration.
            use_fused_linear_logprobs: Whether the model returns precomputed
                next-token logprobs instead of logits.
            opd_full: Optional full-vocabulary distillation configuration.
            seq_logprob_error_threshold: Required when
                ``cfg.seq_logprob_error_in_loss`` is enabled; otherwise unused.
                The loss emits survivor counts but normalizes by
                the original global counts. The worker must aggregate survivor
                counts across all microbatches and data-parallel ranks, then
                rescale gradients and normalized metrics before clipping or
                stepping the optimizer. See ``requires_survivor_normalization``.
        """
        self.seq_logprob_error_in_loss = cfg.seq_logprob_error_in_loss
        self.seq_logprob_error_threshold = seq_logprob_error_threshold
        if self.seq_logprob_error_in_loss:
            if seq_logprob_error_threshold is None:
                raise ValueError(
                    "loss_fn.seq_logprob_error_in_loss requires seq_logprob_error_threshold"
                )
            if not cfg.force_on_policy_ratio or not cfg.token_level_loss:
                raise ValueError(
                    "In-loss sequence filtering requires force_on_policy_ratio "
                    "and token_level_loss"
                )
            if cfg.positive_example_nll_weight != 0:
                raise ValueError("In-loss sequence filtering does not support NLL")
        # When True, the model forward is patched to return precomputed next-token
        # logprobs (via chunked linear CE fusion) instead of full logits. This is
        # consumed by prepare_loss_input, which short-circuits the logits->logprobs
        # conversion. See nemo_rl/distributed/model_utils.py for the fused forward.
        self.use_fused_linear_logprobs = use_fused_linear_logprobs
        self.opd_full = opd_full if opd_full is not None and opd_full.enabled else None
        self.input_type = (
            LossInputType.OPD_FULL
            if self.opd_full is not None
            else LossInputType.LOGPROB
        )
        if self.opd_full is not None:
            _validate_opd_full_loss_config(cfg, use_fused_linear_logprobs)
        self.disable_ppo_ratio = cfg.disable_ppo_ratio
        self.ratio_clip_min = cfg.ratio_clip_min
        self.ratio_clip_max = cfg.ratio_clip_max
        self.ratio_clip_c = cfg.ratio_clip_c  # set to None to disable dual-clipping
        if self.ratio_clip_c is not None:
            assert self.ratio_clip_c > 1, (
                f"ratio_clip_c must exceed 1 representing a lower bound of the ratios, got {self.ratio_clip_c}."
            )
        self.reference_policy_kl_penalty = (
            cfg.reference_policy_kl_penalty if not cfg.use_kl_in_reward else 0
        )
        self.reference_policy_kl_type = cfg.reference_policy_kl_type
        self.kl_input_clamp_value = cfg.kl_input_clamp_value
        self.kl_output_clamp_value = cfg.kl_output_clamp_value
        self.use_importance_sampling_correction = cfg.use_importance_sampling_correction
        # Type of truncated importance sampling: "tis" | "icepop" | "seq-mask-tis"
        self.truncated_importance_sampling_type = cfg.truncated_importance_sampling_type
        self.truncated_importance_sampling_ratio = (
            cfg.truncated_importance_sampling_ratio
        )
        # Lower bound for TIS clipping, ICE-POP filtering, or seq-mask-tis filtering
        self.truncated_importance_sampling_ratio_min = (
            cfg.truncated_importance_sampling_ratio_min
        )
        self.use_on_policy_kl_approximation = cfg.use_on_policy_kl_approximation
        self.force_on_policy_ratio = cfg.force_on_policy_ratio  # Force ratio to 1.0

        # Whether to compute importance weights per-sequence instead of per-token.
        self.sequence_level_importance_ratios = cfg.sequence_level_importance_ratios
        self.positive_example_nll_weight = cfg.positive_example_nll_weight
        self.metrics_level = cfg.metrics_level
        self.enable_torch_compile = cfg.enable_torch_compile
        # Compile the actor objective lazily on the worker after Ray has
        # serialized the loss function.
        self._compiled_actor_objective = None
        self.loss_type = (
            LossType.TOKEN_LEVEL if cfg.token_level_loss else LossType.SEQUENCE_LEVEL
        )
        if self.sequence_level_importance_ratios:
            assert self.loss_type == LossType.SEQUENCE_LEVEL, (
                "sequence-level importance sampling (e.g. GSPO) is mutually exclusive with token-level loss"
            )

        self.use_cispo = cfg.use_cispo
        if self.use_cispo:
            assert not self.disable_ppo_ratio, (
                "use_cispo is incompatible with disable_ppo_ratio; "
                "CISPO needs the pi_theta/pi_theta_old ratio but disable_ppo_ratio removes it"
            )
            assert not self.force_on_policy_ratio, (
                "use_cispo is incompatible with force_on_policy_ratio; "
                "forcing ratio=1 removes the clipped IS-weight that CISPO optimizes"
            )
            assert not self.sequence_level_importance_ratios, (
                "use_cispo is incompatible with sequence_level_importance_ratios; "
                "CISPO uses token-level importance weights"
            )
            assert self.ratio_clip_c is None, (
                "use_cispo is incompatible with dual clipping (ratio_clip_c); "
                "the dual-clip block runs after the CISPO loss assembly and would "
                "silently overwrite it. Set ratio_clip_c=null when use_cispo=True."
            )
            assert self.loss_type == LossType.TOKEN_LEVEL, (
                "use_cispo requires token_level_loss=True (LossType.TOKEN_LEVEL)."
            )
        if self.truncated_importance_sampling_type is not None:
            assert self.use_importance_sampling_correction, (
                "truncated importance sampling is only supported when use_importance_sampling_correction is True"
            )
            assert self.truncated_importance_sampling_type in (
                "tis",
                "icepop",
                "seq-mask-tis",
            ), (
                f"truncated_importance_sampling_type must be 'tis', 'icepop', or 'seq-mask-tis', "
                f"got {self.truncated_importance_sampling_type}"
            )
            assert (
                self.truncated_importance_sampling_ratio is not None
                and self.truncated_importance_sampling_ratio > 0
            ), "truncated_importance_sampling_ratio should be positive"
            if self.truncated_importance_sampling_ratio_min is not None:
                assert (
                    self.truncated_importance_sampling_ratio_min
                    <= self.truncated_importance_sampling_ratio
                ), (
                    "truncated_importance_sampling_ratio_min must be <= "
                    "truncated_importance_sampling_ratio"
                )
            if self.truncated_importance_sampling_type in ("icepop", "seq-mask-tis"):
                assert self.truncated_importance_sampling_ratio_min is not None, (
                    "truncated_importance_sampling_ratio_min should be set when truncated_importance_sampling_type is 'icepop' or 'seq-mask-tis'"
                )
            if self.truncated_importance_sampling_type == "seq-mask-tis":
                assert not self.sequence_level_importance_ratios, (
                    "seq-mask-tis uses token-level IS correction with sequence-level masking, "
                    "and is incompatible with sequence_level_importance_ratios=True"
                )

        # Advertise, per returned metric, the global denominator it was
        # normalized by (see MetricNormalizer). Built here — next to the flags
        # that pick the denominators — so split-API trainers can undo the
        # placeholder global_valid_*=1 normalization without maintaining a
        # consumer-side table. Keep in sync with __call__'s return dict.
        grad_normalizer = (
            MetricNormalizer.TOKENS
            if self.loss_type == LossType.TOKEN_LEVEL
            else MetricNormalizer.SEQUENCES
        )
        self.metric_normalizations: dict[str, MetricNormalizer] = {
            # Normalized like the gradient (loss_type-dependent).
            "loss": grad_normalizer,
            KL_PENALTY_KEY: grad_normalizer,
            # Token-normalized diagnostics, independent of loss_type.
            "probs_ratio": MetricNormalizer.TOKENS,
            "probs_ratio_clamped": MetricNormalizer.TOKENS,
            "token_mult_prob_error": MetricNormalizer.TOKENS,
            "gen_kl_error": MetricNormalizer.TOKENS,
            "policy_kl_error": MetricNormalizer.TOKENS,
            "js_divergence_error": MetricNormalizer.TOKENS,
            APPROX_ENTROPY_KEY: MetricNormalizer.TOKENS,
            # Keyed on sequence_level_importance_ratios, NOT loss_type.
            "sampling_importance_ratio": (
                MetricNormalizer.SEQUENCES
                if self.sequence_level_importance_ratios
                else MetricNormalizer.TOKENS
            ),
            # Raw count — the downstream per-microbatch sum IS the value.
            "policy_support_excluded_tokens": MetricNormalizer.NONE,
            "num_valid_samples": MetricNormalizer.NONE,
            "seq_logprob_error_valid_tokens": MetricNormalizer.NONE,
            "seq_logprob_error_valid_seqs": MetricNormalizer.NONE,
            "num_masked_seqs_by_logprob_error": MetricNormalizer.NONE,
            # Normalized by the microbatch's own correct-token count, not a
            # global factor — already a per-microbatch mean.
            "positive_nll_loss": MetricNormalizer.NONE,
            # Extrema — combined downstream with min/max, never scaled.
            "probs_ratio_min": MetricNormalizer.NONE,
            "probs_ratio_max": MetricNormalizer.NONE,
            "probs_ratio_clamped_min": MetricNormalizer.NONE,
            "probs_ratio_clamped_max": MetricNormalizer.NONE,
        }
        if self.metrics_level == "minimal":
            minimal_metrics = {
                "loss",
                "kl_penalty",
                "num_valid_samples",
                "policy_support_excluded_tokens",
                "positive_nll_loss",
                "token_mult_prob_error",
            }
            if self.use_importance_sampling_correction:
                minimal_metrics.add("sampling_importance_ratio")
            if self.seq_logprob_error_in_loss:
                minimal_metrics.update(
                    {
                        "seq_logprob_error_valid_tokens",
                        "seq_logprob_error_valid_seqs",
                        "num_masked_seqs_by_logprob_error",
                    }
                )
            self.metric_normalizations = {
                name: normalizer
                for name, normalizer in self.metric_normalizations.items()
                if name in minimal_metrics
            }
        if self.truncated_importance_sampling_type is not None:
            # Keyed on the TIS type, NOT loss_type: seq-mask-tis masks whole
            # sequences (÷ global_valid_seqs); tis/icepop are token-level.
            self.metric_normalizations["is_oob_ratio"] = (
                MetricNormalizer.SEQUENCES
                if self.truncated_importance_sampling_type == "seq-mask-tis"
                else MetricNormalizer.TOKENS
            )
        if self.opd_full is not None:
            # opd_full returns its own metric set from _opd_full_call; the
            # policy-gradient diagnostics above are never produced there.
            self.metric_normalizations = {
                "loss": grad_normalizer,
                KL_PENALTY_KEY: grad_normalizer,
                "num_valid_samples": MetricNormalizer.NONE,
                "opd_full_reverse_kl": MetricNormalizer.TOKENS,
                "opd_full_reverse_kl_min": MetricNormalizer.NONE,
                "opd_full_reverse_kl_max": MetricNormalizer.NONE,
            }
            if self.opd_full.validate_decomposition:
                self.metric_normalizations.update(
                    {
                        "opd_full_entropy": MetricNormalizer.TOKENS,
                        "opd_full_cross_entropy": MetricNormalizer.TOKENS,
                        "opd_full_decomposition_error": MetricNormalizer.NONE,
                    }
                )

    @property
    def requires_survivor_normalization(self) -> bool:
        """Whether the worker must renormalize using in-loss survivor counts.

        Uses the same switch as filtering so the loss and worker agree.
        """
        return self.seq_logprob_error_in_loss

    def __call__(
        self,
        next_token_logprobs: Optional[Tensor] = None,
        data: Optional[BatchedDataDict[ClippedPGLossDataDict]] = None,
        global_valid_seqs: Optional[torch.Tensor] = None,
        global_valid_toks: Optional[torch.Tensor] = None,
        opd_full_divergence: Optional[Tensor] = None,
        opd_full_entropy: Optional[Tensor] = None,
        opd_full_cross_entropy: Optional[Tensor] = None,
        policy_support_mask: Optional[Tensor] = None,
    ) -> tuple[torch.Tensor, dict]:
        """Clipped Policy Gradient RL loss, or the full-vocabulary MOPD reverse KL.

        Which objective runs is fixed at construction by ``opd_full``.

        Args:
            next_token_logprobs: Sampled-token log-probabilities ``[B, S - 1]``.
                Required on the policy-gradient branch; on the ``opd_full``
                branch only when ``reference_policy_kl_penalty`` is non-zero.
            data: Microbatch with masks, advantages, and prior log-probabilities.
            global_valid_seqs: Global valid-sequence count for normalization.
            global_valid_toks: Global valid-token count for normalization.
            policy_support_mask: Current top-k/top-p support, supplied only with filtering.
                Negative infinity in other policy logprobs is also excluded.
            opd_full_divergence: Per-token reverse KL ``[B, S - 1]``, required on
                the ``opd_full`` branch.
            opd_full_entropy: Optional ``sum_v p_s log p_s`` diagnostic.
            opd_full_cross_entropy: Optional ``-sum_v p_s log p_t`` diagnostic.

        Returns:
            Tuple of the scalar loss and its metric dict.
        """
        assert data is not None, "ClippedPGLossFn requires data"
        assert global_valid_seqs is not None, (
            "ClippedPGLossFn requires global_valid_seqs"
        )
        assert global_valid_toks is not None, (
            "ClippedPGLossFn requires global_valid_toks"
        )
        if self.opd_full is not None:
            return self._opd_full_call(
                data=data,
                global_valid_seqs=global_valid_seqs,
                global_valid_toks=global_valid_toks,
                next_token_logprobs=next_token_logprobs,
                opd_full_divergence=opd_full_divergence,
                opd_full_entropy=opd_full_entropy,
                opd_full_cross_entropy=opd_full_cross_entropy,
            )
        assert next_token_logprobs is not None, (
            "ClippedPGLossFn requires next_token_logprobs"
        )
        curr_logprobs = next_token_logprobs
        token_mask = data["token_mask"][:, 1:]
        sample_mask = data["sample_mask"]
        advantages = data["advantages"][:, 1:]
        # Skip loading prev_logprobs when force_on_policy_ratio=True (will use curr_logprobs instead)
        prev_logprobs = (
            None if self.force_on_policy_ratio else data["prev_logprobs"][:, 1:]
        )
        generation_logprobs = data["generation_logprobs"][:, 1:]
        if self.reference_policy_kl_penalty != 0:
            reference_policy_logprobs = data["reference_policy_logprobs"][:, 1:]
            curr_logprobs_unfiltered = data.get(
                "curr_logprobs_unfiltered", curr_logprobs
            )

        mask = token_mask * sample_mask.unsqueeze(-1)
        full_metrics = self.metrics_level == "full"
        need_importance_weights = (
            full_metrics or self.use_importance_sampling_correction
        )

        # For truly on-policy training, use curr_logprobs as prev_logprobs
        # This avoids computing prev_logprobs upstream
        if self.force_on_policy_ratio:
            prev_logprobs = curr_logprobs.detach()

        # Only top-k/top-p can mark -inf as an expected support mismatch.
        # Keep NaN/+inf observable, and preserve strict checks without filtering.
        if policy_support_mask is not None:
            policy_support_mask = (
                policy_support_mask
                & ~torch.isneginf(prev_logprobs)
                & ~torch.isneginf(generation_logprobs)
            )
        seq_error_token_mask = token_mask
        if policy_support_mask is not None:
            seq_error_token_mask = token_mask * policy_support_mask.to(token_mask.dtype)
        seq_error_metrics = {}
        if self.seq_logprob_error_in_loss:
            errors, _ = compute_seq_logprob_errors(
                policy_logprobs=curr_logprobs.detach(),
                generation_logprobs=generation_logprobs,
                token_mask=seq_error_token_mask,
                sample_mask=sample_mask,
            )
            filtered_sample_mask = sample_mask * (
                errors <= self.seq_logprob_error_threshold
            ).to(sample_mask.dtype)
            seq_error_metrics["num_masked_seqs_by_logprob_error"] = (
                (sample_mask - filtered_sample_mask).sum().item()
            )
            sample_mask = filtered_sample_mask
            mask = token_mask * sample_mask.unsqueeze(-1)
            seq_error_metrics["seq_logprob_error_valid_tokens"] = mask.sum().item()
            seq_error_metrics["seq_logprob_error_valid_seqs"] = sample_mask.sum().item()
            # A rejected nonfinite logprob must not poison a zero-weight loss.
            curr_logprobs = torch.where(mask.bool(), curr_logprobs, 0.0)
            prev_logprobs = curr_logprobs.detach()
            generation_logprobs = torch.where(mask.bool(), generation_logprobs, 0.0)
            if self.reference_policy_kl_penalty != 0:
                curr_logprobs_unfiltered = torch.where(
                    mask.bool(),
                    data.get("curr_logprobs_unfiltered", next_token_logprobs),
                    0.0,
                )
                reference_policy_logprobs = torch.where(
                    mask.bool(), data["reference_policy_logprobs"][:, 1:], 0.0
                )

        # Reference KL uses the unfiltered policy and retains its own token set.
        # Actor support mismatches must not discard otherwise valid KL terms.
        kl_token_mask = token_mask
        kl_sample_mask = sample_mask
        kl_mask = mask
        kl_generation_logprobs = torch.where(kl_mask.bool(), generation_logprobs, 0.0)

        # Keep support information until all three policy distributions are
        # available. Replacing -inf with zero earlier turns an impossible token
        # into log-probability one in token/sequence importance ratios.
        if policy_support_mask is not None:
            token_mask = token_mask * policy_support_mask.to(token_mask.dtype)
            sample_mask = sample_mask * token_mask.bool().any(dim=-1).to(
                sample_mask.dtype
            )
        mask = token_mask * sample_mask.unsqueeze(-1)
        curr_logprobs = torch.where(mask.bool(), curr_logprobs, 0.0)
        prev_logprobs = torch.where(mask.bool(), prev_logprobs, 0.0)
        generation_logprobs = torch.where(mask.bool(), generation_logprobs, 0.0)
        if self.seq_logprob_error_in_loss:
            seq_error_metrics["seq_logprob_error_valid_tokens"] = mask.sum().item()
            seq_error_metrics["seq_logprob_error_valid_seqs"] = sample_mask.sum().item()
        # Calculate KL regularization.
        if self.reference_policy_kl_penalty != 0:
            # When top-k/top-p filtering is enabled, we need special handling for KL:
            # - reference_policy_logprobs is computed **without** filtering (see use_reference_model)
            # - curr_logprobs/prev_logprobs are computed **with** filtering (for actor loss compatibility)
            # - For KL, we need curr_logprobs **without** filtering to be consistent with ref logprobs
            # - For importance weights, we also use unfiltered curr_logprobs_unfiltered since we're
            #   reweighting samples from π_gen_filtered to π_curr_unfiltered

            # On-policy KL approximation
            # KL samples come from the optimized policy, so the KL loss must include
            # the score-function gradient through the sampling probability; see
            # https://arxiv.org/abs/2506.09477v1. In the non-IS case,
            # exp(x - x.detach()) has forward value 1 while preserving that gradient.
            if self.use_on_policy_kl_approximation:
                # See: docs/guides/grpo.md#on-policy-kl-approximation
                kl_importance_weights = torch.exp(
                    curr_logprobs_unfiltered - kl_generation_logprobs
                )
            else:
                kl_importance_weights = torch.exp(
                    curr_logprobs_unfiltered - curr_logprobs_unfiltered.detach()
                )
            kl_importance_weights = torch.nan_to_num(
                kl_importance_weights, nan=0.0, posinf=0.0, neginf=0.0
            )

            # Compute KL loss
            kl = self.reference_policy_kl_penalty * calculate_kl(
                logprobs=curr_logprobs_unfiltered,
                logprobs_reference=reference_policy_logprobs,
                kl_type=self.reference_policy_kl_type,
                input_clamp_value=self.kl_input_clamp_value,
                output_clamp_value=self.kl_output_clamp_value,
                importance_sampling_weights=kl_importance_weights,
            )

            # Reduce KL loss
            if self.loss_type == LossType.TOKEN_LEVEL:
                kl = masked_mean(
                    kl, kl_mask, global_normalization_factor=global_valid_toks
                )
            else:
                kl = masked_mean(
                    masked_mean(kl, kl_token_mask, dim=-1),
                    kl_sample_mask,
                    global_normalization_factor=global_valid_seqs,
                )
        else:
            kl = curr_logprobs.new_zeros(())

        # -------------------------------------------------------------
        # Off-policy (actor) importance-sampling correction
        # -------------------------------------------------------------
        _is_filter_metrics: dict = {}  # populated for icepop / seq-mask-tis
        # See: docs/guides/grpo.md#importance-sampling-correction
        if not need_importance_weights:
            # No objective or emitted metric needs the actor/generation
            # importance weights in this mode.
            actor_importance_weights_expanded = None
        elif self.sequence_level_importance_ratios:
            # importance weight w_i = exp(Σ_t (log π_actor − log π_behaviour))
            seq_lp_diff = ((prev_logprobs - generation_logprobs) * mask).sum(dim=-1)
            actor_importance_weights = torch.exp(seq_lp_diff).detach()
            actor_importance_weights = torch.nan_to_num(
                actor_importance_weights, nan=0.0, posinf=0.0, neginf=0.0
            )
            # Broadcast to token dimension so we can reuse existing reduction
            actor_importance_weights_expanded = actor_importance_weights.unsqueeze(-1)
        else:
            # Token-level correction
            actor_importance_weights_expanded = torch.exp(
                prev_logprobs - generation_logprobs
            )
            actor_importance_weights_expanded = torch.nan_to_num(
                actor_importance_weights_expanded, nan=0.0, posinf=0.0, neginf=0.0
            )
        # ---- Truncated Importance Sampling ----
        # "tis"          – clamp IS weights to [min, max], where min defaults to 0
        # "icepop"       – zero out tokens whose IS weight ∉ [min, max]   (ref bounds: 0.5–5)
        # "seq-mask-tis" – zero out entire sequences whose geometric-mean
        #                  IS ratio ∉ [min, max]; retained sequences keep
        #                  raw (non-truncated) token-level IS weights      (ref bounds: 0.999–1.002)
        #   Blog: https://yingru.notion.site/When-Speed-Kills-Stability-Demystifying-RL-Collapse-from-the-Training-Inference-Mismatch-271211a558b7808d8b12d403fd15edda
        # is_oob_ratio: fraction of tokens (tis/icepop) or sequences (seq-mask-tis)
        # whose importance weight falls outside the truncation bounds. Each microbatch
        # contributes its out-of-bounds count divided by the *global* valid token/seq
        # count, so the np.sum aggregation in grpo.py recovers the correct global fraction.
        if self.truncated_importance_sampling_ratio is not None:
            if self.truncated_importance_sampling_type == "tis":
                tis_min = self.truncated_importance_sampling_ratio_min
                if tis_min is None:
                    tis_min = 0.0
                token_oob_mask = (
                    actor_importance_weights_expanded
                    > self.truncated_importance_sampling_ratio
                ) | (actor_importance_weights_expanded < tis_min)
                _is_filter_metrics = {
                    "is_oob_ratio": masked_mean(
                        token_oob_mask.float(),
                        mask,
                        global_normalization_factor=global_valid_toks,
                    ),
                }
                actor_importance_weights_expanded = torch.clamp(
                    actor_importance_weights_expanded,
                    min=tis_min,
                    max=self.truncated_importance_sampling_ratio,
                )
            elif self.truncated_importance_sampling_type == "icepop":
                token_kept_mask = (
                    actor_importance_weights_expanded
                    >= self.truncated_importance_sampling_ratio_min
                ) & (
                    actor_importance_weights_expanded
                    <= self.truncated_importance_sampling_ratio
                )
                _is_filter_metrics = {
                    "is_oob_ratio": masked_mean(
                        (~token_kept_mask).float(),
                        mask,
                        global_normalization_factor=global_valid_toks,
                    ),
                }
                actor_importance_weights_expanded = torch.where(
                    token_kept_mask,
                    actor_importance_weights_expanded,
                    torch.zeros_like(actor_importance_weights_expanded),
                )
            elif self.truncated_importance_sampling_type == "seq-mask-tis":
                # geo_mean_i = exp( mean_t( log(π_prev / π_gen) ) )
                log_is_ratio = torch.nan_to_num(
                    prev_logprobs - generation_logprobs,
                    nan=0.0,
                    posinf=0.0,
                    neginf=0.0,
                )
                seq_log_is_ratio_mean = masked_mean(
                    log_is_ratio, token_mask, dim=-1
                )  # [B]
                seq_geomean_is_ratio = torch.exp(seq_log_is_ratio_mean).detach()  # [B]
                seq_kept_mask = (
                    (
                        seq_geomean_is_ratio
                        >= self.truncated_importance_sampling_ratio_min
                    )
                    & (seq_geomean_is_ratio <= self.truncated_importance_sampling_ratio)
                ).float()  # [B]
                _is_filter_metrics = {
                    "is_oob_ratio": masked_mean(
                        1.0 - seq_kept_mask,
                        sample_mask,
                        global_normalization_factor=global_valid_seqs,
                    ),
                }
                actor_importance_weights_expanded = (
                    actor_importance_weights_expanded * seq_kept_mask.unsqueeze(-1)
                )
            else:
                raise ValueError(
                    f"Invalid truncated importance sampling type: {self.truncated_importance_sampling_type}"
                )

        actor_importance_weights = actor_importance_weights_expanded
        del actor_importance_weights_expanded
        if self.use_importance_sampling_correction:
            importance_weights_to_use = actor_importance_weights
        else:
            importance_weights_to_use = None

        if self.enable_torch_compile and self._compiled_actor_objective is None:
            self._compiled_actor_objective = torch.compile(
                clipped_pg_actor_objective,
                dynamic=True,
                fullgraph=True,
                mode="max-autotune-no-cudagraphs",
            )
        actor_objective = self._compiled_actor_objective or clipped_pg_actor_objective
        actor_loss, ratios, ratios_clamped = actor_objective(
            curr_logprobs,
            prev_logprobs=prev_logprobs,
            advantages=advantages,
            token_mask=token_mask,
            mask=mask,
            sample_mask=sample_mask,
            importance_weights=importance_weights_to_use,
            global_valid_toks=global_valid_toks,
            global_valid_seqs=global_valid_seqs,
            ratio_clip_min=self.ratio_clip_min,
            ratio_clip_max=self.ratio_clip_max,
            ratio_clip_c=(self.ratio_clip_c if self.ratio_clip_c is not None else -1.0),
            disable_ppo_ratio=self.disable_ppo_ratio,
            force_on_policy_ratio=self.force_on_policy_ratio,
            sequence_level_importance_ratios=self.sequence_level_importance_ratios,
            use_cispo=self.use_cispo,
            token_level_loss=self.loss_type == LossType.TOKEN_LEVEL,
        )

        # Metric: sampling importance ratio (mean over samples)
        # See: docs/guides/grpo.md#sampling-importance-ratio
        if need_importance_weights:
            if self.sequence_level_importance_ratios:
                sample_importance_ratio = masked_mean(
                    actor_importance_weights.squeeze(-1),
                    sample_mask,
                    global_normalization_factor=global_valid_seqs,
                )
            else:
                sample_importance_ratio = masked_mean(
                    actor_importance_weights,
                    mask,
                    global_normalization_factor=global_valid_toks,
                )

        # -----------------------------------------------------------------
        # VAPO: positive-example NLL loss on correct samples (reward > 0)
        # L = L_PPO + μ · L_NLL(correct)
        # -----------------------------------------------------------------
        nll_loss = curr_logprobs.new_zeros(())
        if self.positive_example_nll_weight > 0 and "rewards" in data:
            correct_sample_mask = (data["rewards"] > 0).float()  # [batch]
            correct_mask = mask * correct_sample_mask.unsqueeze(-1)
            correct_valid_toks = correct_mask.sum()
            nll_loss = torch.where(
                correct_valid_toks > 0,
                masked_mean(
                    -curr_logprobs,
                    correct_mask,
                    global_normalization_factor=correct_valid_toks,
                ),
                curr_logprobs.new_zeros(()),
            )

        loss = actor_loss + kl + self.positive_example_nll_weight * nll_loss
        # Keep all reductions on device and perform one scalar materialization
        # for the Python-facing metric dictionary. This avoids a synchronization
        # for every individual metric while preserving the historical values.

        # If you provided a global_valid_{seqs/toks}, all metrics here are globally normalized
        # by either sequence or token count, depending on particular metric.
        # To get the true metric, you'll need to sum over the microbatch.
        metric_tensors = {"loss": loss.detach()}
        with torch.no_grad():
            lp_error = torch.abs(generation_logprobs - prev_logprobs)
            metric_tensors["token_mult_prob_error"] = masked_mean(
                torch.exp(lp_error * mask),
                mask,
                global_normalization_factor=global_valid_toks,
            )
        if full_metrics:
            metric_tensors.update(
                clipped_pg_diagnostic_metrics(
                    curr_logprobs=curr_logprobs,
                    prev_logprobs=prev_logprobs,
                    generation_logprobs=generation_logprobs,
                    ratios=ratios,
                    ratios_clamped=ratios_clamped,
                    mask=mask,
                    global_valid_toks=global_valid_toks,
                    reference_policy_kl_type=self.reference_policy_kl_type,
                )
            )
        if need_importance_weights:
            metric_tensors["sampling_importance_ratio"] = sample_importance_ratio

        metric_tensors.update(
            {
                KL_PENALTY_KEY: (
                    kl / self.reference_policy_kl_penalty
                    if self.reference_policy_kl_penalty != 0
                    else curr_logprobs.new_zeros(())
                ),
                "num_valid_samples": sample_mask.sum(),
                "positive_nll_loss": nll_loss,
                **_is_filter_metrics,
            }
        )
        metric_tensors["policy_support_excluded_tokens"] = (kl_mask - mask).sum()
        return loss, {
            **materialize_scalar_metrics(metric_tensors, loss.dtype),
            **seq_error_metrics,
        }

    def _opd_full_call(
        self,
        *,
        data: BatchedDataDict[ClippedPGLossDataDict],
        global_valid_seqs: torch.Tensor,
        global_valid_toks: torch.Tensor,
        next_token_logprobs: Optional[Tensor],
        opd_full_divergence: Optional[Tensor],
        opd_full_entropy: Optional[Tensor],
        opd_full_cross_entropy: Optional[Tensor],
    ) -> tuple[torch.Tensor, dict]:
        """Full-vocabulary MOPD objective: the exact reverse KL is the whole loss.

        ``prepare_loss_input`` reconstructs the teacher distribution and runs the
        distributed divergence kernel, so this only masks and normalizes. No
        policy-gradient ratio, advantage, or importance-sampling term applies:
        the reverse KL is an exact expectation over the student distribution and
        does not depend on which token was sampled.

        Args:
            data: Microbatch with ``token_mask`` and ``sample_mask``.
            global_valid_seqs: Global valid-sequence count for normalization.
            global_valid_toks: Global valid-token count for normalization.
            next_token_logprobs: Current-policy sampled-token log-probabilities,
                only needed when ``reference_policy_kl_penalty`` is non-zero.
            opd_full_divergence: Per-token reverse KL ``[B, S - 1]``.
            opd_full_entropy: Optional ``sum_v p_s log p_s`` diagnostic.
            opd_full_cross_entropy: Optional ``-sum_v p_s log p_t`` diagnostic.

        Returns:
            Tuple of the scalar loss and its metric dict.

        Raises:
            ValueError: If the divergence tensor is missing.
        """
        if opd_full_divergence is None:
            raise ValueError(
                "opd_full requires opd_full_divergence from prepare_loss_input."
            )
        token_mask = data["token_mask"][:, 1:]
        sample_mask = data["sample_mask"]
        mask = token_mask * sample_mask.unsqueeze(-1)

        def reduce_like_loss(per_token_values: Tensor) -> torch.Tensor:
            """Reduce a per-token tensor the way the gradient is normalized.

            SEQUENCE_LEVEL averages within each sequence first, so every sequence
            contributes equally regardless of length; a flat masked_mean against
            ``global_valid_seqs`` would instead divide a token sum by a sequence
            count and scale the loss by the mean sequence length.
            """
            if self.loss_type == LossType.TOKEN_LEVEL:
                return masked_mean(
                    per_token_values,
                    mask,
                    global_normalization_factor=global_valid_toks,
                )
            return masked_mean(
                masked_mean(per_token_values, token_mask, dim=-1),
                sample_mask,
                global_normalization_factor=global_valid_seqs,
            )

        actor_loss = reduce_like_loss(opd_full_divergence)

        kl = torch.tensor(0.0, device=actor_loss.device)
        if self.reference_policy_kl_penalty != 0:
            if next_token_logprobs is None:
                raise ValueError(
                    "opd_full with reference_policy_kl_penalty != 0 requires "
                    "next_token_logprobs from prepare_loss_input: the penalty must "
                    "be differentiable through the current policy."
                )
            # Unlike the exact divergence above, this penalty is sampled, so its
            # gradient needs the score-function term through the sampling
            # probability, as in the policy-gradient branch.
            # exp(x - x.detach()) is 1 in the forward pass and supplies it.
            kl_importance_weights = torch.nan_to_num(
                torch.exp(next_token_logprobs - next_token_logprobs.detach()),
                nan=0.0,
                posinf=0.0,
                neginf=0.0,
            )
            kl = self.reference_policy_kl_penalty * calculate_kl(
                logprobs=next_token_logprobs,
                logprobs_reference=data["reference_policy_logprobs"][:, 1:],
                kl_type=self.reference_policy_kl_type,
                input_clamp_value=self.kl_input_clamp_value,
                output_clamp_value=self.kl_output_clamp_value,
                importance_sampling_weights=kl_importance_weights,
            )
            kl = reduce_like_loss(kl)

        loss = actor_loss + kl

        with torch.no_grad():
            masked_divergence = opd_full_divergence.detach()[mask.bool()]
            # Token-normalized diagnostics use global_valid_toks regardless of
            # loss_type, matching their declared MetricNormalizer.TOKENS.
            reverse_kl_metric = masked_mean(
                opd_full_divergence.detach(),
                mask,
                global_normalization_factor=global_valid_toks,
            )
            metrics: dict[str, Any] = {
                "loss": loss.item(),
                # Report the raw KL, undoing the coefficient. Guard on the
                # coefficient: `kl` is 0-dim, so its `numel()` is always 1 and
                # cannot stand in for "the penalty is active".
                KL_PENALTY_KEY: (
                    kl.item() / self.reference_policy_kl_penalty
                    if self.reference_policy_kl_penalty
                    else 0
                ),
                "num_valid_samples": sample_mask.sum().item(),
                "opd_full_reverse_kl": reverse_kl_metric.item(),
                "opd_full_reverse_kl_min": (
                    masked_divergence.min().item()
                    if masked_divergence.numel() > 0
                    else float("inf")
                ),
                "opd_full_reverse_kl_max": (
                    masked_divergence.max().item()
                    if masked_divergence.numel() > 0
                    else float("-inf")
                ),
            }
            if (
                opd_full_entropy is not None
                and opd_full_cross_entropy is not None
                and self.opd_full is not None
                and self.opd_full.validate_decomposition
            ):
                # reverse_kl == sum_v p_s log p_s + (-sum_v p_s log p_t)
                decomposition = (
                    opd_full_entropy.detach() + opd_full_cross_entropy.detach()
                )
                metrics["opd_full_entropy"] = masked_mean(
                    -opd_full_entropy.detach(),
                    mask,
                    global_normalization_factor=global_valid_toks,
                ).item()
                metrics["opd_full_cross_entropy"] = masked_mean(
                    opd_full_cross_entropy.detach(),
                    mask,
                    global_normalization_factor=global_valid_toks,
                ).item()
                metrics["opd_full_decomposition_error"] = (
                    (decomposition - opd_full_divergence.detach())
                    .abs()
                    .mul(mask)
                    .max()
                    .item()
                )

        return loss, metrics


def _validate_opd_full_loss_config(
    cfg: ClippedPGLossConfig, use_fused_linear_logprobs: bool
) -> None:
    """Reject loss knobs that have no code path under the full-vocabulary objective.

    ``opd_full`` replaces the whole policy-gradient objective with an exact
    reverse KL, so the ratio, CISPO, and importance-sampling machinery is dead
    code on that branch. Erroring out beats silently ignoring user config.

    Args:
        cfg: The clipped-PG loss config supplied by the user.
        use_fused_linear_logprobs: Whether the fused linear-CE forward is on.

    Raises:
        ValueError: If any incompatible option is set.
    """
    if use_fused_linear_logprobs:
        raise ValueError(
            "opd_full is incompatible with use_fused_linear_logprobs: the fused "
            "forward returns precomputed logprobs, but the full-vocabulary "
            "reverse KL needs the raw vocabulary-parallel logits."
        )

    # Policy-gradient ratio machinery: no ratio exists under opd_full.
    if not cfg.disable_ppo_ratio:
        raise ValueError("opd_full requires loss_fn.disable_ppo_ratio=true.")
    if cfg.ratio_clip_c is not None:
        raise ValueError("opd_full is incompatible with dual PPO clipping.")
    if cfg.use_cispo:
        raise ValueError("opd_full is incompatible with CISPO.")
    if cfg.force_on_policy_ratio:
        raise ValueError("opd_full is incompatible with force_on_policy_ratio.")
    if cfg.use_on_policy_kl_approximation:
        raise ValueError(
            "opd_full is incompatible with use_on_policy_kl_approximation: it "
            "reweights the reference-KL term by a generation-to-current ratio "
            "that only the policy-gradient branch computes. The opd_full branch "
            "never reads it, so leaving it on would be silently ignored."
        )
    if cfg.sequence_level_importance_ratios:
        raise ValueError(
            "opd_full is incompatible with sequence_level_importance_ratios."
        )

    # Importance-sampling machinery: the reverse KL is an exact expectation over
    # the student distribution, so no sampling-distribution correction applies.
    if cfg.use_importance_sampling_correction:
        raise ValueError(
            "opd_full is incompatible with use_importance_sampling_correction: the "
            "full-vocabulary reverse KL is an exact expectation and carries no "
            "sampled-token score-function term to correct."
        )
    if cfg.truncated_importance_sampling_type is not None:
        raise ValueError(
            "opd_full is incompatible with truncated_importance_sampling_type."
        )

    # Advantage/reward-shaped terms: opd_full never reads advantages or rewards.
    if cfg.positive_example_nll_weight != 0:
        raise ValueError(
            "opd_full is incompatible with a non-zero positive_example_nll_weight."
        )

    if cfg.use_kl_in_reward:
        raise ValueError(
            "opd_full is incompatible with use_kl_in_reward: it routes the "
            "reference KL through the reward, which reaches the objective only "
            "via advantages, and opd_full ignores advantages entirely. It also "
            "zeroes the loss-side reference_policy_kl_penalty, so the penalty "
            "would be dropped on both paths. Set use_kl_in_reward=false; the "
            "loss-side penalty is applied directly."
        )

    if cfg.reference_policy_kl_penalty != 0:
        warnings.warn(
            "opd_full is already a KL objective against the teacher; "
            "reference_policy_kl_penalty adds a second KL against the frozen "
            "reference policy on top of it.",
            stacklevel=3,
        )


class NLLLossFn(LossFunction):
    """Negative Log Likelihood Loss function."""

    loss_type = LossType.TOKEN_LEVEL
    input_type = LossInputType.LOGPROB

    def __init__(self, use_fused_linear_logprobs: bool = False):
        self.use_fused_linear_logprobs = use_fused_linear_logprobs
        # See MetricNormalizer — split-API trainers use this to undo the
        # placeholder global_valid_*=1 normalization per metric.
        self.metric_normalizations: dict[str, MetricNormalizer] = {
            "loss": MetricNormalizer.TOKENS,
            "num_unmasked_tokens": MetricNormalizer.NONE,
            "num_valid_samples": MetricNormalizer.NONE,
        }

    def __call__(
        self,
        next_token_logprobs: Tensor,
        data: BatchedDataDict[Any],
        global_valid_seqs: Tensor | None,
        global_valid_toks: Tensor,
        dpo_loss: bool = False,
        dpo_average_log_probs: bool = False,
    ) -> tuple[torch.Tensor, dict[str, Any]]:
        # logits shape: [batch_size, seq_len, vocab_size]
        # Get the next token logits for each position
        token_mask = data["token_mask"][:, 1:]
        sample_mask = data["sample_mask"]
        mask = token_mask * sample_mask.unsqueeze(-1)

        if dpo_loss:
            ## shape: [batch_size]
            num_unmasked_tokens = torch.sum(mask, -1)
            ## multiply by sample_mask to zero out invalid samples
            loss = -torch.sum(next_token_logprobs * mask, dim=-1)
            if dpo_average_log_probs:
                loss = loss / num_unmasked_tokens.clamp(min=1)
        else:
            ## single scalar loss
            ## scale by the total number of tokens in the batch
            loss = -masked_mean(
                next_token_logprobs,
                mask,
                global_normalization_factor=global_valid_toks,
            )

        return loss, {
            "loss": loss.item() if loss.ndim == 0 else loss,
            "num_unmasked_tokens": mask.sum().item(),
            "num_valid_samples": sample_mask.sum().item(),
        }


class PreferenceLossDataDict(TypedDict):
    """Keys for a preference loss.

    ``pair_index`` and ``is_chosen`` make pair reconstruction robust to row
    reordering by sequence packing. They remain optional for compatibility
    with existing interleaved DPO batches.
    """

    input_ids: torch.Tensor
    token_mask: torch.Tensor
    sample_mask: torch.Tensor
    pair_index: NotRequired[torch.Tensor]
    is_chosen: NotRequired[torch.Tensor]


class PreferenceLossFn(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:
        tuple[torch.Tensor, dict]: 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
    """

    loss_type = LossType.SEQUENCE_LEVEL
    input_type = LossInputType.LOGIT

    def split_output_tensor(
        self,
        tensor: Tensor,
        data: Optional[PreferenceLossDataDict] = None,
    ) -> tuple[Tensor, Tensor]:
        """Return pair-aligned chosen and rejected rows.

        Sequence packing may reorder rows, so positional slicing is only safe
        for legacy batches that do not carry pair metadata.
        """
        if data is None or "pair_index" not in data or "is_chosen" not in data:
            return tensor[::2], tensor[1::2]

        is_chosen = data["is_chosen"].to(device=tensor.device, dtype=torch.bool)
        pair_index = data["pair_index"].to(device=tensor.device)
        if (
            is_chosen.numel() != tensor.shape[0]
            or pair_index.numel() != tensor.shape[0]
        ):
            raise ValueError(
                "pair_index and is_chosen must contain one entry per preference row."
            )
        chosen_rows = torch.nonzero(is_chosen, as_tuple=True)[0]
        rejected_rows = torch.nonzero(~is_chosen, as_tuple=True)[0]
        chosen_pairs = pair_index.index_select(0, chosen_rows)
        rejected_pairs = pair_index.index_select(0, rejected_rows)
        chosen_order = chosen_pairs.argsort(stable=True)
        rejected_order = rejected_pairs.argsort(stable=True)
        chosen_pairs = chosen_pairs.index_select(0, chosen_order)
        rejected_pairs = rejected_pairs.index_select(0, rejected_order)
        if (
            not torch.equal(chosen_pairs, rejected_pairs)
            or torch.unique(chosen_pairs).numel() != chosen_pairs.numel()
        ):
            raise ValueError(
                "Preference batch must contain exactly one chosen and one rejected "
                "row for every pair_index."
            )
        return (
            tensor.index_select(0, chosen_rows.index_select(0, chosen_order)),
            tensor.index_select(0, rejected_rows.index_select(0, rejected_order)),
        )

    def _preference_loss(
        self,
        rewards: Tensor,
        sample_mask: Tensor,
        global_valid_seqs: Tensor,
        beta: float = 1.0,
        data: Optional[PreferenceLossDataDict] = None,
    ) -> tuple[Tensor, Tensor, Tensor, Tensor]:
        rewards_chosen, rewards_rejected = self.split_output_tensor(rewards, data)
        mask_chosen, mask_rejected = self.split_output_tensor(sample_mask, data)
        rewards_delta = rewards_chosen - rewards_rejected

        per_sample_loss = (
            -torch.nn.functional.logsigmoid(beta * rewards_delta) * mask_chosen
        )  ## zero out invalid samples

        ## divide by 2 because each preference example corresponds to 2 samples (chosen, rejected)
        return (
            masked_mean(
                per_sample_loss,
                mask_chosen,
                global_normalization_factor=global_valid_seqs / 2,
            ),
            masked_mean(
                rewards_chosen > rewards_rejected,
                mask_chosen,
                global_normalization_factor=global_valid_seqs / 2,
            ),
            masked_mean(
                rewards_chosen,
                mask_chosen,
                global_normalization_factor=global_valid_seqs / 2,
            ),
            masked_mean(
                rewards_rejected,
                mask_rejected,
                global_normalization_factor=global_valid_seqs / 2,
            ),
        )

    def __call__(
        self,
        logits: Tensor,
        data: BatchedDataDict[PreferenceLossDataDict],
        global_valid_seqs: Tensor,
        global_valid_toks: Tensor | None,
    ) -> tuple[torch.Tensor, dict[str, Any]]:
        sample_mask = data["sample_mask"]

        rewards = logits.squeeze(-1)

        (
            preference_loss,
            accuracy,
            rewards_chosen_mean,
            rewards_rejected_mean,
        ) = self._preference_loss(rewards, sample_mask, global_valid_seqs, data=data)

        ## divide by 2 because we're summing over (chosen, rejected) pairs
        num_valid_samples = sample_mask.sum() / 2

        return preference_loss, {
            "loss": preference_loss.item(),
            "accuracy": accuracy.item(),
            "rewards_chosen_mean": rewards_chosen_mean.item(),
            "rewards_rejected_mean": rewards_rejected_mean.item(),
            "num_valid_samples": num_valid_samples.item(),
        }


class DPOLossConfig(BaseModel, extra="allow"):
    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 DPOLossDataDict(TypedDict):
    """Required keys for the DPO loss function."""

    input_ids: torch.Tensor
    reference_policy_logprobs: torch.Tensor
    token_mask: torch.Tensor
    sample_mask: torch.Tensor
    pair_index: NotRequired[torch.Tensor]
    is_chosen: NotRequired[torch.Tensor]


class DPOLossFn(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:
    1. Preference Loss: Optimizes the model to prefer chosen responses over rejected ones
    2. 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.

    Args:
        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:
        tuple[torch.Tensor, dict]: 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
    """

    loss_type = LossType.SEQUENCE_LEVEL
    input_type = LossInputType.LOGPROB

    def __init__(self, cfg: DPOLossConfig, use_fused_linear_logprobs: bool = False):
        self.reference_policy_kl_penalty = cfg.reference_policy_kl_penalty
        self.preference_loss_weight = cfg.preference_loss_weight
        self.sft_loss_weight = cfg.sft_loss_weight
        self.preference_average_log_probs = cfg.preference_average_log_probs
        self.sft_average_log_probs = cfg.sft_average_log_probs
        self.use_fused_linear_logprobs = use_fused_linear_logprobs
        self.sft_loss = NLLLossFn(use_fused_linear_logprobs=use_fused_linear_logprobs)
        self.metric_normalizations = {
            "loss": MetricNormalizer.SEQUENCES,
            "sft_loss": MetricNormalizer.SEQUENCES,
            "preference_loss": MetricNormalizer.SEQUENCES,
            "accuracy": MetricNormalizer.SEQUENCES,
            "rewards_chosen_mean": MetricNormalizer.SEQUENCES,
            "rewards_rejected_mean": MetricNormalizer.SEQUENCES,
            "num_valid_samples": MetricNormalizer.NONE,
        }

    def _dpo_loss(
        self,
        next_token_logprobs: Tensor,
        data: BatchedDataDict[DPOLossDataDict],
        global_valid_seqs: Tensor,
    ) -> tuple[Tensor, Tensor, Tensor, Tensor]:
        ## TODO(@ashors): there's some duplicate code here with the NLLLossFn function. We should refactor
        token_mask = data["token_mask"][:, 1:]
        sample_mask = data["sample_mask"]

        ref_logprobs = data["reference_policy_logprobs"][:, :-1]
        diff = (next_token_logprobs - ref_logprobs) * token_mask

        rewards = diff.sum(-1)
        if self.preference_average_log_probs:
            rewards = rewards / token_mask.sum(-1).clamp(min=1)

        return self._preference_loss(
            rewards,
            sample_mask,
            global_valid_seqs,
            self.reference_policy_kl_penalty,
            data=data,
        )

    # TODO a cleaner typing fix would be required (probably that DPOLossFn should not inherit from PreferenceLossFn)
    def __call__(  # type: ignore
        self,
        next_token_logprobs: Tensor,
        data: BatchedDataDict[DPOLossDataDict],
        global_valid_seqs: Tensor,
        global_valid_toks: Tensor | None,
    ) -> tuple[torch.Tensor, dict[str, Any]]:
        sft_loss_chosen = torch.tensor(0.0)
        if self.sft_loss_weight > 0:
            assert global_valid_toks is not None, (
                "global_valid_toks must be provided for SFT loss"
            )
            sft_loss, _ = self.sft_loss(
                next_token_logprobs,
                data,
                global_valid_seqs=global_valid_seqs,
                global_valid_toks=global_valid_toks,  ## unused because sft loss returned is at the sample level
                dpo_loss=True,
                dpo_average_log_probs=self.sft_average_log_probs,
            )
            sft_loss_chosen, _ = self.split_output_tensor(sft_loss, data)
            sample_mask_chosen, _ = self.split_output_tensor(data["sample_mask"], data)
            sft_loss_chosen = masked_mean(
                sft_loss_chosen,
                sample_mask_chosen,
                global_normalization_factor=global_valid_seqs / 2,
            )

        (
            preference_loss,
            accuracy,
            rewards_chosen_mean,
            rewards_rejected_mean,
        ) = self._dpo_loss(next_token_logprobs, data, global_valid_seqs)

        dpo_loss = (
            self.sft_loss_weight * sft_loss_chosen
            + self.preference_loss_weight * preference_loss
        )

        ## divide by 2 because we're summing over (chosen, rejected) pairs
        num_valid_samples = data["sample_mask"].sum() / 2

        return dpo_loss, {
            "loss": dpo_loss.item(),
            "sft_loss": sft_loss_chosen.item(),
            "preference_loss": preference_loss.item(),
            "accuracy": accuracy.item(),
            "rewards_chosen_mean": rewards_chosen_mean.item(),
            "rewards_rejected_mean": rewards_rejected_mean.item(),
            "num_valid_samples": num_valid_samples.item(),
        }


class MPOLossConfig(BaseModel, extra="allow"):
    """Mixed Preference Optimization loss configuration."""

    reference_policy_kl_penalty: float = 0.05
    preference_loss_weight: float = 1.0
    sft_loss_weight: float = 0.0
    bco_loss_weight: float = 1.0
    preference_average_log_probs: bool = False
    sft_average_log_probs: bool = False
    quality_average_log_probs: bool = False
    reward_shift_momentum: float = 0.99
    reward_shift: float = 0.0


class MPOLossFn(PreferenceLossFn):
    """MPO loss with driver-owned BCO reward-shift state.

    The loss combines a DPO preference term, a chosen-response SFT term, and
    a BCO quality term. The BCO shift is intentionally not mutated here:
    loss objects are serialized to distributed workers, so worker-local
    mutation would be lost between optimizer steps and would diverge across
    data-parallel ranks. Instead, the loss returns a raw reward sum/count;
    the algorithm aggregates those values and updates ``reward_shift`` once
    on the driver after each successful training step.
    """

    loss_type = LossType.SEQUENCE_LEVEL
    input_type = LossInputType.LOGPROB

    def __init__(
        self, cfg: MPOLossConfig, use_fused_linear_logprobs: bool = False
    ) -> None:
        self.reference_policy_kl_penalty = cfg.reference_policy_kl_penalty
        self.preference_loss_weight = cfg.preference_loss_weight
        self.sft_loss_weight = cfg.sft_loss_weight
        self.bco_loss_weight = cfg.bco_loss_weight
        self.preference_average_log_probs = cfg.preference_average_log_probs
        self.sft_average_log_probs = cfg.sft_average_log_probs
        self.quality_average_log_probs = cfg.quality_average_log_probs
        self.reward_shift_momentum = cfg.reward_shift_momentum
        self.reward_shift = float(cfg.reward_shift)
        self.use_fused_linear_logprobs = use_fused_linear_logprobs

        if not 0.0 <= self.reward_shift_momentum < 1.0:
            raise ValueError("reward_shift_momentum must be in [0, 1)")

        self.metric_normalizations = {
            "loss": MetricNormalizer.SEQUENCES,
            "sft_loss": MetricNormalizer.SEQUENCES,
            "preference_loss": MetricNormalizer.SEQUENCES,
            "bco_loss": MetricNormalizer.SEQUENCES,
            "accuracy": MetricNormalizer.SEQUENCES,
            "rewards_chosen_mean": MetricNormalizer.SEQUENCES,
            "rewards_rejected_mean": MetricNormalizer.SEQUENCES,
            "bco_rewards_chosen_mean": MetricNormalizer.SEQUENCES,
            "bco_rewards_rejected_mean": MetricNormalizer.SEQUENCES,
            "num_valid_samples": MetricNormalizer.NONE,
            # These are sufficient statistics, aggregated by summation before
            # the single driver-side EMA update.
            "bco_reward_sum": MetricNormalizer.NONE,
            "bco_reward_count": MetricNormalizer.NONE,
        }

    def update_reward_shift(self, reward_sum: float, reward_count: float) -> float:
        """Apply one EMA update from globally aggregated BCO rewards."""
        if reward_count <= 0:
            return self.reward_shift
        batch_mean = reward_sum / reward_count
        self.reward_shift = (
            self.reward_shift_momentum * self.reward_shift
            + (1.0 - self.reward_shift_momentum) * batch_mean
        )
        return self.reward_shift

    @staticmethod
    def _sequence_logratios(
        next_token_logprobs: Tensor,
        data: BatchedDataDict[DPOLossDataDict],
        *,
        average: bool,
    ) -> Tensor:
        token_mask = data["token_mask"][:, 1:]
        reference_logprobs = data["reference_policy_logprobs"][:, :-1]
        logratios = ((next_token_logprobs - reference_logprobs) * token_mask).sum(
            dim=-1
        )
        if average:
            logratios = logratios / token_mask.sum(dim=-1).clamp(min=1)
        return logratios

    def _sft_loss(
        self,
        next_token_logprobs: Tensor,
        data: BatchedDataDict[DPOLossDataDict],
        global_valid_seqs: Tensor,
    ) -> Tensor:
        token_mask = data["token_mask"][:, 1:]
        sample_mask = data["sample_mask"]
        per_sequence_loss = -(next_token_logprobs * token_mask).sum(dim=-1)
        if self.sft_average_log_probs:
            per_sequence_loss = per_sequence_loss / token_mask.sum(dim=-1).clamp(min=1)
        chosen_loss, _ = self.split_output_tensor(per_sequence_loss, data)
        chosen_mask, _ = self.split_output_tensor(sample_mask, data)
        return masked_mean(
            chosen_loss,
            chosen_mask,
            global_normalization_factor=global_valid_seqs / 2,
        )

    def _bco_loss(
        self,
        rewards: Tensor,
        data: BatchedDataDict[DPOLossDataDict],
        global_valid_seqs: Tensor,
    ) -> tuple[Tensor, Tensor, Tensor]:
        sample_mask = data["sample_mask"]
        chosen_rewards, rejected_rewards = self.split_output_tensor(rewards, data)
        chosen_mask, rejected_mask = self.split_output_tensor(sample_mask, data)
        shift = torch.as_tensor(
            self.reward_shift, dtype=rewards.dtype, device=rewards.device
        )
        chosen_loss = (
            -torch.nn.functional.logsigmoid(chosen_rewards - shift) * chosen_mask
        )
        rejected_loss = (
            -torch.nn.functional.logsigmoid(-(rejected_rewards - shift)) * rejected_mask
        )
        pair_count = global_valid_seqs / 2
        loss = masked_mean(
            chosen_loss,
            chosen_mask,
            global_normalization_factor=pair_count,
        ) + masked_mean(
            rejected_loss,
            rejected_mask,
            global_normalization_factor=pair_count,
        )
        chosen_mean = masked_mean(
            chosen_rewards,
            chosen_mask,
            global_normalization_factor=pair_count,
        )
        rejected_mean = masked_mean(
            rejected_rewards,
            rejected_mask,
            global_normalization_factor=pair_count,
        )
        return loss, chosen_mean, rejected_mean

    def __call__(
        self,
        next_token_logprobs: Tensor,
        data: BatchedDataDict[DPOLossDataDict],
        global_valid_seqs: Tensor,
        global_valid_toks: Tensor | None,
    ) -> tuple[Tensor, dict[str, Any]]:
        preference_logratios = self._sequence_logratios(
            next_token_logprobs,
            data,
            average=self.preference_average_log_probs,
        )
        (
            preference_loss,
            accuracy,
            rewards_chosen_mean,
            rewards_rejected_mean,
        ) = self._preference_loss(
            preference_logratios,
            data["sample_mask"],
            global_valid_seqs,
            beta=self.reference_policy_kl_penalty,
            data=data,
        )

        quality_logratios = self._sequence_logratios(
            next_token_logprobs,
            data,
            average=self.quality_average_log_probs,
        )
        quality_rewards = self.reference_policy_kl_penalty * quality_logratios
        bco_loss, bco_chosen_mean, bco_rejected_mean = self._bco_loss(
            quality_rewards, data, global_valid_seqs
        )

        sft_loss = torch.zeros(
            (), dtype=next_token_logprobs.dtype, device=next_token_logprobs.device
        )
        if self.sft_loss_weight > 0:
            sft_loss = self._sft_loss(next_token_logprobs, data, global_valid_seqs)

        total_loss = (
            self.preference_loss_weight * preference_loss
            + self.sft_loss_weight * sft_loss
            + self.bco_loss_weight * bco_loss
        )
        sample_mask = data["sample_mask"].to(dtype=quality_rewards.dtype)
        bco_reward_sum = (quality_rewards.detach() * sample_mask).sum()
        bco_reward_count = sample_mask.sum()

        return total_loss, {
            "loss": total_loss.item(),
            "sft_loss": sft_loss.item(),
            "preference_loss": preference_loss.item(),
            "bco_loss": bco_loss.item(),
            "accuracy": accuracy.item(),
            "rewards_chosen_mean": rewards_chosen_mean.item(),
            "rewards_rejected_mean": rewards_rejected_mean.item(),
            "bco_rewards_chosen_mean": bco_chosen_mean.item(),
            "bco_rewards_rejected_mean": bco_rejected_mean.item(),
            "num_valid_samples": (data["sample_mask"].sum() / 2).item(),
            "bco_reward_sum": bco_reward_sum.item(),
            "bco_reward_count": bco_reward_count.item(),
        }


class DistillationLossConfig(BaseModel, extra="allow"):
    kl_type: str = "mixed"
    mixed_kl_weight: float = 0.5
    zero_outside_topk: bool = False


class DistillationLossDataDict(TypedDict):
    input_ids: torch.Tensor
    input_lengths: torch.Tensor
    token_mask: torch.Tensor
    sample_mask: torch.Tensor
    teacher_topk_logits: torch.Tensor
    teacher_topk_indices: torch.Tensor


class DistillationLossFn(LossFunction):
    """Distillation loss function."""

    loss_type = LossType.TOKEN_LEVEL
    input_type = LossInputType.DISTILLATION

    def __init__(self, cfg: DistillationLossConfig):
        self.kl_type = cfg.kl_type
        self.mixed_kl_weight = cfg.mixed_kl_weight
        self.zero_outside_topk = cfg.zero_outside_topk
        self.log_infinitesimal = -100

        assert self.kl_type in ["forward", "reverse", "mixed"], "Invalid KL type"
        assert self.mixed_kl_weight >= 0 and self.mixed_kl_weight <= 1, (
            "Invalid mixed KL weight"
        )

    def __call__(
        self,
        student_topk_logprobs: torch.Tensor,
        teacher_topk_logprobs: torch.Tensor,
        H_all: torch.Tensor | None,
        data: 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."""
        student_probs = student_topk_logprobs.exp()  # [B, S-1, k]
        teacher_probs = teacher_topk_logprobs.exp()  # [B, S-1, k]

        loss_correction_term = torch.zeros_like(student_probs[..., 0])  # [B, S-1]
        if self.zero_outside_topk and self.kl_type != "forward":
            H_rest = H_all - (student_probs * student_topk_logprobs).sum(-1)
            P_rest = 1 - (student_probs.sum(-1))
            # The entropy and prob of the rest of the tokens [B, S-1]
            loss_correction_term = H_rest - self.log_infinitesimal * P_rest  # [B, S-1]
            if self.kl_type == "mixed":
                loss_correction_term = loss_correction_term * (
                    1.0 - self.mixed_kl_weight
                )

        if self.kl_type == "forward":
            per_token_kl = teacher_probs * (
                teacher_topk_logprobs - student_topk_logprobs
            )
        elif self.kl_type == "reverse":
            per_token_kl = student_probs * (
                student_topk_logprobs - teacher_topk_logprobs
            )
        else:
            # mixed KL
            kl_forward = teacher_probs * (teacher_topk_logprobs - student_topk_logprobs)
            kl_reverse = student_probs * (student_topk_logprobs - teacher_topk_logprobs)
            per_token_kl = (
                self.mixed_kl_weight * kl_forward
                + (1.0 - self.mixed_kl_weight) * kl_reverse
            )

        per_token_kl = per_token_kl.sum(dim=-1) + loss_correction_term  # [B, S-1]

        # Masking and reduction
        if "token_mask" in data and "sample_mask" in data:
            token_mask = data["token_mask"][:, 1:]
            sample_mask = data["sample_mask"]
            # Align mask length to current per_token_kl
            max_len = per_token_kl.shape[1]
            token_mask = token_mask[:, :max_len]
            mask = token_mask * sample_mask.unsqueeze(-1)  # [B, S-1]
            # align mask shape to per_token_kl
            kl_loss = masked_mean(
                per_token_kl,
                mask,
                global_normalization_factor=global_valid_toks,
            )
        else:
            kl_loss = per_token_kl.mean()

        metrics = {
            "loss": float(kl_loss.item()) if kl_loss.ndim == 0 else kl_loss,
            "num_valid_samples": data["input_ids"].shape[0],
        }

        return kl_loss, metrics


class MseValueLossConfig(BaseModel, extra="forbid"):
    """Config for the MSE value loss used by PPO's value model."""

    # Scaling factor applied to the value loss before it is added to the policy loss.
    scale: float = 1.0
    # Clipping range for value predictions (PPO-style). Set to None to disable clipping.
    cliprange: Optional[float] = None


class MseValueLossFn(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.
    """

    input_type = LossInputType.LOGIT

    def __init__(self, cfg: MseValueLossConfig):
        self.scale = cfg.scale
        self.cliprange = cfg.cliprange
        self.loss_type = LossType.TOKEN_LEVEL

    def __call__(
        self,
        logits: torch.Tensor,
        data: 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."""
        # Squeeze trailing singleton from value head output: [B, S, 1] -> [B, S]
        if logits.ndim > 2 and logits.shape[-1] == 1:
            logits = logits.squeeze(-1)
        values = logits

        token_mask = data["token_mask"]
        sample_mask = data["sample_mask"]
        returns = data["returns"]
        mask = token_mask * sample_mask.unsqueeze(-1)

        if self.cliprange and self.cliprange > 0:
            old_values = data["values"]
            vpred_clipped = torch.clamp(
                values,
                old_values - self.cliprange,
                old_values + self.cliprange,
            )
            vf_losses_unclipped = (values - returns) ** 2
            vf_losses_clipped = (vpred_clipped - returns) ** 2
            vf_losses = torch.max(vf_losses_unclipped, vf_losses_clipped)
            loss = (
                0.5
                * self.scale
                * masked_mean(
                    vf_losses,
                    mask,
                    global_normalization_factor=global_valid_toks,
                )
            )
            vf_clipfrac = masked_mean(
                (vf_losses_clipped > vf_losses_unclipped).float(),
                mask,
                global_normalization_factor=global_valid_toks,
            ).item()
        else:
            loss = torch.nn.functional.mse_loss(values, returns, reduction="none")
            loss = (
                0.5
                * self.scale
                * masked_mean(
                    loss,
                    mask,
                    global_normalization_factor=global_valid_toks,
                )
            )
            vf_clipfrac = 0.0

        with torch.no_grad():
            # Use global_valid_toks so each MB contributes local_sum/global_total.
            # Summing across MBs in ppo.py then gives the correct global mean.
            returns_mean = masked_mean(
                returns,
                mask,
                global_normalization_factor=global_valid_toks,
            ).item()
            values_mean = masked_mean(
                values,
                mask,
                global_normalization_factor=global_valid_toks,
            ).item()

            # Min/max are per-MB; ppo.py takes min/max across MBs.
            masked_values = values[mask.bool()]
            values_min = (
                masked_values.min().item() if masked_values.numel() > 0 else 0.0
            )
            values_max = (
                masked_values.max().item() if masked_values.numel() > 0 else 0.0
            )

            # Explained variance sufficient statistics.
            # EV = 1 - Var(residual) / Var(returns)
            # We export E[r²] and E[(r-v)²] (both / global_valid_toks so they
            # sum correctly across MBs).  ppo.py combines them with returns_mean
            # and values_mean to compute exact global EV.
            returns_sq_mean = masked_mean(
                returns**2,
                mask,
                global_normalization_factor=global_valid_toks,
            ).item()
            residual_sq_mean = masked_mean(
                (returns - values) ** 2,
                mask,
                global_normalization_factor=global_valid_toks,
            ).item()

        metrics = {
            "loss": float(loss.item()),
            "vf_clipfrac": vf_clipfrac,
            "returns_mean": returns_mean,
            "values_mean": values_mean,
            "values_min": values_min,
            "values_max": values_max,
            "returns_sq_mean": returns_sq_mean,
            "residual_sq_mean": residual_sq_mean,
            "num_valid_samples": int(values.shape[0]),
        }

        return loss, metrics


# =====================================================================
# Cross-tokenizer distillation
# =====================================================================


class CrossTokenizerDistillationLossConfig(TypedDict):
    """Config for cross-tokenizer distillation loss.

    Attributes:
        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.
        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).
        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``.
        temperature: Softmax temperature applied symmetrically to student
            and teacher logits before KL.
        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``.
        uncommon_topk: Cap on the L1 uncommon-tail sort in the gold path.
            Defaults to 8192. Inert when ``gold_loss=False``.
        reverse_kl: If True, compute KL(student || teacher) instead of
            KL(teacher || student).
        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.
        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.
        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.
        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.
        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.
        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.
    """

    gold_loss: bool
    xtoken_loss: bool
    temperature: float
    vocab_topk: int
    uncommon_topk: int
    reverse_kl: bool
    exact_token_match_only: bool
    kl_loss_weight: float
    ce_loss_scale: float
    dynamic_loss_scaling: bool
    # Multi-teacher aggregation (user loss_fn knobs). gold_loss/xtoken_loss above
    # are the global defaults; teachers[i].gold_loss / .xtoken_loss can override
    # them per teacher in kd_loss_mode="sum".
    kd_loss_mode: str  # "sum" | "averaged_logits" | "select_teacher"
    normalize_teacher_by_vocab: (
        bool  # sum-mode only: scale each teacher's KD by log(V_t_i)/log(min_j V_t_j)
    )
    alpha: float  # softmax temperature on dynamic teacher-weight scores (sum_weights_metric)
    sum_weights_metric: NotRequired[
        Optional[str]
    ]  # "ce" | "entropy" | "max_prob"; None => static teacher_weights. sum-mode only.
    # Runtime-injected by xtoken_off_policy_distillation.setup (parallel
    # per-teacher lists + the student vocab size); not user loss_fn keys.
    student_vocab_size: NotRequired[int]
    teacher_vocab_sizes: NotRequired[list[int]]
    projection_matrix_paths: NotRequired[list[Optional[str]]]
    teacher_weights: NotRequired[list[float]]
    teacher_gold_loss: NotRequired[list[Optional[bool]]]
    teacher_xtoken_loss: NotRequired[list[Optional[bool]]]


class CrossTokenizerDistillationLossDataDict(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).
    """

    input_ids: torch.Tensor
    input_lengths: torch.Tensor
    token_mask: torch.Tensor
    sample_mask: torch.Tensor


class CrossTokenizerDistillationLossFn(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, metrics)``. Aggregate metrics: ``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``.
    """

    loss_type = LossType.TOKEN_LEVEL
    input_type = LossInputType.DISTILLATION_CROSS_TOKENIZER

    def __init__(self, cfg: CrossTokenizerDistillationLossConfig):
        if cfg["xtoken_loss"] and not cfg["gold_loss"]:
            raise ValueError(
                "xtoken_loss=True requires gold_loss=True; xtoken_loss is "
                "a modifier inside the gold path (relaxes the exact-map "
                "threshold and adds collision resolution) and is undefined "
                "in the P-KL path."
            )
        # Dynamic teacher weighting (sum_weights_metric) and normalize_teacher_by_vocab
        # are only applied in kd_loss_mode="sum"; reject the combo instead of
        # silently ignoring them under the other modes.
        if cfg.get("sum_weights_metric") is not None and cfg["kd_loss_mode"] != "sum":
            raise ValueError(
                f"sum_weights_metric={cfg['sum_weights_metric']!r} is only applied "
                f"in kd_loss_mode='sum'; it is ignored by '{cfg['kd_loss_mode']}'. "
                "Unset one of them."
            )
        if cfg.get("normalize_teacher_by_vocab") and cfg["kd_loss_mode"] != "sum":
            raise ValueError(
                "normalize_teacher_by_vocab is only applied in kd_loss_mode='sum'; "
                f"it is ignored by '{cfg['kd_loss_mode']}'. Unset one of them."
            )
        # averaged_logits forms a convex combination of teacher logits
        # (weight_i / sum(weights)); a zero weight-sum (all zeros, or a signed
        # set cancelling to 0) makes that division undefined. Reject it here
        # rather than fail with a deep ZeroDivisionError mid-step.
        _weights = cfg.get("teacher_weights")
        if (
            cfg["kd_loss_mode"] == "averaged_logits"
            and _weights is not None
            and sum(_weights) == 0
        ):
            raise ValueError(
                "teacher_weights must not sum to zero in "
                "kd_loss_mode='averaged_logits' (they form a convex combination "
                f"of teacher logits); got teacher_weights={list(_weights)}."
            )
        # Global loss knobs (shared across all teachers).
        self.gold_loss = cfg["gold_loss"]
        self.xtoken_loss = cfg["xtoken_loss"]
        self.temperature = cfg["temperature"]
        self.vocab_topk = cfg["vocab_topk"]
        self.uncommon_topk = cfg["uncommon_topk"]
        self.reverse_kl = cfg["reverse_kl"]
        self.exact_token_match_only = cfg["exact_token_match_only"]
        self.kl_loss_weight = cfg["kl_loss_weight"]
        self.ce_loss_scale = cfg["ce_loss_scale"]
        self.dynamic_loss_scaling = cfg["dynamic_loss_scaling"]
        self.student_vocab_size = cfg["student_vocab_size"]
        # Multi-teacher aggregation knobs.
        self.kd_loss_mode = cfg["kd_loss_mode"]
        self.normalize_teacher_by_vocab = cfg["normalize_teacher_by_vocab"]
        self.alpha = cfg["alpha"]
        # sum_weights_metric is NotRequired -> None means static teacher_weights.
        self.sum_weights_metric = cfg.get("sum_weights_metric")
        # 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.
        self.projection_matrix_paths = list(cfg["projection_matrix_paths"])
        self.teacher_vocab_sizes = list(cfg["teacher_vocab_sizes"])
        self.teacher_weights = list(cfg["teacher_weights"])
        self.teacher_gold_loss = list(cfg["teacher_gold_loss"])
        self.teacher_xtoken_loss = list(cfg["teacher_xtoken_loss"])
        # Every per-teacher list must have the same length (one entry per
        # teacher); a mismatch would otherwise surface as a deep IndexError
        # mid-training instead of a clear error here.
        per_teacher_lens = {
            "projection_matrix_paths": len(self.projection_matrix_paths),
            "teacher_vocab_sizes": len(self.teacher_vocab_sizes),
            "teacher_weights": len(self.teacher_weights),
            "teacher_gold_loss": len(self.teacher_gold_loss),
            "teacher_xtoken_loss": len(self.teacher_xtoken_loss),
        }
        if len(set(per_teacher_lens.values())) != 1:
            raise ValueError(
                f"per-teacher lists must be equal length, got {per_teacher_lens}"
            )
        self.num_teachers = len(self.projection_matrix_paths)
        # The materialized projection matrix and the derived exact-map
        # partition both live in process-local caches in
        # ``x_token.loss_utils`` (see ``get_sparse_projection_matrix``,
        # ``get_topk_projection``, ``build_exact_token_map``), not on
        # this instance. That keeps the driver-side ``loss_fn`` free of
        # any large CUDA tensors and lets multiple loss instances on
        # the same worker share one load.

    def _teacher_is_same_vocab(self, i: int) -> bool:
        """A teacher is same-vocab (direct KL, no projection) iff its path is None."""
        return self.projection_matrix_paths[i] is None

    def __call__(
        self,
        data: BatchedDataDict[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, LocalizedAlignment],
        *,
        student_next_token_logprobs: Optional[torch.Tensor] = None,
        student_next_token_mask: Optional[torch.Tensor] = None,
        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 Automodel CP path also supplies its sequence-local CE inputs.
        """
        ce_loss = self._compute_ce(
            logits,
            data,
            global_valid_toks,
            student_next_token_logprobs=student_next_token_logprobs,
            student_next_token_mask=student_next_token_mask,
            cp_group=cp_group,
        )

        if self.kd_loss_mode == "sum":
            total_kd, per_teacher_metrics = self._sum_kd(
                student_logits_contig,
                data,
                teacher_full_logits_by_idx,
                aligns_by_idx,
                global_valid_toks,
                tp_group=tp_group,
                cp_group=cp_group,
            )
        elif self.kd_loss_mode == "averaged_logits":
            total_kd, per_teacher_metrics = self._averaged_logits_kd(
                student_logits_contig,
                data,
                teacher_full_logits_by_idx,
                aligns_by_idx,
                global_valid_toks,
                tp_group=tp_group,
                cp_group=cp_group,
            )
        elif self.kd_loss_mode == "select_teacher":
            total_kd, per_teacher_metrics = self._select_teacher_kd(
                student_logits_contig,
                data,
                teacher_full_logits_by_idx,
                aligns_by_idx,
                global_valid_toks,
                tp_group=tp_group,
                cp_group=cp_group,
            )
        else:
            raise ValueError(f"Unknown kd_loss_mode: {self.kd_loss_mode!r}")

        # Combine the aggregated KD term with the single student CE term.
        if self.dynamic_loss_scaling:
            # loss = sg(ce/kd) * kd + ce; user kl_loss_weight / ce_loss_scale
            # are intentionally ignored in this branch.
            kd_detached = total_kd.detach().abs()
            ce_detached = ce_loss.detach().abs()
            kl_scale = torch.where(
                kd_detached > 0,
                ce_detached / kd_detached,
                torch.ones_like(kd_detached),
            )
            loss = kl_scale * total_kd + ce_loss
        else:
            kl_scale = torch.tensor(1.0, device=total_kd.device, dtype=total_kd.dtype)
            loss = self.kl_loss_weight * total_kd + self.ce_loss_scale * ce_loss

        # Next-token accuracy on the student side (quick per-step signal), masked
        # to valid tokens. Computed once on the student from the shared CP-relaid
        # fields (carried on every teacher's align); the CP-aware shift pairs
        # predictors with the right labels under load-balanced sharding.
        align0 = aligns_by_idx[0]
        accuracy = next_token_accuracy(
            student_logits_contig,
            input_ids=align0.student_input_ids,
            token_mask=align0.student_token_mask,
            sample_mask=data["sample_mask"],
            tp_group=tp_group,
            cp_group=cp_group,
        )

        metrics: dict[str, Any] = {
            "loss": loss.item(),
            # Aggregate KD term (kept under ``kl_loss`` so existing trainer
            # metric handling continues to work); per-teacher terms are suffixed
            # ``_t{i}``.
            "kl_loss": total_kd.item(),
            "ce_loss": ce_loss.item(),
            "kl_loss_scale": kl_scale.item(),
            "accuracy": accuracy.item(),
            "num_valid_samples": data["input_ids"].shape[0],
        }
        metrics.update(per_teacher_metrics)
        return loss, metrics

    # ------------------------------------------------------------------ #
    # Multi-teacher aggregation
    # ------------------------------------------------------------------ #
    def _resolve_gold_xtoken(self, 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.
        """
        if not use_per_teacher:
            return self.gold_loss, self.xtoken_loss
        gold = (
            self.teacher_gold_loss[i]
            if self.teacher_gold_loss[i] is not None
            else self.gold_loss
        )
        xtoken = (
            self.teacher_xtoken_loss[i]
            if self.teacher_xtoken_loss[i] is not None
            else self.xtoken_loss
        )
        return gold, xtoken

    def _compute_teacher_kd(
        self,
        i: int,
        student_logits_contig: torch.Tensor,
        data: BatchedDataDict[CrossTokenizerDistillationLossDataDict],
        teacher_full_logits_by_idx: dict[int, torch.Tensor],
        aligns_by_idx: dict[int, 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.
        """
        if self._teacher_is_same_vocab(i):
            return self._compute_same_vocab_kl(
                i,
                student_logits_contig,
                teacher_full_logits_by_idx[i],
                aligns_by_idx[i],
                global_valid_toks,
                tp_group=tp_group,
                cp_group=cp_group,
            )

        teacher_logits = teacher_full_logits_by_idx[i]
        align = aligns_by_idx[i]
        proj_path = self.projection_matrix_paths[i]
        v_t = self.teacher_vocab_sizes[i]
        gold, xtoken = self._resolve_gold_xtoken(i, use_per_teacher_flags)
        if xtoken and not gold:
            raise ValueError(f"teacher {i}: xtoken_loss=True requires gold_loss=True.")

        if gold:
            kd, kl_common, l1_uncommon, num_valid_chunks, top1 = self._compute_gold(
                student_logits_contig,
                teacher_logits,
                align,
                projection_matrix_path=proj_path,
                teacher_vocab_size=v_t,
                xtoken_loss=xtoken,
                tp_group=tp_group,
                cp_group=cp_group,
            )
            return kd, {
                "kl_loss": kd.item(),
                "kl_common": kl_common.item(),
                "l1_uncommon": l1_uncommon.item(),
                "proj_accuracy": top1.item(),
                "num_valid_chunks": int(num_valid_chunks.item()),
            }

        kl, num_valid_pairs, proj_acc = self._compute_p_kl(
            student_logits_contig,
            teacher_logits,
            align,
            projection_matrix_path=proj_path,
            teacher_vocab_size=v_t,
            tp_group=tp_group,
            cp_group=cp_group,
        )
        return kl, {
            "kl_loss": kl.item(),
            "proj_accuracy": proj_acc.item(),
            "num_valid_pairs": int(num_valid_pairs.item()),
        }

    def _compute_same_vocab_kl(
        self,
        i: int,
        student_logits_contig: torch.Tensor,
        teacher_full_logits: torch.Tensor,
        align: 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``.
        """
        kd = self._direct_topk_kl(
            student_logits_contig,
            teacher_full_logits,
            align,
            global_valid_toks,
            tp_group=tp_group,
            cp_group=cp_group,
        )
        return kd, {"kl_loss": kd.item()}

    def _direct_topk_kl(
        self,
        student_logits: torch.Tensor,
        teacher_full_logits: torch.Tensor,
        align: 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's K columns are
        gathered TP-aware (``vocab_parallel_gather_columns``) and softmaxed
        within the subset directly: because both sides renormalize within K,
        the full-vocab partition function cancels, so
        ``log_softmax_K(logits[..., idx] / T)`` equals the previous
        full-vocab-log-softmax-then-renorm form exactly — without materializing
        the all-gathered ``[B, T, V]`` student log-probs. 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.
        """
        T = self.temperature
        # Drop HF lm_head padding beyond the shared tokenizer vocab.
        v_s = self.student_vocab_size
        teacher = teacher_full_logits
        if teacher.shape[-1] > v_s:
            teacher = teacher[..., :v_s]
        vocab_topk = min(self.vocab_topk, teacher.shape[-1])
        topk_idx = select_teacher_topk_indices(teacher, vocab_topk, cp_group=cp_group)

        # [B, T, K] instead of the all-gathered [B, T, V]: the K-subset
        # renormalization cancels the full-vocab partition function, so the
        # subset log_softmax is exactly equivalent (see docstring).
        student_topk_logits = vocab_parallel_gather_columns(
            student_logits, topk_idx, tp_group=tp_group
        )
        student_log_probs_k = torch.log_softmax(student_topk_logits / T, dim=-1)
        teacher_log_probs_k = torch.log_softmax(
            teacher[..., topk_idx].float() / T, dim=-1
        )
        if self.reverse_kl:
            per_pos = torch.nn.functional.kl_div(
                teacher_log_probs_k,
                student_log_probs_k,
                reduction="none",
                log_target=True,
            ).sum(dim=-1)
        else:
            per_pos = torch.nn.functional.kl_div(
                student_log_probs_k,
                teacher_log_probs_k,
                reduction="none",
                log_target=True,
            ).sum(dim=-1)
        return self._same_vocab_masked_kl(per_pos, align, global_valid_toks, cp_group)

    def _direct_full_vocab_kl(
        self,
        student_logits: torch.Tensor,
        teacher_logits: torch.Tensor,
        align: 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.
        """
        T = self.temperature
        student_log_probs = vocab_parallel_full_log_softmax(
            student_logits, T, tp_group=tp_group
        )
        v_s = student_log_probs.shape[-1]
        teacher = teacher_logits.float()
        if teacher.shape[-1] > v_s:
            teacher = teacher[..., :v_s]
        teacher_log_probs = torch.log_softmax(teacher / T, dim=-1)
        if self.reverse_kl:
            per_pos = torch.nn.functional.kl_div(
                teacher_log_probs, student_log_probs, reduction="none", log_target=True
            ).sum(dim=-1)
        else:
            per_pos = torch.nn.functional.kl_div(
                student_log_probs, teacher_log_probs, reduction="none", log_target=True
            ).sum(dim=-1)
        return self._same_vocab_masked_kl(per_pos, align, global_valid_toks, cp_group)

    def _same_vocab_masked_kl(
        self,
        per_pos: torch.Tensor,
        align: 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``.
        """
        T = self.temperature
        next_mask = cp_shift_next(
            to_local_if_dtensor(align.student_token_mask), cp_group, fill=0
        )
        sample_mask = to_local_if_dtensor(align.sample_mask)
        mask = next_mask.float() * sample_mask.unsqueeze(-1).float()
        return (
            masked_mean(per_pos, mask, global_normalization_factor=global_valid_toks)
            * T
            * T
        )

    def _sum_kd(
        self,
        student_logits_contig: torch.Tensor,
        data: BatchedDataDict[CrossTokenizerDistillationLossDataDict],
        teacher_full_logits_by_idx: dict[int, torch.Tensor],
        aligns_by_idx: dict[int, 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)``.
        """
        device = student_logits_contig.device
        if self.sum_weights_metric is not None:
            weights = self._compute_dynamic_weights(
                data, teacher_full_logits_by_idx, aligns_by_idx
            )
        else:
            weights = [
                torch.tensor(
                    self.teacher_weights[i],
                    device=device,
                    dtype=student_logits_contig.dtype,
                )
                for i in range(self.num_teachers)
            ]

        if self.normalize_teacher_by_vocab:
            temp_weight = torch.log(
                torch.tensor(float(min(self.teacher_vocab_sizes)), device=device)
            )

        total_kd: Optional[torch.Tensor] = None
        per_metrics: dict[str, Any] = {}
        # Deterministic teacher order: each teacher's KD fires its own
        # collectives, so the order must match across ranks.
        for i in range(self.num_teachers):
            kd_i, m_i = self._compute_teacher_kd(
                i,
                student_logits_contig,
                data,
                teacher_full_logits_by_idx,
                aligns_by_idx,
                global_valid_toks,
                use_per_teacher_flags=True,
                tp_group=tp_group,
                cp_group=cp_group,
            )
            weighted = kd_i * weights[i]
            if self.normalize_teacher_by_vocab:
                v_scale = (
                    torch.log(
                        torch.tensor(float(self.teacher_vocab_sizes[i]), device=device)
                    )
                    / temp_weight
                )
                weighted = weighted * v_scale
            total_kd = weighted if total_kd is None else total_kd + weighted
            for k, v in m_i.items():
                per_metrics[f"{k}_t{i}"] = v
            per_metrics[f"weight_t{i}"] = float(weights[i].item())
        assert total_kd is not None
        return total_kd, per_metrics

    def _averaged_logits_kd(
        self,
        student_logits_contig: torch.Tensor,
        data: BatchedDataDict[CrossTokenizerDistillationLossDataDict],
        teacher_full_logits_by_idx: dict[int, torch.Tensor],
        aligns_by_idx: dict[int, 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``).
        """
        full = [teacher_full_logits_by_idx.get(i) for i in range(self.num_teachers)]
        # Direct per-position KL is only valid when every teacher shares the
        # student's tokenizer (no projection matrix) *and* ships full logits of
        # identical shape. Two cross-tokenizer teachers can have matching shapes
        # yet still need the projection/alignment path, so the shape check alone
        # is insufficient.
        same_tokenizer = all(p is None for p in self.projection_matrix_paths)
        same_shape = all(f is not None for f in full) and (
            len({tuple(f.shape) for f in full if f is not None}) == 1
        )
        if not (same_tokenizer and same_shape):
            total_kd: Optional[torch.Tensor] = None
            per_metrics: dict[str, Any] = {}
            for i in range(self.num_teachers):
                kd_i, m_i = self._compute_teacher_kd(
                    i,
                    student_logits_contig,
                    data,
                    teacher_full_logits_by_idx,
                    aligns_by_idx,
                    global_valid_toks,
                    use_per_teacher_flags=False,
                    tp_group=tp_group,
                    cp_group=cp_group,
                )
                w = self.teacher_weights[i]
                weighted = kd_i * w
                total_kd = weighted if total_kd is None else total_kd + weighted
                for k, v in m_i.items():
                    per_metrics[f"{k}_t{i}"] = v
                per_metrics[f"weight_t{i}"] = float(w)
            assert total_kd is not None
            return total_kd, per_metrics

        total_w = sum(self.teacher_weights)
        avg: Optional[torch.Tensor] = None
        for i, f in enumerate(full):
            assert f is not None
            contrib = f.float() * (self.teacher_weights[i] / total_w)
            avg = contrib if avg is None else avg + contrib
        assert avg is not None
        kd = self._direct_full_vocab_kl(
            student_logits_contig,
            avg,
            aligns_by_idx[0],
            global_valid_toks,
            tp_group=tp_group,
            cp_group=cp_group,
        )
        return kd, {"kl_loss": kd.item()}

    def _dp_global_masked_mean(
        self, 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).
        """
        num = group_all_reduce_sum(
            (values * mask).sum(), group=torch.distributed.group.WORLD
        )
        den = group_all_reduce_sum(
            mask.sum(), group=torch.distributed.group.WORLD
        ).clamp(min=1.0)
        return num / den

    def _select_teacher_kd(
        self,
        student_logits_contig: torch.Tensor,
        data: BatchedDataDict[CrossTokenizerDistillationLossDataDict],
        teacher_full_logits_by_idx: dict[int, torch.Tensor],
        aligns_by_idx: dict[int, 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."""
        with torch.no_grad():
            ces: list[float] = []
            for i in range(self.num_teachers):
                t_logits, t_ids, t_mask = self._teacher_score_inputs(
                    i, data, teacher_full_logits_by_idx, aligns_by_idx
                )
                ce_pos = torch.nn.functional.cross_entropy(
                    t_logits[:, :-1].reshape(-1, t_logits.shape[-1]).float(),
                    t_ids[:, 1:].reshape(-1),
                    reduction="none",
                )
                mask = (
                    t_mask[:, 1:].float() * data["sample_mask"].unsqueeze(-1).float()
                ).reshape(-1)
                ces.append(self._dp_global_masked_mean(ce_pos, mask).item())
            best = int(min(range(self.num_teachers), key=lambda j: ces[j]))

        kd, m = self._compute_teacher_kd(
            best,
            student_logits_contig,
            data,
            teacher_full_logits_by_idx,
            aligns_by_idx,
            global_valid_toks,
            use_per_teacher_flags=False,
            tp_group=tp_group,
            cp_group=cp_group,
        )
        per_metrics: dict[str, Any] = {f"{k}_t{best}": v for k, v in m.items()}
        per_metrics["selected_teacher"] = best
        return kd, per_metrics

    def _teacher_score_inputs(
        self,
        i: int,
        data: BatchedDataDict[CrossTokenizerDistillationLossDataDict],
        teacher_full_logits_by_idx: dict[int, torch.Tensor],
        aligns_by_idx: dict[int, 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.
        """
        if self._teacher_is_same_vocab(i):
            align = aligns_by_idx[i]
            ids = to_local_if_dtensor(align.student_input_ids)
            token_mask = to_local_if_dtensor(align.student_token_mask)
        else:
            ids = to_local_if_dtensor(data[f"teacher_{i}_input_ids"])
            token_mask = to_local_if_dtensor(data[f"teacher_{i}_token_mask"])
        return teacher_full_logits_by_idx[i], ids, token_mask

    def _compute_dynamic_weights(
        self,
        data: BatchedDataDict[CrossTokenizerDistillationLossDataDict],
        teacher_full_logits_by_idx: dict[int, torch.Tensor],
        aligns_by_idx: dict[int, 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.
        """
        device = data["input_ids"].device
        if self.normalize_teacher_by_vocab:
            temp_weight = torch.log(
                torch.tensor(float(min(self.teacher_vocab_sizes)), device=device)
            )
        scores: list[torch.Tensor] = []
        for i in range(self.num_teachers):
            t_logits, t_ids, t_mask = self._teacher_score_inputs(
                i, data, teacher_full_logits_by_idx, aligns_by_idx
            )
            score = self._teacher_weight_score(
                t_logits, t_ids, t_mask, data["sample_mask"]
            )
            if self.normalize_teacher_by_vocab:
                v_log = torch.log(
                    torch.tensor(float(self.teacher_vocab_sizes[i]), device=device)
                )
                score = score * (v_log / temp_weight)
            scores.append(score)
        weights = torch.softmax(self.alpha * torch.stack(scores), dim=0)
        return [weights[i] for i in range(self.num_teachers)]

    def _teacher_weight_score(
        self,
        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.
        """
        samp = to_local_if_dtensor(sample_mask).unsqueeze(-1).float()
        if self.sum_weights_metric == "ce":
            ce_pos = torch.nn.functional.cross_entropy(
                t_logits[:, :-1].reshape(-1, t_logits.shape[-1]).float(),
                t_ids[:, 1:].reshape(-1),
                reduction="none",
            )
            mask = (t_mask[:, 1:].float() * samp).reshape(-1)
            return -self._dp_global_masked_mean(ce_pos, mask)
        mask = t_mask.float() * samp
        if self.sum_weights_metric == "entropy":
            probs = torch.softmax(t_logits.float(), dim=-1)
            entropy = -(probs * torch.log(probs + 1e-10)).sum(dim=-1)
            return -self._dp_global_masked_mean(entropy, mask)
        if self.sum_weights_metric == "max_prob":
            probs = torch.softmax(t_logits.float(), dim=-1)
            return self._dp_global_masked_mean(probs.max(dim=-1).values, mask)
        raise ValueError(f"Unknown sum_weights_metric: {self.sum_weights_metric!r}")

    # ------------------------------------------------------------------ #
    # Loss-mode implementations
    # ------------------------------------------------------------------ #
    def _compute_p_kl(
        self,
        student_logits: torch.Tensor,
        teacher_full_logits: torch.Tensor,
        align: 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:

        1. Use the full teacher logits materialized by ``prepare_loss_input``.
        2. 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.
        3. Restrict ``M`` to those ``k`` teacher columns and project the
           student probs through it, so only the ``k`` columns are ever
           produced. Each teacher column of ``M.t() @ p`` is an independent
           contraction over the student axis, so slicing before the matmul
           is value-preserving.
        4. Slice the teacher logits to the same ``k`` columns.
        5. Build per-token chunk masks from ``alignment_*_chunk_id`` and
           chunk-average via ``bmm`` (shared helper).
        6. Renormalize student chunk distributions inside the top-k subset
           (avg-then-renormalize, log).
        7. Forward (or reverse) KL between chunk distributions.
        """
        T = self.temperature
        device = student_logits.device
        eps = 1e-10

        # Vocab-sharded (TP) students keep the shard with globally-correct
        # normalization; the result is projected onto the teacher vocab next.
        student_log_probs = vocab_parallel_log_softmax(
            student_logits, T, tp_group=tp_group
        )
        student_probs = student_log_probs.exp()  # [B, T_s_local, V_s_local]

        sparse_projection = get_sparse_projection_matrix(
            projection_matrix_path,
            device,
            student_vocab_size=self.student_vocab_size,
            teacher_vocab_size=teacher_vocab_size,
        )  # [V_s, V_t] sparse COO, fp32
        # The projection itself is deferred until the top-k columns are known
        # (see below); its teacher-vocab width is available from the matrix.
        full_teacher_vocab_size = sparse_projection.size(1)

        # `teacher_full_logits` [B, T_t, V_t_model] is materialized by
        # `prepare_loss_input` (rebuilt from the IPC handles). Same transport
        # as the gold path consumes; here we additionally compute a
        # microbatch-global top-k inline.
        # HF models commonly pad lm_head out_features beyond len(tokenizer)
        # for embedding/FFN alignment (e.g. Qwen3: tokenizer 151669,
        # lm_head 151936). The projection matrix is sized to the real
        # tokenizer vocab (the ``teacher_vocab_size`` arg); the padded
        # columns aren't real tokens and the projection has no entries
        # there. Slice to the projection's V_t to keep the projected
        # student probs and the teacher logits on the same vocab axis.
        if teacher_full_logits.shape[-1] > full_teacher_vocab_size:
            teacher_full_logits = teacher_full_logits[..., :full_teacher_vocab_size]

        # Chunk ids (localized + next-token-shifted) come from `prepare_loss_input`.
        student_chunk_id = align.student_chunk_id
        teacher_chunk_id = align.teacher_chunk_id
        pair_valid = align.pair_valid  # [B, max_pairs]
        if self.exact_token_match_only:
            pair_valid = pair_valid & align.pair_is_correct
        max_chunks = pair_valid.shape[1]

        # One microbatch-global top-k teacher subset (CP-reduced so every rank
        # agrees on the same vocab columns), shared across all samples/positions.
        vocab_topk = min(self.vocab_topk, full_teacher_vocab_size)
        global_top_indices = select_teacher_topk_indices(
            teacher_full_logits, vocab_topk, cp_group=cp_group
        )  # [k]

        # Slice both sides to the shared [k] columns. On the student side the
        # slice is folded into the projection matrix rather than applied to its
        # output: every teacher column of ``M.t() @ p`` is an independent
        # contraction over the student axis, so restricting M's columns first is
        # value-preserving, and the only renormalization here is within the [k]
        # subset (below), never over V_t. Projecting all V_t columns and then
        # discarding all but k built -- and all-reduced, under TP -- a
        # [B, T_s, V_t] fp32 tensor to read k of its columns.
        if vocab_topk < full_teacher_vocab_size:
            sparse_projection = slice_sparse_projection_cols(
                sparse_projection, global_top_indices
            )
        projected_topk = project_student_to_teacher_vocab(
            student_probs, sparse_projection, tp_group=tp_group
        )  # [B, T_s_local, k]
        teacher_topk_logits = teacher_full_logits[
            ..., global_top_indices
        ]  # [B, T_t, k]
        target_log_probs = torch.log_softmax(
            teacher_topk_logits / T, dim=-1
        )  # [B, T_t, k] (renormalized within the [k] subset).

        # Chunk-average both sides via the shared helper.
        proj_chunks, proj_sizes = chunk_average_log_probs(
            projected_topk, student_chunk_id, max_chunks, cp_group=cp_group
        )
        tgt_log_chunks, tgt_sizes = chunk_average_log_probs(
            target_log_probs, teacher_chunk_id, max_chunks, cp_group=cp_group
        )

        # Renormalize the projected chunk distribution within the top-k
        # subset, then take log. Teacher side is already log-probs (avg of
        # log_softmaxes; not a true log of mean).
        proj_chunks = proj_chunks / (proj_chunks.sum(dim=-1, keepdim=True) + eps)
        proj_log_chunks = (proj_chunks + eps).log()

        chunk_mask = valid_chunk_mask(proj_sizes, tgt_sizes, pair_valid)
        # Compute the DP-global valid-chunk count BEFORE the early
        # return so the collective fires on every rank (a local-only
        # `chunk_mask.any()` check would deadlock when one rank skips
        # and others do not). The reduction at the bottom uses this
        # global count so the KL is normalized by
        # `sum(global_valid_chunks)` rather than a per-rank mean —
        # mirrors the `global_valid_toks` convention used by CE.
        sample_mask_bool = align.sample_mask.bool()
        valid_bool = chunk_mask & sample_mask_bool.unsqueeze(-1)
        global_valid_chunks = group_all_reduce_sum(
            valid_bool.sum().to(torch.float32), group=torch.distributed.group.WORLD
        )
        if global_valid_chunks.item() == 0:
            zero = torch.zeros((), device=device, dtype=proj_log_chunks.dtype)
            return (
                zero,
                torch.zeros((), device=device, dtype=torch.long),
                zero.detach(),
            )

        # Projection top-1 accuracy: per-chunk argmax of the student-side
        # projected distribution vs the teacher's argmax over the same
        # top-k subset.
        with torch.no_grad():
            proj_top1 = proj_chunks.argmax(dim=-1)  # [B, C]
            tgt_top1 = torch.exp(tgt_log_chunks).argmax(dim=-1)  # [B, C]
            proj_matches = (proj_top1 == tgt_top1) & chunk_mask
            proj_acc = proj_matches.sum().float() / chunk_mask.sum().float().clamp(
                min=1.0
            )

        # KL between chunk-averaged distributions.
        if self.reverse_kl:
            # KL(student || teacher)
            per_chunk_kl = torch.nn.functional.kl_div(
                tgt_log_chunks, proj_log_chunks, reduction="none", log_target=True
            ).sum(dim=-1)
        else:
            # Forward KL(teacher || student)
            per_chunk_kl = torch.nn.functional.kl_div(
                proj_log_chunks, tgt_log_chunks, reduction="none", log_target=True
            ).sum(dim=-1)

        sample_mask = align.sample_mask.to(per_chunk_kl.dtype)  # [B]
        valid = chunk_mask.to(per_chunk_kl.dtype) * sample_mask.unsqueeze(-1)
        denom = global_valid_chunks.to(per_chunk_kl.dtype).clamp(min=1.0)
        kl_loss = (per_chunk_kl * valid).sum() / denom * (T * T)

        return kl_loss, valid.sum().detach(), proj_acc.detach()

    def _compute_gold(
        self,
        student_logits: torch.Tensor,
        teacher_full_logits: torch.Tensor,
        align: 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``.

        1. Lazy-build the exact-token map (cached per device).
        2. Use the full teacher logits materialized by ``prepare_loss_input``.
        3. ``log_softmax`` on full vocab both sides; chunk-average via the
           shared helper using the precomputed next-token-shifted chunk ids.
        4. 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).
        5. 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.
        6. Combine: ``loss = (kl_common + l1_uncommon) * T**2``.
        7. 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.
        """
        T = self.temperature
        device = student_logits.device

        exact_map = build_exact_token_map(
            projection_matrix_path,
            device,
            xtoken_loss=xtoken_loss,
            teacher_vocab_size=teacher_vocab_size,
        )
        common_s = exact_map["common_student"]
        common_t = exact_map["common_teacher"]
        uncommon_s = exact_map["uncommon_student"]
        uncommon_t = exact_map["uncommon_teacher"]
        v_teacher = teacher_vocab_size

        # `teacher_full_logits` [B, T_t, V_t_model] is materialized by
        # `prepare_loss_input` (rebuilt from the IPC handles).
        # Drop any padded lm_head vocab beyond the real tokenizer vocab —
        # the exact-token map's t-axis is bounded by `teacher_vocab_size`,
        # so chunked teacher log-probs must use the same axis. See the
        # matching note in `_compute_p_kl` for why the model vocab can
        # exceed `len(tokenizer)`.
        if teacher_full_logits.shape[-1] > v_teacher:
            teacher_full_logits = teacher_full_logits[..., :v_teacher]

        # common_s / uncommon_s are arbitrary V_s indices, so the gold path needs
        # full-vocab student log-probs (TP-sharded students are all-gathered).
        student_log_probs = vocab_parallel_full_log_softmax(
            student_logits, T, tp_group=tp_group
        )
        # teacher_full_logits is already vocab-full (consumer-side routing
        # P2P-concat'd across TP siblings); local log_softmax is correct.
        teacher_log_probs = torch.log_softmax(
            teacher_full_logits / T, dim=-1
        )  # [B, T_t_local, V_t]

        # Chunk ids (localized + next-token-shifted) come from `prepare_loss_input`.
        student_chunk_id = align.student_chunk_id
        teacher_chunk_id = align.teacher_chunk_id
        pair_valid = align.pair_valid
        max_chunks = pair_valid.shape[1]

        student_chunks, s_sizes = chunk_average_log_probs(
            student_log_probs, student_chunk_id, max_chunks, cp_group=cp_group
        )
        teacher_chunks, t_sizes = chunk_average_log_probs(
            teacher_log_probs, teacher_chunk_id, max_chunks, cp_group=cp_group
        )

        chunk_mask = valid_chunk_mask(s_sizes, t_sizes, pair_valid)
        # Match the P-KL path: a chunk only contributes if its alignment is
        # geometrically valid AND its sample isn't masked out by sample_mask.
        sample_mask = align.sample_mask  # [B]
        valid_chunk = chunk_mask & sample_mask.bool().unsqueeze(-1)
        zero_dtype = student_log_probs.dtype
        # Compute the DP-global valid-chunk count BEFORE any potentially
        # divergent early return so the collective fires on every rank;
        # both `kl_common` and `l1_uncommon` use this as their denom so
        # the loss is normalized by `sum(global_valid_chunks)`, not a
        # per-rank mean.
        global_valid_chunks = group_all_reduce_sum(
            valid_chunk.sum().to(torch.float32), group=torch.distributed.group.WORLD
        )
        if global_valid_chunks.item() == 0:
            zero = torch.zeros((), device=device, dtype=zero_dtype)
            return (
                zero,
                zero.detach(),
                zero.detach(),
                torch.zeros((), device=device, dtype=torch.long),
                zero.detach(),
            )

        # ---------------------- KL on common ----------------------
        if common_s.numel() > 0:
            student_common = student_chunks[:, :, common_s]  # [B, C, N_common]
            teacher_common = teacher_chunks[:, :, common_t]  # [B, C, N_common]
            if self.reverse_kl:
                kl_per_elem = torch.nn.functional.kl_div(
                    teacher_common,
                    student_common,
                    reduction="none",
                    log_target=True,
                )
            else:
                kl_per_elem = torch.nn.functional.kl_div(
                    student_common,
                    teacher_common,
                    reduction="none",
                    log_target=True,
                )
            kl_per_chunk = kl_per_elem.sum(dim=-1) * valid_chunk  # [B, C]
            kl_common = kl_per_chunk.sum() / global_valid_chunks.to(
                kl_per_chunk.dtype
            ).clamp(min=1.0)
        else:
            kl_common = torch.zeros(
                (), device=device, dtype=zero_dtype, requires_grad=True
            )
            student_common = None
            teacher_common = None

        # -------------------- L1 on uncommon ----------------------
        uncommon_topk = self.uncommon_topk
        if uncommon_s.numel() > 0 or uncommon_t.numel() > 0:
            student_unc = student_chunks[:, :, uncommon_s][
                valid_chunk
            ]  # [N_valid, N_u_s]
            teacher_unc = teacher_chunks[:, :, uncommon_t][
                valid_chunk
            ]  # [N_valid, N_u_t]
            n_valid = student_unc.shape[0]
            max_uncommon = min(
                student_unc.shape[-1],
                teacher_unc.shape[-1],
                uncommon_topk,
            )
            if n_valid > 0 and max_uncommon > 0:
                student_unc_probs = student_unc.exp()
                teacher_unc_probs = teacher_unc.exp()
                if student_unc_probs.shape[-1] > max_uncommon:
                    student_sorted = torch.topk(
                        student_unc_probs, k=max_uncommon, dim=-1, largest=True
                    ).values
                else:
                    student_sorted = student_unc_probs.sort(
                        dim=-1, descending=True
                    ).values
                if teacher_unc_probs.shape[-1] > max_uncommon:
                    teacher_sorted = torch.topk(
                        teacher_unc_probs, k=max_uncommon, dim=-1, largest=True
                    ).values
                else:
                    teacher_sorted = teacher_unc_probs.sort(
                        dim=-1, descending=True
                    ).values
                min_len = min(student_sorted.shape[-1], teacher_sorted.shape[-1])
                student_sorted = student_sorted[:, :min_len]
                teacher_sorted = teacher_sorted[:, :min_len]
                l1_per_chunk = torch.nn.functional.l1_loss(
                    student_sorted, teacher_sorted, reduction="none"
                ).sum(dim=-1)
                l1_uncommon = l1_per_chunk.sum() / global_valid_chunks.to(
                    l1_per_chunk.dtype
                ).clamp(min=1.0)
            else:
                l1_uncommon = torch.zeros(
                    (), device=device, dtype=zero_dtype, requires_grad=True
                )
        else:
            l1_uncommon = torch.zeros(
                (), device=device, dtype=zero_dtype, requires_grad=True
            )

        # -------------------- Top-1 accuracy ----------------------
        with torch.no_grad():
            if student_common is not None:
                s_common_valid = student_common[valid_chunk]
                t_common_valid = teacher_common[valid_chunk]
                matches = (
                    (s_common_valid.argmax(dim=-1) == t_common_valid.argmax(dim=-1))
                    .sum()
                    .float()
                )
                top1_acc = matches / valid_chunk.sum().float().clamp(min=1.0)
            else:
                top1_acc = torch.zeros((), device=device, dtype=zero_dtype)

        loss = (kl_common + l1_uncommon) * (T * T)
        return (
            loss,
            kl_common.detach(),
            l1_uncommon.detach(),
            valid_chunk.sum().detach(),
            top1_acc.detach(),
        )

    def _compute_ce(
        self,
        logits: torch.Tensor,
        data: BatchedDataDict[CrossTokenizerDistillationLossDataDict],
        global_valid_toks: torch.Tensor,
        *,
        student_next_token_logprobs: Optional[torch.Tensor] = None,
        student_next_token_mask: Optional[torch.Tensor] = None,
        cp_group: Optional[torch.distributed.ProcessGroup] = None,
    ) -> torch.Tensor:
        """Next-token CE on the student side (TP/CP handled by the helpers)."""
        if student_next_token_logprobs is not None:
            assert student_next_token_mask is not None
            label_mask = student_next_token_mask.to(
                student_next_token_logprobs.dtype
            ) * to_local_if_dtensor(data["sample_mask"]).to(
                student_next_token_logprobs.device
            ).unsqueeze(-1)
            local_ce = masked_mean(
                -student_next_token_logprobs,
                label_mask,
                global_normalization_factor=global_valid_toks,
            )
            # Forward SUM gives every CP rank the same full-sequence CE (and
            # therefore the same dynamic KD scale). Identity backward keeps the
            # disjoint rank-local windows at a single gradient fanout.
            return group_all_reduce_sum_with_grad(local_ce, cp_group)

        per_token_ce = student_next_token_ce(logits, input_ids=data["input_ids"])
        label_mask = ce_label_mask(
            token_mask=data["token_mask"],
            sample_mask=data["sample_mask"],
            ce_seq_len=per_token_ce.shape[1],
            dtype=per_token_ce.dtype,
        )
        return masked_mean(
            per_token_ce,
            label_mask,
            global_normalization_factor=global_valid_toks,
        )