Alignment

Functional, mechanism-agnostic primitives for DP-safe preference learning and supervised fine-tuning: per-example SFT and DPO losses, per-sequence log-probability helpers, batch collators, dataset transforms, reference-model handling, and reward metrics. Every public symbol is a pure function or a factory returning a callable, so each composes under vmap(grad(...)) — the per-example differentiation DP-SGD and DP-FTRL clipping is built on. Import everything from the opake.alignment.* façades; the mechanism and optimizer are chosen at the call site, never inside this package. For end-to-end walkthroughs see the SFT and DPO guides.

Overview

Losses

Per-example loss functions. SFT losses return the per-example mean over an example's own non-ignored tokens (a DP-safe divisor, not a batch aggregate). DPO per-pair heads take (chosen_logratio, rejected_logratio, *, beta, ...) and return a per-example scalar; a DPO loss is head(sequence_logp(...) - ref_logp, ...). Eager losses take logits; the fused_* twins take hidden_states + the lm_head weight and never materialize (T, V) logits (CUDA + half-precision fused kernel with opake-alignment[patches], eager fallback otherwise).

SFT:

• nll_loss() (opake.alignment.sft.loss) — per-example causal-LM cross-entropy with a per-example token-count divisor.
• dft_loss() (opake.alignment.sft.loss) — Dynamic Fine-Tuning: NLL weighted by the detached softmax probability of each target token.
• fused_nll_loss() (opake.alignment.sft.loss) — memory-efficient nll_loss over hidden states + lm_head weight.
• fused_dft_loss() (opake.alignment.sft.loss) — memory-efficient dft_loss over hidden states + lm_head weight.

DPO per-pair heads (opake.alignment.dpo.loss):

• sigmoid_loss() — standard DPO sigmoid (logistic) loss; the default.
• hinge_loss() — DPO hinge loss.
• robust_loss() — robust (label-smoothed) DPO loss.
• ipo_loss() — IPO loss (Azar 2024).
• discopop_loss() — DiscoPOP loss.
• chosen_nll_loss() — chosen-completion NLL regularizer for MPO/RPO blends.
• apo_zero_loss() / apo_down_loss() — APO-zero / APO-down (arXiv:2408.06266).
• exo_loss() — EXO pairwise loss (Towards Efficient Exact Optimization of Language Model Alignment).
• nca_loss() — NCA pairwise loss.
• bco_loss() — BCO pairwise loss.
• sppo_loss() — SPPO hard-label loss.

Reference-free heads (opake.alignment.dpo.loss) — no reference model; score the policy log-prob directly:

• simpo_loss() — SimPO: length-normalized sigmoid with a target margin γ (Meng 2024).
• odds_ratio_loss() — ORPO odds-ratio loss on length-normalized log-probs (Hong 2024); pair with chosen_nll_loss via mpo_combine.

DPO log-ratio combinators (opake.alignment.dpo.loss) for composite objectives:

• f_divergence_remap() — remap a per-example log-ratio under the chosen f-divergence.
• f_divergence_logits() — form the remapped preference logits under an f-divergence.
• mpo_combine() — weighted sum of per-example loss tensors (MPO / TRL loss_type=list).
• wpo_weights() — per-example WPO weight (arXiv:2406.11827).
• ld_dpo_split() — length-desensitized sequence-logp split (LD-DPO, arXiv:2409.06411).

Primary papers for the implemented alignment objectives:

• DPO — Rafailov et al., Direct Preference Optimization: Your Language Model is Secretly a Reward Model.
• IPO — Azar et al., A General Theoretical Paradigm to Understand Learning from Human Preferences.
• DiscoPOP — Lu et al., Discovering Preference Optimization Algorithms with and for Large Language Models.
• SimPO — Meng et al., SimPO: Simple Preference Optimization with a Reference-Free Reward.
• ORPO — Hong et al., ORPO: Monolithic Preference Optimization without Reference Model.
• WPO — Zhou et al., WPO: Enhancing RLHF with Weighted Preference Optimization.
• LD-DPO — Liu et al., Length Desensitization in Direct Preference Optimization.
• APO — D'Oosterlinck et al., Anchored Preference Optimization and Contrastive Revisions: Addressing Underspecification in Alignment.
• NCA — Chen et al., Noise Contrastive Alignment of Language Models with Explicit Rewards.
• BCO — Jhunjhunwala et al., Binary Classifier Optimization for Large Language Model Alignment.
• SPPO — Wu et al., Self-Play Preference Optimization for Language Model Alignment.

Log-probabilities

• sequence_logp() (opake.alignment.dpo.loss) — sum of completion-token log-probabilities per sequence; applies the causal-LM shift, masks to the completion span, and sums. Works per-example or on a batch axis. Pass length_normalized=True for the per-token mean reward log π(y)/|y| (SimPO / ORPO), or ld_alpha for the LD-DPO length-desensitized split. Its compatibility argument shared_prefix_len is the paper's public length (the shorter completion length), not a token-identical prefix.
• fused_sequence_logp() (opake.alignment.dpo.loss) — memory-efficient sequence_logp over hidden states + lm_head weight (per-example only; same length_normalized option).

Collators

Factory functions returning a collate(examples) callable.

• language_modeling_collator() (opake.alignment.sft.collator) — SFT/causal-LM collation: input_ids/attention_mask/labels (B, L) with pad and (optionally) prompt tokens masked to -100, optional completion_mask, keep-start truncation to max_length. Output schema LMBatch (opake.alignment.sft.collator.types).
• preference_collator() (opake.alignment.dpo.collator) — DPO collation in a (B, ...) layout: the chosen and rejected input_ids/attention_mask/completion_mask trios, plus optional precomputed ref_chosen_logps / ref_rejected_logps (B,). Chosen and rejected stay separate (not TRL's concatenated (2B, L)) so one pair maps to one clipped gradient.
• extract_prompt() (opake.alignment.dpo.data) — split an implicit shared prompt out of a preference example before tokenizing.

Data

Shared chat-template data prep (opake.alignment.data). Install a training chat template, then tokenize chat turns into input_ids + a completion_mask for completion-only loss.

• clone_chat_template() — copy a chat template (and its special tokens) from a source tokenizer onto a destination tokenizer and resize the model's embedding matrix to match. Returns the 3-tuple (model, tokenizer, added_token_ids) — the (possibly resized) model, the updated tokenizer, and the ids of any newly added special tokens (mark these trainable when the rest of the model is frozen). Call it before make_functional (the embedding resize must be captured in the functional snapshot).
• get_training_chat_template() — return a chat-template string carrying validated assistant-turn {% generation %} markers so the assistant-token mask is recoverable at tokenization time. Unsupported templates raise instead of falling back to a potentially incorrect mask.
• apply_chat_template_with_mask() — tokenize a chat conversation into input_ids + completion_mask (1 on assistant tokens), the mask consumed by language_modeling_collator(..., completion_only_loss=True).

Reference handling

Reference-model helpers for DPO. These run outside the per-example vmap(grad(...)) region — a separate forward pass, or PEFT adapter toggles.

• compute_ref_logprobs_for_dataset() (opake.alignment.dpo.reference) — run the reference once over a dataset, attach per-example logp columns, and cache to a content-addressed .safetensors file keyed by dataset identity, cache_identity, and output_columns (a cache hit skips the forward). Cache directories and archives are owner-only; callers remove selected cache directories when their private per-example values are no longer needed. Set use_cache=False for non-reusable one-shot values; TR-DPO uses this mode for its seed columns. Under a live process group each rank scores its own shard and the columns are gathered back into dataset order, so the result matches a single-process run — see Across ranks.
• null_ref_context() (opake.alignment.dpo.reference) — context manager that turns a model into its own reference, dispatching over the separate-model / "ref"-adapter / disabled-adapter / no-op configurations.
• with_disabled_adapter() (opake.alignment.dpo.reference) — context manager that disables a PEFT adapter so the base model serves as the reference.
• ema_update_reference() (opake.alignment.dpo.reference) — TR-DPO: leafwise EMA step (1 - alpha) * ref + alpha * policy to move the reference toward the policy.

Metrics

These metrics are un-noised and outside Opake's DP accounting. See Telemetry outside the guarantee. Shared token metrics (opake.alignment.metric):

• mean_token_accuracy() — mean next-token argmax accuracy over the supervised (non-ignored) positions.
• entropy_from_logits() — mean per-token predictive entropy.

Preference reward telemetry (opake.alignment.dpo.metric):

• reward_metrics() — rewards/chosen, rewards/rejected, rewards/accuracies, rewards/margins from per-example log-ratios.

See also: Alignment overview, SFT end-to-end, DPO end-to-end.

API Documentation

loss

opake.alignment.sft.loss façade — re-exports the SFT loss functions.

nll_loss / dft_loss (on logits) and their memory-efficient fused twins fused_nll_loss / fused_dft_loss (on hidden states + lm_head weight).

dft_loss(logits, labels)

Per-example DFT loss with detached softmax weighting and DP-safe divisor.

Shifts logits[..., :-1, :] against labels[..., 1:], computes the per-token DFT loss -detach(softmax_prob) * logp, masks positions where the shifted label is -100, and returns the per-example mean over non-ignored tokens using a per-example divisor (DP-safe).

The weighting probability p is detached (stop-gradient), so the gradient of the loss w.r.t. logits equals that of -p.detach() * logp. This matches the TRL dft formula in spirit but uses the DP-corrected per-example divisor instead of TRL's batch-level num_items_in_batch.

Parameters:

Name	Type	Description	Default
logits	Tensor	Float tensor of shape (..., T, V) where T is the sequence length and V is the vocabulary size.	required
labels	Tensor	Integer tensor of shape (..., T). Values of -100 are ignored (set by the collator for prompt / pad tokens).	required

Returns:

Type	Description
Tensor	Float tensor of shape (...) (0-dim for a single example, (B,)
Tensor	for a batch). Each entry is
Tensor	the mean per-token loss over non-ignored shifted labels.

Source code in packages/opake-alignment/src/opake/api/alignment/sft/loss/_dft.py

def dft_loss(
    logits: torch.Tensor,
    labels: torch.Tensor,
) -> torch.Tensor:
    """Per-example DFT loss with detached softmax weighting and DP-safe divisor.

    Shifts ``logits[..., :-1, :]`` against ``labels[..., 1:]``, computes
    the per-token DFT loss ``-detach(softmax_prob) * logp``, masks positions
    where the shifted label is ``-100``, and returns the **per-example mean**
    over non-ignored tokens using a per-example divisor (DP-safe).

    The weighting probability ``p`` is detached (stop-gradient), so the
    gradient of the loss w.r.t. ``logits`` equals that of ``-p.detach() *
    logp``.  This matches the TRL ``dft`` formula in spirit but uses the
    DP-corrected per-example divisor instead of TRL's batch-level
    ``num_items_in_batch``.

    Args:
        logits: Float tensor of shape ``(..., T, V)`` where ``T`` is the
            sequence length and ``V`` is the vocabulary size.
        labels: Integer tensor of shape ``(..., T)``.  Values of ``-100``
            are ignored (set by the collator for prompt / pad tokens).

    Returns:
        Float tensor of shape ``(...)`` (0-dim for a single example, ``(B,)``
        for a batch).  Each entry is
        the mean per-token loss over non-ignored shifted labels.
    """
    # Shift: logits[..., :-1, :] predicts labels[..., 1:]
    shifted_logits = logits[..., :-1, :]  # (..., T-1, V)
    shifted_labels = labels[..., 1:]  # (..., T-1)

    # Mask: 1 where the label is a real token (not -100)
    mask = (shifted_labels != _IGNORE_INDEX).to(shifted_logits.dtype)  # (..., T-1)

    # Clamp -100 → 0 so gather does not receive an out-of-range index.
    # The mask zeroes out the contribution of these positions afterwards.
    clamped_labels = shifted_labels.clamp(min=0)  # (..., T-1)

    # Per-token log-probabilities (differentiable through logits)
    logp = selective_log_softmax(shifted_logits, clamped_labels)  # (..., T-1)

    # Detached softmax probability at the target token (stop-gradient)
    # torch.softmax + gather, then detach — gradient does NOT flow through p.
    p = (
        torch.softmax(shifted_logits, dim=-1)
        .gather(-1, clamped_labels.unsqueeze(-1))
        .squeeze(-1)
        .detach()
    )  # (..., T-1), no gradient

    # Per-token DFT loss: -p * logp  (p is detached; logp is differentiable)
    per_token = -p * logp  # (..., T-1)

    # Accumulate in >= fp32 with an exact token count so eager matches fused.
    acc_dtype = _compute_dtype(per_token)
    n_valid = (shifted_labels != _IGNORE_INDEX).sum(-1).clamp(min=1)
    loss = (per_token * mask).to(acc_dtype).sum(-1) / n_valid.to(acc_dtype)  # (...)
    return loss

fused_dft_loss(hidden_states, lm_head_weight, labels)

Memory-efficient :func:dft_loss over hidden states (fused, per-example).

Mathematically dft_loss(hidden_states @ lm_head_weight.T, labels), but on CUDA + half precision with opake-alignment[patches] installed it computes the detached-confidence-weighted per-token CE against the projection via the patches fused linear-CE kernel in its use_token_scaling mode — no (T, V) logits are materialised and the LSE is recomputed in the backward. Otherwise it falls back to the eager form.

Per-example. Pass one example (T, H) and drive with vmap(grad(...)) (the clipped_grad DP-SGD path); see :func:~opake.api.alignment.sft.loss.fused_nll_loss for the per-example contract.

Parameters:

Name	Type	Description	Default
hidden_states	Tensor	Last-layer hidden states (T, H) for one example (batched to (B, T, H) only by an outer vmap).	required
lm_head_weight	Tensor	LM-head weight (V, H); logits are hidden_states @ lm_head_weight.T.	required
labels	Tensor	Token-id targets (T,); -100 positions are ignored.	required

Returns:

Type	Description
Tensor	Scalar per-example DFT loss (matches :func:dft_loss).

Source code in packages/opake-alignment/src/opake/api/alignment/sft/loss/_dft.py

def fused_dft_loss(
    hidden_states: torch.Tensor,
    lm_head_weight: torch.Tensor,
    labels: torch.Tensor,
) -> torch.Tensor:
    """Memory-efficient :func:`dft_loss` over hidden states (fused, per-example).

    Mathematically ``dft_loss(hidden_states @ lm_head_weight.T, labels)``, but on
    CUDA + half precision with ``opake-alignment[patches]`` installed it computes
    the detached-confidence-weighted per-token CE against the projection via the
    patches fused linear-CE kernel in its ``use_token_scaling`` mode — no ``(T,
    V)`` logits are materialised and the LSE is recomputed in the backward.
    Otherwise it falls back to the eager form.

    **Per-example.** Pass one example ``(T, H)`` and drive with
    ``vmap(grad(...))`` (the ``clipped_grad`` DP-SGD path); see
    :func:`~opake.api.alignment.sft.loss.fused_nll_loss` for the per-example
    contract.

    Args:
        hidden_states: Last-layer hidden states ``(T, H)`` for one example
            (batched to ``(B, T, H)`` only by an outer ``vmap``).
        lm_head_weight: LM-head weight ``(V, H)``; logits are
            ``hidden_states @ lm_head_weight.T``.
        labels: Token-id targets ``(T,)``; ``-100`` positions are ignored.

    Returns:
        Scalar per-example DFT loss (matches :func:`dft_loss`).
    """
    if lce_available(hidden_states):
        ce_sum = linear_nll_sum(
            hidden_states, lm_head_weight, labels, use_token_scaling=True
        )
        n_valid = (labels[..., 1:] != _IGNORE_INDEX).sum(-1).clamp(min=1)
        return ce_sum / n_valid
    return dft_loss(hidden_states @ lm_head_weight.transpose(-2, -1), labels)

fused_nll_loss(hidden_states, lm_head_weight, labels)

Memory-efficient :func:nll_loss over hidden states (fused, per-example).

Mathematically nll_loss(hidden_states @ lm_head_weight.T, labels), but on CUDA + half precision with opake-alignment[patches] installed it computes the per-token CE against the projection via the patches fused linear-CE kernel — no (T, V) logits are materialised and the LSE is recomputed in the backward. Otherwise it falls back to the eager form.

Per-example. Pass one example (T, H) and drive with vmap(grad(...)) (the clipped_grad DP-SGD path); the kernel's merged vmap rules then make the whole microbatch one forward + one backward kernel. Unlike :func:nll_loss, this is not meant to be called directly on a batch axis (the fused path would collapse it to a single scalar).

Parameters:

Name	Type	Description	Default
hidden_states	Tensor	Last-layer hidden states (T, H) for one example (batched to (B, T, H) only by an outer vmap).	required
lm_head_weight	Tensor	LM-head weight (V, H); logits are hidden_states @ lm_head_weight.T.	required
labels	Tensor	Token-id targets (T,); -100 positions are ignored.	required

Returns:

Type	Description
Tensor	Scalar per-example mean NLL (matches :func:nll_loss).

Source code in packages/opake-alignment/src/opake/api/alignment/sft/loss/_nll.py

def fused_nll_loss(
    hidden_states: torch.Tensor,
    lm_head_weight: torch.Tensor,
    labels: torch.Tensor,
) -> torch.Tensor:
    """Memory-efficient :func:`nll_loss` over hidden states (fused, per-example).

    Mathematically ``nll_loss(hidden_states @ lm_head_weight.T, labels)``, but on
    CUDA + half precision with ``opake-alignment[patches]`` installed it computes
    the per-token CE against the projection via the patches fused linear-CE
    kernel — no ``(T, V)`` logits are materialised and the LSE is recomputed in
    the backward. Otherwise it falls back to the eager form.

    **Per-example.** Pass one example ``(T, H)`` and drive with
    ``vmap(grad(...))`` (the ``clipped_grad`` DP-SGD path); the kernel's merged
    vmap rules then make the whole microbatch one forward + one backward kernel.
    Unlike :func:`nll_loss`, this is *not* meant to be called directly on a batch
    axis (the fused path would collapse it to a single scalar).

    Args:
        hidden_states: Last-layer hidden states ``(T, H)`` for one example
            (batched to ``(B, T, H)`` only by an outer ``vmap``).
        lm_head_weight: LM-head weight ``(V, H)``; logits are
            ``hidden_states @ lm_head_weight.T``.
        labels: Token-id targets ``(T,)``; ``-100`` positions are ignored.

    Returns:
        Scalar per-example mean NLL (matches :func:`nll_loss`).
    """
    if lce_available(hidden_states):
        ce_sum = linear_nll_sum(hidden_states, lm_head_weight, labels)
        n_valid = (labels[..., 1:] != _IGNORE_INDEX).sum(-1).clamp(min=1)
        return ce_sum / n_valid
    return nll_loss(hidden_states @ lm_head_weight.transpose(-2, -1), labels)

nll_loss(logits, labels)

Per-example causal-LM NLL loss with per-example DP-safe divisor.

Shifts logits[..., :-1, :] against labels[..., 1:] (standard causal-LM convention), masks positions where the shifted label is -100, computes selective_log_softmax at the (clamped) shifted labels, and returns the per-example mean of the negative per-token log-probabilities over non-ignored tokens.

The exact shifted-label token count is computed per example from this example's data only.

Parameters:

Name	Type	Description	Default
logits	Tensor	Float tensor of shape (..., T, V) where T is the sequence length and V is the vocabulary size. Leading dims (...) may be absent (per-example) or a batch axis (B,).	required
labels	Tensor	Integer tensor of shape (..., T) matching the leading dims of logits. Values of -100 are ignored.	required

Returns:

Type	Description
Tensor	Float tensor of shape (...) (0-dim for a single example, shape
Tensor	(B,) for a batch). Each entry is the per-example mean NLL over
Tensor	non-ignored shifted tokens.

Source code in packages/opake-alignment/src/opake/api/alignment/sft/loss/_nll.py

def nll_loss(
    logits: torch.Tensor,
    labels: torch.Tensor,
) -> torch.Tensor:
    """Per-example causal-LM NLL loss with per-example DP-safe divisor.

    Shifts ``logits[..., :-1, :]`` against ``labels[..., 1:]`` (standard
    causal-LM convention), masks positions where the shifted label is
    ``-100``, computes ``selective_log_softmax`` at the (clamped) shifted
    labels, and returns the **per-example mean** of the negative per-token
    log-probabilities over non-ignored tokens.

    The exact shifted-label token count is computed per example from this
    example's data only.

    Args:
        logits: Float tensor of shape ``(..., T, V)`` where ``T`` is the
            sequence length and ``V`` is the vocabulary size.  Leading dims
            ``(...)`` may be absent (per-example) or a batch axis ``(B,)``.
        labels: Integer tensor of shape ``(..., T)`` matching the leading
            dims of ``logits``.  Values of ``-100`` are ignored.

    Returns:
        Float tensor of shape ``(...)`` (0-dim for a single example, shape
        ``(B,)`` for a batch).  Each entry is the per-example mean NLL over
        non-ignored shifted tokens.
    """
    # Shift: logits[..., :-1, :] predicts labels[..., 1:]
    shifted_logits = logits[..., :-1, :]  # (..., T-1, V)
    shifted_labels = labels[..., 1:]  # (..., T-1)

    # Mask: 1 where the label is a real token (not -100)
    mask = (shifted_labels != _IGNORE_INDEX).to(shifted_logits.dtype)  # (..., T-1)

    # Clamp -100 → 0 so gather does not receive an out-of-range index.
    # The mask zeros out the contribution of these positions.
    clamped_labels = shifted_labels.clamp(min=0)  # (..., T-1)

    # Per-token log-probabilities at the (clamped) target indices
    logp = selective_log_softmax(shifted_logits, clamped_labels)  # (..., T-1)

    # Accumulate in >= fp32 with an exact token count so eager matches fused.
    acc_dtype = _compute_dtype(logp)
    n_valid = (shifted_labels != _IGNORE_INDEX).sum(-1).clamp(min=1)
    nll = (-logp * mask).to(acc_dtype).sum(-1) / n_valid.to(acc_dtype)  # (...)
    return nll

collator

opake.alignment.sft.collator façade — the language-modeling collator.

Output schema (:class:LMBatch) lives in :mod:opake.alignment.sft.collator.types.

language_modeling_collator(pad_token_id, max_length, *, completion_only_loss=False, pad_to_multiple_of=None)

Return a callable that collates per-example dicts into a padded :class:LMBatch.

This is a factory function: it returns a plain callable, not a user-instantiated class. The returned callable is stateless and deterministic — the same input list always produces identical output tensors.

Parameters:

Name	Type	Description	Default
pad_token_id	int	Token id used to right-pad sequences to the batch length.	required
max_length	int	Maximum sequence length. Examples longer than this are truncated from the right (keep-start) before padding.	required
completion_only_loss	bool	When True, positions where completion_mask == 0 are set to -100 in labels so that prompt tokens do not contribute to the language-modelling loss. Defaults to False (standard next-token prediction over the full sequence).	False
pad_to_multiple_of	int | None	When set, the padded length L is rounded up to the nearest multiple of this value before the batch tensors are allocated. Useful for tensor-core alignment (e.g. 8 or 64). None disables rounding.	None

Returns:

Type	Description
Callable[[list[LMExample]], LMBatch]	A collate(examples: list[LMExample]) -> LMBatch
Callable[[list[LMExample]], LMBatch]	callable. See :class:~opake.api.alignment.sft.collator.types.LMBatch
Callable[[list[LMExample]], LMBatch]	for the output schema.

Example::

collate = language_modeling_collator(pad_token_id=0, max_length=16)
batch = collate([{"input_ids": [1, 2, 3]}, {"input_ids": [4, 5]}])
# batch["input_ids"].shape == (2, 3)

Source code in packages/opake-alignment/src/opake/api/alignment/sft/collator/_language_modeling.py

def language_modeling_collator(
    pad_token_id: int,
    max_length: int,
    *,
    completion_only_loss: bool = False,
    pad_to_multiple_of: int | None = None,
) -> Callable[[list[LMExample]], LMBatch]:
    """Return a callable that collates per-example dicts into a padded :class:`LMBatch`.

    This is a factory function: it returns a
    plain callable, not a user-instantiated class.  The returned callable is
    stateless and deterministic — the same input list always produces identical
    output tensors.

    Args:
        pad_token_id: Token id used to right-pad sequences to the batch
            length.
        max_length: Maximum sequence length.  Examples longer than this are
            truncated from the right (keep-start) before padding.
        completion_only_loss: When ``True``, positions where
            ``completion_mask == 0`` are set to ``-100`` in ``labels`` so that
            prompt tokens do not contribute to the language-modelling loss.
            Defaults to ``False`` (standard next-token prediction over the
            full sequence).
        pad_to_multiple_of: When set, the padded length ``L`` is rounded up to
            the nearest multiple of this value before the batch tensors are
            allocated.  Useful for tensor-core alignment (e.g. ``8`` or
            ``64``).  ``None`` disables rounding.

    Returns:
        A ``collate(examples: list[LMExample]) -> LMBatch``
        callable.  See :class:`~opake.api.alignment.sft.collator.types.LMBatch`
        for the output schema.

    Example::

        collate = language_modeling_collator(pad_token_id=0, max_length=16)
        batch = collate([{"input_ids": [1, 2, 3]}, {"input_ids": [4, 5]}])
        # batch["input_ids"].shape == (2, 3)
    """
    return _LMCollator(
        pad_token_id=pad_token_id,
        max_length=max_length,
        completion_only_loss=completion_only_loss,
        pad_to_multiple_of=pad_to_multiple_of,
    )

types

opake.alignment.sft.collator types façade — re-exports the output schema.

LMBatch

Bases: TypedDict

Output schema for :func:language_modeling_collator.

All present fields are :class:torch.Tensor instances with dtype=torch.long.

Required keys (always present):

input_ids Token ids, shape (B, L). Right-padded to L with the pad_token_id supplied to the factory; each example is first truncated (keep-start) to max_length before padding. L = min(max_length, max_example_length_in_batch), rounded up to the nearest multiple of pad_to_multiple_of when that argument is set.

attention_mask Binary mask, shape (B, L). 1 on real tokens, 0 on pad positions.

labels Token ids with pad positions replaced by -100 (the ignore_index recognised by :class:torch.nn.CrossEntropyLoss), shape (B, L). When the collator was created with completion_only_loss=True the non-completion positions (where the completion_mask field of the raw example is 0) are also set to -100.

Optional keys (present only when applicable):

completion_mask Binary mask, shape (B, L). Included when at least one input example in the batch carries a "completion_mask" field. 1 on completion tokens, 0 elsewhere (including pad positions). If an example in such a mixed batch has no mask, all its real tokens are marked as completion tokens to match its fully supervised labels.

Source code in packages/opake-alignment/src/opake/api/alignment/sft/collator/types.py

class LMBatch(TypedDict):
    """Output schema for :func:`language_modeling_collator`.

    All present fields are :class:`torch.Tensor` instances with ``dtype=torch.long``.

    Required keys (always present):

    ``input_ids``
        Token ids, shape ``(B, L)``.  Right-padded to ``L`` with the
        ``pad_token_id`` supplied to the factory; each example is first
        truncated (keep-start) to ``max_length`` before padding.
        ``L = min(max_length, max_example_length_in_batch)``, rounded up to
        the nearest multiple of ``pad_to_multiple_of`` when that argument is
        set.

    ``attention_mask``
        Binary mask, shape ``(B, L)``.  ``1`` on real tokens, ``0`` on pad
        positions.

    ``labels``
        Token ids with pad positions replaced by ``-100`` (the
        ``ignore_index`` recognised by :class:`torch.nn.CrossEntropyLoss`),
        shape ``(B, L)``.  When the collator was created with
        ``completion_only_loss=True`` the non-completion positions (where the
        ``completion_mask`` field of the raw example is ``0``) are also set to
        ``-100``.

    Optional keys (present only when applicable):

    ``completion_mask``
        Binary mask, shape ``(B, L)``.  Included when **at least one** input
        example in the batch carries a ``"completion_mask"`` field.  ``1`` on
        completion tokens, ``0`` elsewhere (including pad positions). If an
        example in such a mixed batch has no mask, all its real tokens are
        marked as completion tokens to match its fully supervised labels.
    """

    input_ids: torch.Tensor
    attention_mask: torch.Tensor
    labels: torch.Tensor
    completion_mask: NotRequired[torch.Tensor]

LMExample

Bases: TypedDict

One tokenized causal-language-model example accepted by the collator.

Source code in packages/opake-alignment/src/opake/api/alignment/sft/collator/types.py

class LMExample(TypedDict):
    """One tokenized causal-language-model example accepted by the collator."""

    input_ids: list[int]
    completion_mask: NotRequired[list[int]]

data

opake.alignment.data — shared chat-template data-prep (method namespace).

Method-agnostic data prep for chat datasets: install a training chat template, then tokenize chat turns into input_ids + a completion_mask for completion-only loss.

• clone_chat_template — copy a chat template (and its special tokens) from a source tokenizer onto a destination tokenizer and resize the model's embedding matrix to match.
• get_training_chat_template — return a chat-template string carrying assistant-turn {% generation %} markers so the assistant-token mask can be recovered at tokenization time.
• apply_chat_template_with_mask — tokenize a chat conversation into input_ids + completion_mask (1 on assistant tokens), the mask consumed by sft.collator.language_modeling_collator(..., completion_only_loss=True).

apply_chat_template_with_mask(tokenizer, conversation, **kwargs)

Tokenize a chat conversation and return the assistant-completion mask.

Wraps tokenizer.apply_chat_template(..., tokenize=True, return_assistant_tokens_mask=True, return_dict=True, **kwargs) and renames Hugging Face's assistant_masks to completion_mask (1 on assistant-turn tokens, 0 elsewhere) so the result drops straight into :func:opake.alignment.sft.collator.language_modeling_collator with completion_only_loss=True.

The active template (including a selected named template) must carry a Transformers-recognized generation tag around assistant content; install canonical markers with :func:opake.api.alignment.data._chat_template.get_training_chat_template. Without the marker HF returns no assistant mask and this function raises a :class:ConfigurationError.

Parameters:

Name	Type	Description	Default
tokenizer	PreTrainedTokenizerBase	A tokenizer whose active chat_template carries a Transformers-recognized generation tag (see :func:get_training_chat_template).	required
conversation	list[dict[str, Any]]	A list of chat-message dicts (each with "role" and "content") for a single conversation.	required
**kwargs	Any	Forwarded verbatim to tokenizer.apply_chat_template (e.g. max_length, truncation, add_generation_prompt). The tokenize, return_assistant_tokens_mask, and return_dict arguments are fixed by this function and must not be overridden.	{}

Returns:

Type	Description
dict[str, Any]	A dict with "input_ids" and "completion_mask" (and
dict[str, Any]	"attention_mask" when HF returns one). Both lists have the same
dict[str, Any]	length.

Raises:

Type	Description
ValueError	If the chat template does not carry {% generation %} markers, or if the tokenized result carries no (or an empty) assistant_masks. Without the markers Hugging Face only logs a warning and returns an all-zero mask (no assistant tokens flagged), so this function checks the template up front and raises instead. Install the markers with :func:get_training_chat_template.

Source code in packages/opake-alignment/src/opake/api/alignment/data/_completion_mask.py

def apply_chat_template_with_mask(
    tokenizer: PreTrainedTokenizerBase,
    conversation: list[dict[str, Any]],
    **kwargs: Any,
) -> dict[str, Any]:
    """Tokenize a chat *conversation* and return the assistant-completion mask.

    Wraps ``tokenizer.apply_chat_template(..., tokenize=True,
    return_assistant_tokens_mask=True, return_dict=True, **kwargs)`` and renames
    Hugging Face's ``assistant_masks`` to ``completion_mask`` (``1`` on
    assistant-turn tokens, ``0`` elsewhere) so the result drops straight into
    :func:`opake.alignment.sft.collator.language_modeling_collator` with
    ``completion_only_loss=True``.

    The active template (including a selected named template) **must** carry a
    Transformers-recognized generation tag around assistant content; install
    canonical markers with
    :func:`opake.api.alignment.data._chat_template.get_training_chat_template`.
    Without the marker HF returns no assistant mask and this function raises a
    :class:`ConfigurationError`.

    Args:
        tokenizer: A tokenizer whose active ``chat_template`` carries a
            Transformers-recognized generation tag (see
            :func:`get_training_chat_template`).
        conversation: A list of chat-message dicts (each with ``"role"`` and
            ``"content"``) for a single conversation.
        **kwargs: Forwarded verbatim to
            ``tokenizer.apply_chat_template`` (e.g. ``max_length``,
            ``truncation``, ``add_generation_prompt``).  The ``tokenize``,
            ``return_assistant_tokens_mask``, and ``return_dict`` arguments are
            fixed by this function and must not be overridden.

    Returns:
        A dict with ``"input_ids"`` and ``"completion_mask"`` (and
        ``"attention_mask"`` when HF returns one).  Both lists have the same
        length.

    Raises:
        ValueError: If the chat template does not carry ``{% generation %}``
            markers, or if the tokenized result carries no (or an empty)
            ``assistant_masks``.  Without the markers Hugging Face only logs a
            warning and returns an all-zero mask (no assistant tokens flagged),
            so this function checks the template up front and raises instead.
            Install the markers with :func:`get_training_chat_template`.
    """
    # Hugging Face only populates assistant_masks when the active chat template
    # contains the '{% generation %}' keyword; otherwise it merely logs a
    # warning and returns an all-zero mask.  Guard up front so the caller gets a
    # clear error (pointing at get_training_chat_template) instead of a silently
    # empty mask.  The template passed via kwargs (if any) takes precedence over
    # the tokenizer's own chat_template, matching apply_chat_template.
    active_template = _resolve_chat_template(
        tokenizer,
        chat_template=kwargs.get("chat_template"),
        tools=kwargs.get("tools"),
    )
    if not _has_generation_marker(active_template):
        raise ConfigurationError(
            *(
                "apply_chat_template_with_mask: the active chat template does not "
                "carry the '{% generation %}' / '{% endgeneration %}' markers that "
                "return_assistant_tokens_mask=True relies on, so no "
                "assistant-token mask can be produced.  Install them first with "
                "opake.alignment.data.get_training_chat_template:\n\n"
                "    tokenizer.chat_template = get_training_chat_template(tokenizer)\n",
            )
        )

    encoded = tokenizer.apply_chat_template(
        conversation,
        tokenize=True,
        return_assistant_tokens_mask=True,
        return_dict=True,
        **kwargs,
    )

    assistant_masks = encoded.get("assistant_masks")
    if not _has_assistant_tokens(assistant_masks):
        raise ConfigurationError(
            *(
                "apply_chat_template_with_mask: the tokenizer returned no "
                "assistant-token mask.  This means the chat template does not "
                "carry the '{% generation %}' / '{% endgeneration %}' markers that "
                "return_assistant_tokens_mask=True relies on.  Install them first "
                "with opake.alignment.data.get_training_chat_template:\n\n"
                "    tokenizer.chat_template = get_training_chat_template(tokenizer)\n",
            )
        )

    result: dict[str, Any] = {
        "input_ids": encoded["input_ids"],
        "completion_mask": assistant_masks,
    }
    attention_mask = encoded.get("attention_mask")
    if attention_mask is not None:
        result["attention_mask"] = attention_mask
    return result

clone_chat_template(model, tokenizer, source_tokenizer_or_path)

Copy a chat template and special tokens from source onto tokenizer.

The function:

1. Loads the source tokenizer (if a path/name string is given, it is loaded via AutoTokenizer.from_pretrained with trust_remote_code=False).
2. Copies source.chat_template onto tokenizer.chat_template.
3. Collects every token that is marked special=True in the source tokenizer's added_tokens_decoder and adds any that are missing from tokenizer via tokenizer.add_special_tokens.
4. Calls model.resize_token_embeddings(len(tokenizer)) to extend (or confirm the size of) the model's input-embedding matrix.

Ordering invariant: this mutates the model's embedding matrix and must be called before any make_functional / torch.func.functional_call snapshot, otherwise the snapshot keeps the old (smaller) matrix and the functional forward errors at runtime (embedding index-out-of-range) when a newly-added token id is first seen.

Parameters:

Name	Type	Description	Default
model	PreTrainedModel	The :class:~transformers.PreTrainedModel whose embedding matrix will be resized to accommodate the cloned special tokens.	required
tokenizer	PreTrainedTokenizerBase	The destination :class:~transformers.PreTrainedTokenizerBase that will receive the chat template and any new special tokens.	required
source_tokenizer_or_path	PreTrainedTokenizerBase | str	Either a loaded :class:~transformers.PreTrainedTokenizerBase whose chat_template and special tokens are to be copied, or a string (local path or Hugging Face Hub model ID) that will be passed to AutoTokenizer.from_pretrained.	required

Returns:

Type	Description
PreTrainedModel	A (model, tokenizer, added_token_ids) tuple. model and
PreTrainedTokenizerBase	tokenizer are mutated in-place and also returned for convenient
list[int]	chaining. added_token_ids is the list of vocabulary indices of the
tuple[PreTrainedModel, PreTrainedTokenizerBase, list[int]]	special tokens newly added to tokenizer (empty when none were added).
tuple[PreTrainedModel, PreTrainedTokenizerBase, list[int]]	Callers fine-tuning with PEFT use it to mark the new embedding rows
tuple[PreTrainedModel, PreTrainedTokenizerBase, list[int]]	trainable, since a frozen base would never learn an embedding for a
tuple[PreTrainedModel, PreTrainedTokenizerBase, list[int]]	token that did not exist at pre-training time.

Raises:

Type	Description
ValueError	If source_tokenizer_or_path does not have a chat_template attribute set (i.e. the attribute is None).
TypeError	If source_tokenizer_or_path is neither a string nor a :class:~transformers.PreTrainedTokenizerBase.

Source code in packages/opake-alignment/src/opake/api/alignment/data/_chat_template.py

def clone_chat_template(
    model: PreTrainedModel,
    tokenizer: PreTrainedTokenizerBase,
    source_tokenizer_or_path: PreTrainedTokenizerBase | str,
) -> tuple[PreTrainedModel, PreTrainedTokenizerBase, list[int]]:
    """Copy a chat template and special tokens from *source* onto *tokenizer*.

    The function:

    1. Loads the *source* tokenizer (if a path/name string is given, it is
       loaded via ``AutoTokenizer.from_pretrained`` with
       ``trust_remote_code=False``).
    2. Copies ``source.chat_template`` onto ``tokenizer.chat_template``.
    3. Collects every token that is marked ``special=True`` in the source
       tokenizer's ``added_tokens_decoder`` and adds any that are **missing**
       from *tokenizer* via ``tokenizer.add_special_tokens``.
    4. Calls ``model.resize_token_embeddings(len(tokenizer))`` to extend (or
       confirm the size of) the model's input-embedding matrix.

    Ordering invariant: this mutates the model's embedding matrix and must be
    called before any ``make_functional`` / ``torch.func.functional_call``
    snapshot, otherwise the snapshot keeps the old (smaller) matrix and the
    functional forward errors at runtime (embedding index-out-of-range) when a
    newly-added token id is first seen.

    Args:
        model: The :class:`~transformers.PreTrainedModel` whose embedding
            matrix will be resized to accommodate the cloned special tokens.
        tokenizer: The destination :class:`~transformers.PreTrainedTokenizerBase`
            that will receive the chat template and any new special tokens.
        source_tokenizer_or_path: Either a loaded
            :class:`~transformers.PreTrainedTokenizerBase` whose ``chat_template``
            and special tokens are to be copied, **or** a string (local path or
            Hugging Face Hub model ID) that will be passed to
            ``AutoTokenizer.from_pretrained``.

    Returns:
        A ``(model, tokenizer, added_token_ids)`` tuple.  ``model`` and
        ``tokenizer`` are mutated in-place and also returned for convenient
        chaining.  ``added_token_ids`` is the list of vocabulary indices of the
        special tokens newly added to *tokenizer* (empty when none were added).
        Callers fine-tuning with PEFT use it to mark the new embedding rows
        trainable, since a frozen base would never learn an embedding for a
        token that did not exist at pre-training time.

    Raises:
        ValueError: If *source_tokenizer_or_path* does not have a
            ``chat_template`` attribute set (i.e. the attribute is ``None``).
        TypeError: If *source_tokenizer_or_path* is neither a string nor a
            :class:`~transformers.PreTrainedTokenizerBase`.
    """
    # Resolve source tokenizer.
    if isinstance(source_tokenizer_or_path, str):
        from transformers import AutoTokenizer  # lazy import

        source_tokenizer = AutoTokenizer.from_pretrained(
            source_tokenizer_or_path,
            trust_remote_code=False,
        )
    else:
        # Duck-type check: must look like a tokenizer.
        if not hasattr(source_tokenizer_or_path, "chat_template"):
            raise InputTypeError(
                *(
                    "source_tokenizer_or_path must be a string path or a "
                    "PreTrainedTokenizerBase instance; "
                    f"got {type(source_tokenizer_or_path)!r}",
                )
            )
        source_tokenizer = source_tokenizer_or_path

    # Copy chat template.
    chat_template = getattr(source_tokenizer, "chat_template", None)
    if chat_template is None:
        raise ConfigurationError(
            *(
                "The source tokenizer does not have a chat_template set "
                "(chat_template is None).  Set it before calling "
                "clone_chat_template.",
            )
        )
    tokenizer.chat_template = chat_template

    # Preserve named roles and AddedToken metadata, rather than treating every
    # special token as an anonymous additional special token.
    _register_model_specific_special_tokens(tokenizer, source_tokenizer)
    source_special_tokens = _source_special_tokens(source_tokenizer)
    existing_vocab = set(tokenizer.get_vocab())
    special_token_strings = {
        str(token)
        for value in source_special_tokens.values()
        for token in (value if isinstance(value, list) else [value])
    }

    added_token_ids: list[int] = []
    if source_special_tokens:
        n_added = tokenizer.add_special_tokens(source_special_tokens)
        # Resolve the vocabulary indices of the tokens that did not exist before
        # cloning so PEFT callers can mark exactly those embedding rows trainable.
        unk_id = getattr(tokenizer, "unk_token_id", None)
        for token_str in special_token_strings - existing_vocab:
            tid = tokenizer.convert_tokens_to_ids(token_str)
            if tid is not None and tid != unk_id:
                added_token_ids.append(int(tid))
        logger.info(
            "clone_chat_template: added %d new special token(s) to tokenizer: %s",
            n_added,
            sorted(special_token_strings - existing_vocab),
        )
    else:
        logger.debug("clone_chat_template: no new special tokens to add.")

    # Resize model embeddings. Called unconditionally so the invariant
    # ``get_input_embeddings().weight.shape[0] == len(tokenizer)`` holds on
    # exit even when no new tokens were added (resize is a no-op then).
    if hasattr(model, "resize_token_embeddings"):
        model.resize_token_embeddings(len(tokenizer))
    else:
        logger.warning(
            "clone_chat_template: model has no resize_token_embeddings method; "
            "skipping embedding resize.  The embedding size may not match the "
            "tokenizer vocabulary size."
        )

    return model, tokenizer, added_token_ids

get_training_chat_template(tokenizer)

Return a chat-template string with assistant-turn generation markers.

Wraps the assistant-turn content of tokenizer's chat template with {% generation %} / {% endgeneration %} Jinja2 markers so that tokenizer.apply_chat_template(..., return_assistant_tokens_mask=True) returns a per-token boolean mask suitable for assistant_only_loss.

The function is idempotent: a template that already contains a Transformers-recognized generation tag is returned unchanged after its rendered spans are validated as assistant-only.

Strategy (close to TRL chat_template_utils.py):

• If the tokenizer's current template already contains a generation marker (including Jinja whitespace-control variants), validate and return it.
• Otherwise, locate the assistant content expression in an explicit assistant branch, a shared Llama/Qwen-style role expression, or a Gemma-style shared render expression that maps the assistant role to model.
• Render each candidate with distinct system, user, and assistant probe text. Accept it only when the generation spans contain assistant content and exclude non-assistant content.
• Raise :class:ValueError when the template cannot be transformed without guessing.

Parameters:

Name	Type	Description	Default
tokenizer	PreTrainedTokenizerBase	A tokenizer whose chat_template attribute must be a non-empty Jinja2 string.	required

Returns:

Type	Description
str	A Jinja2 template string guaranteed to contain
str	{% generation %} ... {% endgeneration %}.

Raises:

Type	Description
ValueError	If tokenizer.chat_template is None or empty, or if its assistant render path cannot be identified and validated.

Source code in packages/opake-alignment/src/opake/api/alignment/data/_chat_template.py

def get_training_chat_template(tokenizer: PreTrainedTokenizerBase) -> str:
    """Return a chat-template string with assistant-turn generation markers.

    Wraps the assistant-turn content of *tokenizer*'s chat template with
    ``{% generation %}`` / ``{% endgeneration %}`` Jinja2 markers so that
    ``tokenizer.apply_chat_template(..., return_assistant_tokens_mask=True)``
    returns a per-token boolean mask suitable for ``assistant_only_loss``.

    The function is **idempotent**: a template that already contains a
    Transformers-recognized generation tag is returned unchanged after its
    rendered spans are validated as assistant-only.

    Strategy (close to TRL ``chat_template_utils.py``):

    - If the tokenizer's current template already contains a generation marker
      (including Jinja whitespace-control variants), validate and return it.
    - Otherwise, locate the assistant content expression in an explicit
      assistant branch, a shared Llama/Qwen-style role expression, or a
      Gemma-style shared render expression that maps the ``assistant`` role to
      ``model``.
    - Render each candidate with distinct system, user, and assistant probe
      text. Accept it only when the generation spans contain assistant content
      and exclude non-assistant content.
    - Raise :class:`ValueError` when the template cannot be transformed without
      guessing.

    Args:
        tokenizer: A tokenizer whose ``chat_template`` attribute must be a
            non-empty Jinja2 string.

    Returns:
        A Jinja2 template string guaranteed to contain
        ``{% generation %}`` ... ``{% endgeneration %}``.

    Raises:
        ValueError: If ``tokenizer.chat_template`` is ``None`` or empty, or if
            its assistant render path cannot be identified and validated.
    """
    if not getattr(tokenizer, "chat_template", None):
        raise ConfigurationError(
            *(
                "tokenizer.chat_template is not set.  Assign a Jinja2 template "
                "string to tokenizer.chat_template before calling "
                "get_training_chat_template.",
            )
        )
    template = _resolve_chat_template(tokenizer)

    # Idempotency: already has the marker.
    if _has_generation_marker(template):
        if _generation_block_marks_only_assistant(template, tokenizer):
            return template
        raise ConfigurationError(
            *(
                "get_training_chat_template: existing generation markers do not "
                "identify a validated assistant-only render path in the tokenizer's "
                "chat template. Ensure they exclude system and user content.",
            )
        )

    # Strategy 1: wrap {{ generation_token }} — TRL v2 canonical placeholder.
    _GENERATION_TOKEN_EXPR = "{{ generation_token }}"
    if _GENERATION_TOKEN_EXPR in template:
        candidate = template.replace(
            _GENERATION_TOKEN_EXPR,
            f"{_GEN_START}{_GENERATION_TOKEN_EXPR}{_GEN_END}",
            1,
        )
        if _generation_block_marks_only_assistant(candidate, tokenizer):
            return candidate

    # Strategy 2: wrap the assistant content expression inside the assistant
    # conditional block. Handles unmarked templates with explicit assistant
    # branches, such as ChatML.
    template_out = _wrap_assistant_content(template)
    if template_out is not None and _generation_block_marks_only_assistant(
        template_out, tokenizer
    ):
        return template_out

    # Strategy 3: Llama 3 constructs every turn in a shared ``content`` variable
    # and renders it outside a role branch.
    template_out = _wrap_shared_content_expression(template)
    if template_out is not None and _generation_block_marks_only_assistant(
        template_out, tokenizer
    ):
        return template_out

    # Strategy 4: Gemma maps assistant to the rendered "model" role, then emits
    # every role through one expression outside the role-selection branch.
    template_out = _wrap_gemma_shared_expression(template)
    if template_out is not None and _generation_block_marks_only_assistant(
        template_out, tokenizer
    ):
        return template_out

    # Strategy 5: shared OR-clause rendering (Qwen2.5-Instruct pattern). When
    # user/system/assistant turns share one expression inside a multi-condition
    # ``or`` clause, the last assistant mention is the tool-call ``elif`` and
    # Strategy 2 would wrap an unreachable expression (empty
    # ``generation_indices``). Rewrite the shared expression with an inner
    # role-guard so only the assistant case renders inside the mask.
    template_out = _wrap_shared_or_clause(template)
    if template_out is not None and _generation_block_marks_only_assistant(
        template_out, tokenizer
    ):
        return template_out

    raise ConfigurationError(
        *(
            "get_training_chat_template: could not identify and validate an "
            "assistant-only render path in the tokenizer's chat template. "
            "Provide a template with explicit '{% generation %}' / "
            "'{% endgeneration %}' markers.",
        )
    )

loss

opake.alignment.dpo.loss — the DPO loss-construction toolkit.

Everything needed to assemble a per-example DPO loss, exposed as direct functions (mapping a config string to one of them is the caller's concern):

• Per-sequence log-probabilities: sequence_logp and its memory-efficient fused drop-in fused_sequence_logp (project hidden states through the lm_head without materialising logits); pass length_normalized=True for the per-token mean reward used by SimPO / ORPO.
• Per-pair heads on log-ratios: sigmoid_loss, hinge_loss, …. Used reference-free by passing the policy logp itself as the log-ratio (CPO is mpo_combine of a reference-free sigmoid_loss and chosen_nll_loss).
• Reference-free heads on (length-normalized) log-probs: simpo_loss (length-normalized sigmoid with a target margin) and odds_ratio_loss (ORPO; takes log-probs directly, not log-ratios).
• Log-ratio combinators for composite objectives: f_divergence_remap / f_divergence_logits, mpo_combine, wpo_weights, ld_dpo_split.

A DPO per_example_loss is head(sequence_logp(...) - ref_logp, …); see examples/train_dpo.py.

Naming is <method>[_<variant>]_loss: the preference/pairwise context is implied by this namespace, and a _<variant> qualifier appears only to disambiguate co-existing variants of one method (apo_zero / apo_down). chosen_nll_loss is the chosen-completion NLL regulariser used in MPO/RPO/CPO blends — not the SFT method's loss (that lives in opake.alignment.sft.loss).

apo_down_loss(chosen_logratio, rejected_logratio, *, beta)

APO-down loss (arXiv:2408.06266).

Pulls the chosen log-ratio down (σ(β·chosen)) while simultaneously enforcing that chosen exceeds rejected by a relative margin (1 − σ(β·Δ)):

loss = σ(β · chosen_lr) + (1 − σ(β · (chosen_lr − rejected_lr)))

Parameters:

Name	Type	Description	Default
chosen_logratio	Tensor	Per-example scalar log π(y_w|x) − log π_ref(y_w|x).	required
rejected_logratio	Tensor	Per-example scalar log π(y_l|x) − log π_ref(y_l|x).	required
beta	float	KL-regularisation temperature (DPO-style).	required

Returns:

Type	Description
Tensor	Per-example scalar loss tensor with the same shape as the inputs.

Source code in packages/opake-alignment/src/opake/api/alignment/dpo/loss/_apo.py

def apo_down_loss(
    chosen_logratio: torch.Tensor,
    rejected_logratio: torch.Tensor,
    *,
    beta: float,
) -> torch.Tensor:
    """APO-down loss (arXiv:2408.06266).

    Pulls the chosen log-ratio down (σ(β·chosen)) while simultaneously
    enforcing that chosen exceeds rejected by a relative margin
    (1 − σ(β·Δ)):

        loss = σ(β · chosen_lr) + (1 − σ(β · (chosen_lr − rejected_lr)))

    Args:
        chosen_logratio: Per-example scalar ``log π(y_w|x) − log π_ref(y_w|x)``.
        rejected_logratio: Per-example scalar ``log π(y_l|x) − log π_ref(y_l|x)``.
        beta: KL-regularisation temperature (DPO-style).

    Returns:
        Per-example scalar loss tensor with the same shape as the inputs.
    """
    losses_chosen = F.sigmoid(beta * chosen_logratio)
    losses_rejected = 1 - F.sigmoid(beta * (chosen_logratio - rejected_logratio))
    return losses_chosen + losses_rejected

apo_zero_loss(chosen_logratio, rejected_logratio, *, beta)

APO-zero loss (arXiv:2408.06266).

Pushes the chosen log-ratio up and the rejected log-ratio down independently of each other. Mathematically:

loss = (1 − σ(β · chosen_lr)) + σ(β · rejected_lr)

where σ is the sigmoid function.

Parameters:

Name	Type	Description	Default
chosen_logratio	Tensor	Per-example scalar log π(y_w|x) − log π_ref(y_w|x).	required
rejected_logratio	Tensor	Per-example scalar log π(y_l|x) − log π_ref(y_l|x).	required
beta	float	KL-regularisation temperature (DPO-style).	required

Returns:

Type	Description
Tensor	Per-example scalar loss tensor with the same shape as the inputs.

Source code in packages/opake-alignment/src/opake/api/alignment/dpo/loss/_apo.py

def apo_zero_loss(
    chosen_logratio: torch.Tensor,
    rejected_logratio: torch.Tensor,
    *,
    beta: float,
) -> torch.Tensor:
    """APO-zero loss (arXiv:2408.06266).

    Pushes the chosen log-ratio up and the rejected log-ratio down
    independently of each other.  Mathematically:

        loss = (1 − σ(β · chosen_lr)) + σ(β · rejected_lr)

    where σ is the sigmoid function.

    Args:
        chosen_logratio: Per-example scalar ``log π(y_w|x) − log π_ref(y_w|x)``.
        rejected_logratio: Per-example scalar ``log π(y_l|x) − log π_ref(y_l|x)``.
        beta: KL-regularisation temperature (DPO-style).

    Returns:
        Per-example scalar loss tensor with the same shape as the inputs.
    """
    losses_chosen = 1 - F.sigmoid(beta * chosen_logratio)
    losses_rejected = F.sigmoid(beta * rejected_logratio)
    return losses_chosen + losses_rejected

bco_loss(chosen_logratio, rejected_logratio, *, beta, delta=0.0)

BCO pairwise loss.

The reward baseline delta corresponds to TRL's running-mean reward estimate. The default 0.0 gives the zero-baseline variant. A nonzero value must be public, separately DP, or derived from prior DP outputs; detaching a raw cross-batch statistic is not sufficient.

Parameters:

Name	Type	Description	Default
chosen_logratio	Tensor	Per-example scalar log π(y_w|x) − log π_ref(y_w|x).	required
rejected_logratio	Tensor	Per-example scalar log π(y_l|x) − log π_ref(y_l|x).	required
beta	float	KL-regularisation temperature (DPO-style).	required
delta	float	Public or DP-safe reward baseline. Default 0.0.	0.0

Returns:

Type	Description
Tensor	Per-example scalar loss tensor with the same shape as the inputs.

Source code in packages/opake-alignment/src/opake/api/alignment/dpo/loss/_bco.py

def bco_loss(
    chosen_logratio: torch.Tensor,
    rejected_logratio: torch.Tensor,
    *,
    beta: float,
    delta: float = 0.0,
) -> torch.Tensor:
    """BCO pairwise loss.

    The reward baseline ``delta`` corresponds to TRL's running-mean reward
    estimate. The default ``0.0`` gives the zero-baseline variant. A nonzero
    value must be public, separately DP, or derived from prior DP outputs;
    detaching a raw cross-batch statistic is not sufficient.

    Args:
        chosen_logratio: Per-example scalar ``log π(y_w|x) − log π_ref(y_w|x)``.
        rejected_logratio: Per-example scalar ``log π(y_l|x) − log π_ref(y_l|x)``.
        beta: KL-regularisation temperature (DPO-style).
        delta: Public or DP-safe reward baseline. Default ``0.0``.

    Returns:
        Per-example scalar loss tensor with the same shape as the inputs.
    """
    return -F.logsigmoid(beta * chosen_logratio - delta) - F.logsigmoid(
        -(beta * rejected_logratio - delta)
    )

chosen_nll_loss(chosen_logp, /, *_args, **_kwargs)

SFT regulariser: negative log-likelihood of the chosen completion.

Returns -chosen_logp, i.e. the NLL of the chosen sequence under the current policy. Extra positional and keyword arguments are silently ignored so this function can be dispatched alongside log-ratio–based variants in the MPO blend loop.

Parameters:

Name	Type	Description	Default
chosen_logp	Tensor	Per-example scalar log-probability log π(y_w | x) (the full completion log-prob, not a log-ratio). May be 0-dim (single example) or (B,) (batch).	required
*_args	object	Ignored. Present so the function can be called with the (chosen_logp, rejected_logratio, ...) signature of the other DPO variants without raising a TypeError.	()
**_kwargs	object	Ignored. beta and other variant kwargs are accepted and discarded.	{}

Returns:

Type	Description
Tensor	Per-example NLL scalar (same shape as chosen_logp).

Source code in packages/opake-alignment/src/opake/api/alignment/dpo/loss/_chosen_nll.py

def chosen_nll_loss(
    chosen_logp: torch.Tensor,
    /,
    *_args: object,
    **_kwargs: object,
) -> torch.Tensor:
    """SFT regulariser: negative log-likelihood of the chosen completion.

    Returns ``-chosen_logp``, i.e. the NLL of the chosen sequence under the
    current policy.  Extra positional and keyword arguments are silently
    ignored so this function can be dispatched alongside log-ratio–based
    variants in the MPO blend loop.

    Args:
        chosen_logp: Per-example scalar log-probability
            ``log π(y_w | x)`` (the full completion log-prob, **not** a
            log-ratio).  May be 0-dim (single example) or ``(B,)`` (batch).
        *_args: Ignored.  Present so the function can be called with the
            ``(chosen_logp, rejected_logratio, ...)`` signature of the other
            DPO variants without raising a ``TypeError``.
        **_kwargs: Ignored.  ``beta`` and other variant kwargs are accepted
            and discarded.

    Returns:
        Per-example NLL scalar (same shape as *chosen_logp*).
    """
    return -chosen_logp

discopop_loss(chosen_logratio, rejected_logratio, *, beta, discopop_tau=0.05)

DiscoPOP per-example loss (Lu et al., 2024, equations 4-5).

Blends a logistic component -log σ(β·Δ) and an exponential component exp(-β·Δ) using a soft gate parameterised by temperature τ::

logits    = (chosen_logratio - rejected_logratio) * beta
gate      = σ(logits / τ)          # → 1 for large positive logits
logistic  = -log σ(logits)
exp_comp  = exp(-logits)            # clamped to avoid overflow
L         = logistic * (1 - gate) + exp_comp * gate

Parameters:

Name	Type	Description	Default
chosen_logratio	Tensor	Per-example scalar log-ratio for the chosen completion. May be 0-dim or (B,).	required
rejected_logratio	Tensor	Per-example scalar log-ratio for the rejected completion. Same shape as chosen_logratio.	required
beta	float	KL-regularisation temperature (DPO β).	required
discopop_tau	float	Modulation temperature τ (default 0.05). Smaller values sharpen the transition between the two components.	0.05

Returns:

Type	Description
Tensor	Per-example scalar loss tensor with the same shape as the inputs.

Note

The exponential component overflows half precision when logits ≪ 0. To prevent NaN/Inf while preserving gradient locality, the logits passed to torch.exp are clamped from below at a dtype-specific bound: 11 for float16 and 80 otherwise. This keeps the exponential finite in the active compute dtype while leaving the representable region unchanged.

Source code in packages/opake-alignment/src/opake/api/alignment/dpo/loss/_discopop.py

def discopop_loss(
    chosen_logratio: torch.Tensor,
    rejected_logratio: torch.Tensor,
    *,
    beta: float,
    discopop_tau: float = 0.05,
) -> torch.Tensor:
    """DiscoPOP per-example loss (Lu et al., 2024, equations 4-5).

    Blends a logistic component ``-log σ(β·Δ)`` and an exponential component
    ``exp(-β·Δ)`` using a soft gate parameterised by temperature *τ*::

        logits    = (chosen_logratio - rejected_logratio) * beta
        gate      = σ(logits / τ)          # → 1 for large positive logits
        logistic  = -log σ(logits)
        exp_comp  = exp(-logits)            # clamped to avoid overflow
        L         = logistic * (1 - gate) + exp_comp * gate

    Args:
        chosen_logratio: Per-example scalar log-ratio for the chosen
            completion.  May be 0-dim or ``(B,)``.
        rejected_logratio: Per-example scalar log-ratio for the rejected
            completion.  Same shape as *chosen_logratio*.
        beta: KL-regularisation temperature (DPO β).
        discopop_tau: Modulation temperature *τ* (default ``0.05``).
            Smaller values sharpen the transition between the two components.

    Returns:
        Per-example scalar loss tensor with the same shape as the inputs.

    Note:
        The exponential component overflows half precision when ``logits ≪ 0``.
        To prevent NaN/Inf while preserving gradient locality, the logits passed
        to ``torch.exp`` are clamped from below at a dtype-specific bound:
        ``11`` for ``float16`` and ``80`` otherwise.  This keeps the
        exponential finite in the active compute dtype while leaving the
        representable region unchanged.
    """
    logits = (chosen_logratio - rejected_logratio) * beta
    exp_clamp = _exp_clamp_for_dtype(logits.dtype)
    logits_clamped = logits.clamp(min=-exp_clamp)
    modulation = torch.sigmoid(logits / discopop_tau)
    logistic_component = -F.logsigmoid(logits)
    exp_component = torch.exp(-logits_clamped)
    return logistic_component * (1.0 - modulation) + exp_component * modulation

exo_loss(chosen_logratio, rejected_logratio, *, beta, label_smoothing=0.001)

EXO pairwise loss (arXiv:2402.00856, simplified pairwise variant).

Parameters:

Name	Type	Description	Default
chosen_logratio	Tensor	Per-example scalar log π(y_w|x) − log π_ref(y_w|x).	required
rejected_logratio	Tensor	Per-example scalar log π(y_l|x) − log π_ref(y_l|x).	required
beta	float	KL-regularisation temperature (DPO-style).	required
label_smoothing	float	Label-smoothing coefficient ε ∈ (0, 0.5). When 0 or negative, silently clamped to 1e-3 (same as TRL) to keep log(ε) finite.	0.001

Returns:

Type	Description
Tensor	Per-example scalar loss tensor with the same shape as the inputs.

Source code in packages/opake-alignment/src/opake/api/alignment/dpo/loss/_exo.py

def exo_loss(
    chosen_logratio: torch.Tensor,
    rejected_logratio: torch.Tensor,
    *,
    beta: float,
    label_smoothing: float = 1e-3,
) -> torch.Tensor:
    """EXO pairwise loss (arXiv:2402.00856, simplified pairwise variant).

    Args:
        chosen_logratio: Per-example scalar ``log π(y_w|x) − log π_ref(y_w|x)``.
        rejected_logratio: Per-example scalar ``log π(y_l|x) − log π_ref(y_l|x)``.
        beta: KL-regularisation temperature (DPO-style).
        label_smoothing: Label-smoothing coefficient ε ∈ (0, 0.5).
            When 0 or negative, silently clamped to ``1e-3`` (same as TRL)
            to keep ``log(ε)`` finite.

    Returns:
        Per-example scalar loss tensor with the same shape as the inputs.
    """
    ls = label_smoothing if label_smoothing > 0 else 1e-3
    logits = beta * (chosen_logratio - rejected_logratio)
    return F.sigmoid(logits) * (F.logsigmoid(logits) - math.log(1 - ls)) + F.sigmoid(
        -logits
    ) * (F.logsigmoid(-logits) - math.log(ls))

f_divergence_logits(chosen_logratio, rejected_logratio, *, f_divergence_type='reverse_kl', alpha=1.0)

Form the remapped preference logits delta under an f-divergence.

Computes::

delta = f_divergence_remap(chosen_logratio) - f_divergence_remap(rejected_logratio)

where both sides are remapped under the same f_divergence_type / alpha. The resulting delta is the generalised drop-in replacement for the plain chosen_logratio - rejected_logratio term consumed by the DPO variants. Under reverse_kl this reduces exactly to that plain difference.

The function is elementwise and works on batched (B,) inputs as well as 0-dim scalars.

Parameters:

Name	Type	Description	Default
chosen_logratio	Tensor	Per-example log-ratio for the chosen response. Shape (B,) or ().	required
rejected_logratio	Tensor	Per-example log-ratio for the rejected response. Same shape as chosen_logratio.	required
f_divergence_type	FDivergence	Which f-divergence remap to apply to both sides. Defaults to "reverse_kl".	'reverse_kl'
alpha	float	The α parameter for alpha_divergence. When alpha == 1 the function falls through to reverse_kl. Ignored by the other divergences. Defaults to 1.0.	1.0

Returns:

Type	Description
Tensor	Remapped preference-logit tensor of the same shape as the inputs.

Raises:

Type	Description
ValueError	Propagated from :func:f_divergence_remap for an unknown f_divergence_type.

Source code in packages/opake-alignment/src/opake/api/alignment/dpo/loss/_f_divergence.py

def f_divergence_logits(
    chosen_logratio: torch.Tensor,
    rejected_logratio: torch.Tensor,
    *,
    f_divergence_type: FDivergence = "reverse_kl",
    alpha: float = 1.0,
) -> torch.Tensor:
    """Form the remapped preference logits ``delta`` under an f-divergence.

    Computes::

        delta = f_divergence_remap(chosen_logratio) - f_divergence_remap(rejected_logratio)

    where both sides are remapped under the same ``f_divergence_type`` /
    ``alpha``. The resulting ``delta`` is the generalised drop-in replacement for
    the plain ``chosen_logratio - rejected_logratio`` term consumed by the DPO
    variants. Under ``reverse_kl`` this reduces exactly to that plain difference.

    The function is elementwise and works on batched ``(B,)`` inputs as well as
    0-dim scalars.

    Args:
        chosen_logratio: Per-example log-ratio for the chosen response. Shape
            ``(B,)`` or ``()``.
        rejected_logratio: Per-example log-ratio for the rejected response. Same
            shape as ``chosen_logratio``.
        f_divergence_type: Which f-divergence remap to apply to both sides.
            Defaults to ``"reverse_kl"``.
        alpha: The α parameter for ``alpha_divergence``. When ``alpha == 1``
            the function falls through to ``reverse_kl``. Ignored by the
            other divergences. Defaults to ``1.0``.

    Returns:
        Remapped preference-logit tensor of the same shape as the inputs.

    Raises:
        ValueError: Propagated from :func:`f_divergence_remap` for an unknown
            ``f_divergence_type``.
    """
    chosen = f_divergence_remap(
        chosen_logratio, f_divergence_type=f_divergence_type, alpha=alpha
    )
    rejected = f_divergence_remap(
        rejected_logratio, f_divergence_type=f_divergence_type, alpha=alpha
    )
    return chosen - rejected

f_divergence_remap(logratio, *, f_divergence_type='reverse_kl', alpha=1.0)

Remap a per-example log-ratio under the chosen f-divergence.

Applies the per-side remap g for the requested f-divergence:

• reverse_kl: g(x) = x (identity — grad-transparent).
• forward_kl: g(x) = 1 - _cap_exp(-x).
• js_divergence: g(x) = log(2) + logsigmoid(x).
• alpha_divergence: g(x) = (_cap_exp((alpha - 1) * x) - 1) / (alpha - 1); near alpha == 1 this uses the identity (reverse_kl).

The function is elementwise and works on batched (B,) inputs as well as 0-dim scalars.

Parameters:

Name	Type	Description	Default
logratio	Tensor	Per-example log-ratio log π(y | x) − log π_ref(y | x). Shape (B,) or ().	required
f_divergence_type	FDivergence	Which f-divergence remap to apply. One of "reverse_kl", "forward_kl", "js_divergence", "alpha_divergence". Defaults to "reverse_kl".	'reverse_kl'
alpha	float	The α parameter for alpha_divergence. When alpha == 1 the function returns the identity (matching TRL's fallback). Ignored by the other divergences. Defaults to 1.0.	1.0

Returns:

Type	Description
Tensor	Remapped log-ratio tensor of the same shape as logratio.

Raises:

Type	Description
ValueError	If f_divergence_type is not a recognised divergence name.

Source code in packages/opake-alignment/src/opake/api/alignment/dpo/loss/_f_divergence.py

def f_divergence_remap(
    logratio: torch.Tensor,
    *,
    f_divergence_type: FDivergence = "reverse_kl",
    alpha: float = 1.0,
) -> torch.Tensor:
    """Remap a per-example log-ratio under the chosen f-divergence.

    Applies the per-side remap ``g`` for the requested f-divergence:

    - ``reverse_kl``: ``g(x) = x`` (identity — grad-transparent).
    - ``forward_kl``: ``g(x) = 1 - _cap_exp(-x)``.
    - ``js_divergence``: ``g(x) = log(2) + logsigmoid(x)``.
    - ``alpha_divergence``:
      ``g(x) = (_cap_exp((alpha - 1) * x) - 1) / (alpha - 1)``;
      near ``alpha == 1`` this uses the identity (``reverse_kl``).

    The function is elementwise and works on batched ``(B,)`` inputs as well as
    0-dim scalars.

    Args:
        logratio: Per-example log-ratio ``log π(y | x) − log π_ref(y | x)``.
            Shape ``(B,)`` or ``()``.
        f_divergence_type: Which f-divergence remap to apply. One of
            ``"reverse_kl"``, ``"forward_kl"``, ``"js_divergence"``,
            ``"alpha_divergence"``. Defaults to ``"reverse_kl"``.
        alpha: The α parameter for ``alpha_divergence``. When ``alpha == 1``
            the function returns the identity (matching TRL's fallback).
            Ignored by the other divergences. Defaults to ``1.0``.

    Returns:
        Remapped log-ratio tensor of the same shape as ``logratio``.

    Raises:
        ValueError: If ``f_divergence_type`` is not a recognised divergence name.
    """
    if f_divergence_type == "reverse_kl":
        # Identity: returned unchanged so the gradient path is fully
        # transparent (d g/d x = 1), reproducing the standard DPO objective.
        return logratio
    if f_divergence_type == "forward_kl":
        return 1 - _cap_exp(-logratio)
    if f_divergence_type == "js_divergence":
        return math.log(2.0) + F.logsigmoid(logratio)
    if f_divergence_type == "alpha_divergence":
        if abs(alpha - 1.0) < _ALPHA_DIVERGENCE_LIMIT_TOLERANCE:
            # Limit of (exp((alpha-1)*x) - 1)/(alpha-1) as alpha->1.
            return logratio
        return (_cap_exp((alpha - 1) * logratio) - 1) / (alpha - 1)
    raise ConfigurationError(
        *(
            f"Unknown f_divergence_type {f_divergence_type!r}; expected one of "
            '"reverse_kl", "forward_kl", "js_divergence", "alpha_divergence".',
        )
    )

hinge_loss(chosen_logratio, rejected_logratio, *, beta)

Per-example DPO hinge loss.

Implements the hinge variant of the DPO objective from Liu et al. 2023. The function is elementwise and works on both batched (B,) inputs and 0-dim scalars.

The hinge loss is zero when the margin β·Δ ≥ 1 (chosen is clearly preferred over rejected by the temperature-scaled margin). It is linear in the deficit below 1 otherwise.

Parameters:

Name	Type	Description	Default
chosen_logratio	Tensor	Per-example log-ratio for the chosen response, log π(y_w | x) − log π_ref(y_w | x). Shape (B,) or ().	required
rejected_logratio	Tensor	Per-example log-ratio for the rejected response, log π(y_l | x) − log π_ref(y_l | x). Same shape as chosen_logratio.	required
beta	float	DPO temperature (reference-deviation strength). Positive float; typical values in [0.01, 0.5].	required

Returns:

Type	Description
Tensor	Per-example loss tensor of the same shape as the inputs.

Source code in packages/opake-alignment/src/opake/api/alignment/dpo/loss/_hinge.py

def hinge_loss(
    chosen_logratio: torch.Tensor,
    rejected_logratio: torch.Tensor,
    *,
    beta: float,
) -> torch.Tensor:
    """Per-example DPO hinge loss.

    Implements the hinge variant of the DPO objective from Liu et al. 2023.
    The function is elementwise and works on both batched ``(B,)`` inputs and
    0-dim scalars.

    The hinge loss is zero when the margin β·Δ ≥ 1 (chosen is clearly
    preferred over rejected by the temperature-scaled margin). It is linear
    in the deficit below 1 otherwise.

    Args:
        chosen_logratio: Per-example log-ratio for the chosen response,
            ``log π(y_w | x) − log π_ref(y_w | x)``. Shape ``(B,)`` or ``()``.
        rejected_logratio: Per-example log-ratio for the rejected response,
            ``log π(y_l | x) − log π_ref(y_l | x)``. Same shape as
            ``chosen_logratio``.
        beta: DPO temperature (reference-deviation strength). Positive float;
            typical values in [0.01, 0.5].

    Returns:
        Per-example loss tensor of the same shape as the inputs.
    """
    logits = chosen_logratio - rejected_logratio
    return torch.relu(1 - beta * logits)

ipo_loss(chosen_logratio, rejected_logratio, *, beta)

IPO per-example loss (Azar 2024).

Computes the squared deviation between the log-ratio difference and the inverse of twice the KL-regularisation temperature::

L = (Δ - 1/(2β))²    where Δ = chosen_logratio - rejected_logratio

When the caller passes length-normalised log-ratios (i.e. each log-ratio divided by the corresponding completion token count), this recovers the original IPO objective from the paper.

Parameters:

Name	Type	Description	Default
chosen_logratio	Tensor	Per-example scalar log π(y_w) - log π_ref(y_w). May be 0-dim (single example) or 1-dim (B,) (batch).	required
rejected_logratio	Tensor	Per-example scalar log π(y_l) - log π_ref(y_l). Same shape as chosen_logratio.	required
beta	float	KL-regularisation temperature (DPO β). Controls the strength of the preference signal; larger β → the model must diverge more from the reference to achieve the same reward margin.	required

Returns:

Type	Description
Tensor	Per-example scalar loss tensor with the same shape as the inputs.

Source code in packages/opake-alignment/src/opake/api/alignment/dpo/loss/_ipo.py

def ipo_loss(
    chosen_logratio: torch.Tensor,
    rejected_logratio: torch.Tensor,
    *,
    beta: float,
) -> torch.Tensor:
    """IPO per-example loss (Azar 2024).

    Computes the squared deviation between the log-ratio difference and the
    inverse of twice the KL-regularisation temperature::

        L = (Δ - 1/(2β))²    where Δ = chosen_logratio - rejected_logratio

    When the caller passes length-normalised log-ratios (i.e. each log-ratio
    divided by the corresponding completion token count), this recovers the
    original IPO objective from the paper.

    Args:
        chosen_logratio: Per-example scalar ``log π(y_w) - log π_ref(y_w)``.
            May be 0-dim (single example) or 1-dim ``(B,)`` (batch).
        rejected_logratio: Per-example scalar ``log π(y_l) - log π_ref(y_l)``.
            Same shape as *chosen_logratio*.
        beta: KL-regularisation temperature (DPO β).  Controls the strength
            of the preference signal; larger β → the model must diverge more
            from the reference to achieve the same reward margin.

    Returns:
        Per-example scalar loss tensor with the same shape as the inputs.
    """
    logits = chosen_logratio - rejected_logratio
    return (logits - 1.0 / (2.0 * beta)).pow(2)

ld_dpo_split(per_token_logps, completion_mask, shared_prefix_len, alpha)

Apply the LD-DPO per-sequence weighting for a supplied public length.

Weights completion tokens within the public length by 1.0 and tokens beyond it by alpha, then masked-sums over the sequence axis::

pos    = completion_mask.cumsum(-1)
weight = where((pos >= 1) & (pos <= shared_prefix_len), 1.0, alpha)
return (per_token_logps * completion_mask * weight).sum(-1)

Parameters:

Name	Type	Description	Default
per_token_logps	Tensor	Per-token log-probabilities of the completion. Shape (..., T); leading dims may be empty (per-example under vmap) or a batch axis.	required
completion_mask	Tensor	Tensor of shape (..., T); non-zero where a token belongs to the completion span. Cast to the logp dtype.	required
shared_prefix_len	Tensor | int	The paper's public length (the shorter completion length), retained under this compatibility name. A Python int (same value for every example) or a per-example tensor that broadcasts against the position axis (T,) (e.g. shape (..., 1)). The first shared_prefix_len completion tokens are weighted 1.0; later completion tokens are weighted alpha.	required
alpha	float	Weight on the length-desensitised tail in [0, 1]. alpha = 1 recovers the plain masked-sum; alpha = 0 keeps only tokens within the public length.	required

Returns:

Type	Description
Tensor	Float tensor of shape (...) (one length-desensitised logp per
Tensor	sequence).

Note

This low-level, single-sequence primitive cannot derive the public length from a chosen/rejected pair. Callers claiming the LD-DPO objective must pass the minimum of the paired completion lengths.

Source code in packages/opake-alignment/src/opake/api/alignment/dpo/loss/_ld_dpo.py

def ld_dpo_split(
    per_token_logps: torch.Tensor,
    completion_mask: torch.Tensor,
    shared_prefix_len: torch.Tensor | int,
    alpha: float,
) -> torch.Tensor:
    """Apply the LD-DPO per-sequence weighting for a supplied public length.

    Weights completion tokens within the public length by ``1.0`` and tokens
    beyond it by ``alpha``, then masked-sums over the sequence
    axis::

        pos    = completion_mask.cumsum(-1)
        weight = where((pos >= 1) & (pos <= shared_prefix_len), 1.0, alpha)
        return (per_token_logps * completion_mask * weight).sum(-1)

    Args:
        per_token_logps: Per-token log-probabilities of the completion. Shape
            ``(..., T)``; leading dims may be empty (per-example under
            ``vmap``) or a batch axis.
        completion_mask: Tensor of shape ``(..., T)``; non-zero where a token
            belongs to the completion span. Cast to the logp dtype.
        shared_prefix_len: The paper's public length (the shorter completion
            length), retained under this compatibility name. A Python ``int``
            (same value for every example) or a per-example tensor that
            broadcasts against the position axis ``(T,)`` (e.g. shape
            ``(..., 1)``). The first ``shared_prefix_len`` completion tokens are
            weighted ``1.0``; later completion tokens are weighted ``alpha``.
        alpha: Weight on the length-desensitised tail in ``[0, 1]``.
            ``alpha = 1`` recovers the plain masked-sum; ``alpha = 0`` keeps
            only tokens within the public length.

    Returns:
        Float tensor of shape ``(...)`` (one length-desensitised logp per
        sequence).

    Note:
        This low-level, single-sequence primitive cannot derive the public
        length from a chosen/rejected pair. Callers claiming the LD-DPO
        objective must pass the minimum of the paired completion lengths.
    """
    if not math.isfinite(alpha) or not 0.0 <= alpha <= 1.0:
        raise ConfigurationError(
            *(f"alpha (LD-DPO) must be finite and in [0, 1], got {alpha}",)
        )

    mask = completion_mask.to(per_token_logps.dtype)
    completion_pos = (completion_mask != 0).cumsum(dim=-1)
    prefix_len = shared_prefix_len
    if isinstance(prefix_len, torch.Tensor):
        while prefix_len.ndim < completion_pos.ndim:
            prefix_len = prefix_len.unsqueeze(-1)
    is_prefix = (completion_pos >= 1) & (completion_pos <= prefix_len)
    weight = torch.where(
        is_prefix,
        torch.ones((), dtype=per_token_logps.dtype, device=per_token_logps.device),
        torch.full(
            (), alpha, dtype=per_token_logps.dtype, device=per_token_logps.device
        ),
    )
    masked = per_token_logps * mask * weight
    return masked.sum(dim=-1)

mpo_combine(losses, weights)

Weighted sum of per-example loss tensors (MPO / TRL loss_type=list).

Computes sum_k weights[k] * losses[k] over the keys of weights. Only the terms named in weights contribute, so a caller may pass a superset of available losses in losses and select a subset via weights.

Parameters:

Name	Type	Description	Default
losses	dict[str, Tensor]	Mapping from loss name to its per-example loss tensor. All selected tensors must broadcast against one another.	required
weights	dict[str, float]	Mapping from loss name to its scalar blend weight. Its keys must be a subset of losses's keys.	required

Returns:

Type	Description
Tensor	The per-example blended loss tensor, with shape equal to the
Tensor	broadcast of the selected loss tensors.

Raises:

Type	Description
ConfigurationError	If weights names a key that is absent from losses.

Source code in packages/opake-alignment/src/opake/api/alignment/dpo/loss/_mpo.py

def mpo_combine(
    losses: dict[str, torch.Tensor],
    weights: dict[str, float],
) -> torch.Tensor:
    """Weighted sum of per-example loss tensors (MPO / TRL ``loss_type=list``).

    Computes ``sum_k weights[k] * losses[k]`` over the keys of *weights*.
    Only the terms named in *weights* contribute, so a caller may pass a
    superset of available losses in *losses* and select a subset via
    *weights*.

    Args:
        losses: Mapping from loss name to its per-example loss tensor. All
            selected tensors must broadcast against one another.
        weights: Mapping from loss name to its scalar blend weight. Its keys
            must be a subset of ``losses``'s keys.

    Returns:
        The per-example blended loss tensor, with shape equal to the
        broadcast of the selected loss tensors.

    Raises:
        ConfigurationError: If *weights* names a key that is absent from *losses*.
    """
    missing = weights.keys() - losses.keys()
    if missing:
        raise ConfigurationError(
            *(
                f"weights keys {sorted(missing)} are not present in losses "
                f"(available: {sorted(losses)})",
            )
        )

    out: torch.Tensor | None = None
    for name, weight in weights.items():
        term = weight * losses[name]
        out = term if out is None else out + term
    if out is None:
        raise ConfigurationError(*("weights must select at least one loss term",))
    return out

nca_loss(chosen_logratio, rejected_logratio, *, beta)

NCA pairwise loss.

Parameters:

Name	Type	Description	Default
chosen_logratio	Tensor	Per-example scalar log π(y_w|x) − log π_ref(y_w|x).	required
rejected_logratio	Tensor	Per-example scalar log π(y_l|x) − log π_ref(y_l|x).	required
beta	float	KL-regularisation temperature (DPO-style).	required

Returns:

Type	Description
Tensor	Per-example scalar loss tensor with the same shape as the inputs.

Source code in packages/opake-alignment/src/opake/api/alignment/dpo/loss/_nca.py

def nca_loss(
    chosen_logratio: torch.Tensor,
    rejected_logratio: torch.Tensor,
    *,
    beta: float,
) -> torch.Tensor:
    """NCA pairwise loss.

    Args:
        chosen_logratio: Per-example scalar ``log π(y_w|x) − log π_ref(y_w|x)``.
        rejected_logratio: Per-example scalar ``log π(y_l|x) − log π_ref(y_l|x)``.
        beta: KL-regularisation temperature (DPO-style).

    Returns:
        Per-example scalar loss tensor with the same shape as the inputs.
    """
    cr = beta * chosen_logratio
    rr = beta * rejected_logratio
    return -F.logsigmoid(cr) - 0.5 * F.logsigmoid(-cr) - 0.5 * F.logsigmoid(-rr)

odds_ratio_loss(chosen_logp, rejected_logp)

Per-example ORPO odds-ratio loss.

Computes -log σ(log odds(y_w) - log odds(y_r)) with log odds(y) = log p(y) - log(1 - p(y)).

Parameters:

Name	Type	Description	Default
chosen_logp	Tensor	Length-normalized log-probability log p(y_w) of the chosen completion (< 0). Shape (B,) or (). Form it with sequence_logp(..., length_normalized=True).	required
rejected_logp	Tensor	Same for the rejected completion.	required

Returns:

Type	Description
Tensor	Per-example scalar odds-ratio loss, same shape as the inputs.

Source code in packages/opake-alignment/src/opake/api/alignment/dpo/loss/_orpo.py

def odds_ratio_loss(
    chosen_logp: torch.Tensor,
    rejected_logp: torch.Tensor,
) -> torch.Tensor:
    """Per-example ORPO odds-ratio loss.

    Computes ``-log σ(log odds(y_w) - log odds(y_r))`` with
    ``log odds(y) = log p(y) - log(1 - p(y))``.

    Args:
        chosen_logp: Length-normalized log-probability ``log p(y_w)`` of the
            chosen completion (``< 0``). Shape ``(B,)`` or ``()``. Form it with
            ``sequence_logp(..., length_normalized=True)``.
        rejected_logp: Same for the rejected completion.

    Returns:
        Per-example scalar odds-ratio loss, same shape as the inputs.
    """
    log_odds = (chosen_logp - _log1mexp(chosen_logp)) - (
        rejected_logp - _log1mexp(rejected_logp)
    )
    return -F.logsigmoid(log_odds)

robust_loss(chosen_logratio, rejected_logratio, *, beta, label_smoothing=0.0)

Per-example DPO robust (label-smoothed) loss.

Implements the robust label-smoothed DPO objective. The function is elementwise and works on both batched (B,) inputs and 0-dim scalars.

Unlike the sigmoid variant's label smoothing (which simply blends the reversed loss in), the robust variant normalises the blended loss by (1 − 2·label_smoothing) so that the gradient scale is preserved across smoothing levels.

Parameters:

Name	Type	Description	Default
chosen_logratio	Tensor	Per-example log-ratio for the chosen response, log π(y_w | x) − log π_ref(y_w | x). Shape (B,) or ().	required
rejected_logratio	Tensor	Per-example log-ratio for the rejected response, log π(y_l | x) − log π_ref(y_l | x). Same shape as chosen_logratio.	required
beta	float	DPO temperature (reference-deviation strength). Positive float; typical values in [0.01, 0.5].	required
label_smoothing	float	Label-smoothing coefficient ε ∈ [0, 0.5). The denominator is (1 − 2·ε); passing ε ≥ 0.5 causes division by zero and is a caller error. Defaults to 0.0.	0.0

Returns:

Type	Description
Tensor	Per-example loss tensor of the same shape as the inputs.

Source code in packages/opake-alignment/src/opake/api/alignment/dpo/loss/_robust.py

def robust_loss(
    chosen_logratio: torch.Tensor,
    rejected_logratio: torch.Tensor,
    *,
    beta: float,
    label_smoothing: float = 0.0,
) -> torch.Tensor:
    """Per-example DPO robust (label-smoothed) loss.

    Implements the robust label-smoothed DPO objective. The function is
    elementwise and works on both batched ``(B,)`` inputs and 0-dim scalars.

    Unlike the sigmoid variant's label smoothing (which simply blends the
    reversed loss in), the robust variant normalises the blended loss by
    ``(1 − 2·label_smoothing)`` so that the gradient scale is preserved
    across smoothing levels.

    Args:
        chosen_logratio: Per-example log-ratio for the chosen response,
            ``log π(y_w | x) − log π_ref(y_w | x)``. Shape ``(B,)`` or ``()``.
        rejected_logratio: Per-example log-ratio for the rejected response,
            ``log π(y_l | x) − log π_ref(y_l | x)``. Same shape as
            ``chosen_logratio``.
        beta: DPO temperature (reference-deviation strength). Positive float;
            typical values in [0.01, 0.5].
        label_smoothing: Label-smoothing coefficient ε ∈ [0, 0.5). The
            denominator is ``(1 − 2·ε)``; passing ε ≥ 0.5 causes division by
            zero and is a caller error. Defaults to ``0.0``.

    Returns:
        Per-example loss tensor of the same shape as the inputs.
    """
    logits = chosen_logratio - rejected_logratio
    return (
        -F.logsigmoid(beta * logits) * (1 - label_smoothing)
        + F.logsigmoid(-beta * logits) * label_smoothing
    ) / (1 - 2 * label_smoothing)

sigmoid_loss(chosen_logratio, rejected_logratio, *, beta, label_smoothing=0.0)

Per-example DPO sigmoid (logistic) loss.

Implements the original DPO objective from Rafailov et al. 2023 with optional label smoothing. The function is elementwise and works on both batched (B,) inputs and 0-dim scalars.

Parameters:

Name	Type	Description	Default
chosen_logratio	Tensor	Per-example log-ratio for the chosen response, log π(y_w | x) − log π_ref(y_w | x). Shape (B,) or ().	required
rejected_logratio	Tensor	Per-example log-ratio for the rejected response, log π(y_l | x) − log π_ref(y_l | x). Same shape as chosen_logratio.	required
beta	float	DPO temperature (reference-deviation strength). Positive float; typical values in [0.01, 0.5].	required
label_smoothing	float	Label-smoothing coefficient ε ∈ [0, 1). At ε = 0 this reduces to the unsmoothed DPO loss. Defaults to 0.0.	0.0

Returns:

Type	Description
Tensor	Per-example loss tensor of the same shape as the inputs.

Source code in packages/opake-alignment/src/opake/api/alignment/dpo/loss/_sigmoid.py

def sigmoid_loss(
    chosen_logratio: torch.Tensor,
    rejected_logratio: torch.Tensor,
    *,
    beta: float,
    label_smoothing: float = 0.0,
) -> torch.Tensor:
    """Per-example DPO sigmoid (logistic) loss.

    Implements the original DPO objective from Rafailov et al. 2023 with
    optional label smoothing. The function is elementwise and works on both
    batched ``(B,)`` inputs and 0-dim scalars.

    Args:
        chosen_logratio: Per-example log-ratio for the chosen response,
            ``log π(y_w | x) − log π_ref(y_w | x)``. Shape ``(B,)`` or ``()``.
        rejected_logratio: Per-example log-ratio for the rejected response,
            ``log π(y_l | x) − log π_ref(y_l | x)``. Same shape as
            ``chosen_logratio``.
        beta: DPO temperature (reference-deviation strength). Positive float;
            typical values in [0.01, 0.5].
        label_smoothing: Label-smoothing coefficient ε ∈ [0, 1). At ε = 0 this
            reduces to the unsmoothed DPO loss. Defaults to ``0.0``.

    Returns:
        Per-example loss tensor of the same shape as the inputs.
    """
    logits = chosen_logratio - rejected_logratio
    return (
        -F.logsigmoid(beta * logits) * (1 - label_smoothing)
        - F.logsigmoid(-beta * logits) * label_smoothing
    )

simpo_loss(chosen_logratio, rejected_logratio, *, beta, gamma=0.0, label_smoothing=0.0)

Per-example SimPO loss.

Computes L = -log σ(m)·(1 - ε) - log σ(-m)·ε with m = β·(chosen_logratio - rejected_logratio) - γ. At gamma=0 and label_smoothing=0 this is the length-normalized DPO sigmoid loss.

Parameters:

Name	Type	Description	Default
chosen_logratio	Tensor	Per-example length-normalized reward for the chosen completion — the policy log-prob divided by the completion length. Reference-free, so no reference is subtracted. Shape (B,) or ().	required
rejected_logratio	Tensor	Same for the rejected completion.	required
beta	float	Reward scale (SimPO β).	required
gamma	float	Target reward margin subtracted inside the sigmoid; 0.0 (default) recovers the marginless length-normalized sigmoid.	0.0
label_smoothing	float	Conservative label-smoothing coefficient ε ∈ [0, 0.5).	0.0

Returns:

Type	Description
Tensor	Per-example scalar loss, same shape as the inputs.

Source code in packages/opake-alignment/src/opake/api/alignment/dpo/loss/_simpo.py

def simpo_loss(
    chosen_logratio: torch.Tensor,
    rejected_logratio: torch.Tensor,
    *,
    beta: float,
    gamma: float = 0.0,
    label_smoothing: float = 0.0,
) -> torch.Tensor:
    """Per-example SimPO loss.

    Computes ``L = -log σ(m)·(1 - ε) - log σ(-m)·ε`` with
    ``m = β·(chosen_logratio - rejected_logratio) - γ``. At ``gamma=0`` and
    ``label_smoothing=0`` this is the length-normalized DPO sigmoid loss.

    Args:
        chosen_logratio: Per-example length-normalized reward for the chosen
            completion — the policy log-prob divided by the completion length.
            Reference-free, so no reference is subtracted. Shape ``(B,)`` or
            ``()``.
        rejected_logratio: Same for the rejected completion.
        beta: Reward scale (SimPO β).
        gamma: Target reward margin subtracted inside the sigmoid; ``0.0``
            (default) recovers the marginless length-normalized sigmoid.
        label_smoothing: Conservative label-smoothing coefficient ε ∈ [0, 0.5).

    Returns:
        Per-example scalar loss, same shape as the inputs.
    """
    margin = beta * (chosen_logratio - rejected_logratio) - gamma
    return (
        -F.logsigmoid(margin) * (1.0 - label_smoothing)
        - F.logsigmoid(-margin) * label_smoothing
    )

sppo_loss(chosen_logratio, rejected_logratio, *, beta)

SPPO hard-label loss.

Parameters:

Name	Type	Description	Default
chosen_logratio	Tensor	Per-example scalar log π(y_w|x) − log π_ref(y_w|x).	required
rejected_logratio	Tensor	Per-example scalar log π(y_l|x) − log π_ref(y_l|x).	required
beta	float	KL-regularisation temperature (DPO-style). The Nash-equilibrium target in log-ratio space is 0.5 / beta; beta must be non-zero.	required

Returns:

Type	Description
Tensor	Per-example scalar loss tensor with the same shape as the inputs.

Source code in packages/opake-alignment/src/opake/api/alignment/dpo/loss/_sppo.py

def sppo_loss(
    chosen_logratio: torch.Tensor,
    rejected_logratio: torch.Tensor,
    *,
    beta: float,
) -> torch.Tensor:
    """SPPO hard-label loss.

    Args:
        chosen_logratio: Per-example scalar ``log π(y_w|x) − log π_ref(y_w|x)``.
        rejected_logratio: Per-example scalar ``log π(y_l|x) − log π_ref(y_l|x)``.
        beta: KL-regularisation temperature (DPO-style).  The Nash-equilibrium
            target in log-ratio space is ``0.5 / beta``; ``beta`` must be
            non-zero.

    Returns:
        Per-example scalar loss tensor with the same shape as the inputs.
    """
    target = 0.5 / beta
    return (chosen_logratio - target).pow(2) + (rejected_logratio + target).pow(2)

wpo_weights(per_token_logps, completion_mask, logits=None)

Per-example WPO weight (arXiv:2406.11827).

When logits are provided, computes the weight-aligned policy probability from Equation 2 of WPO and returns its detached geometric mean::

log_denom = logsumexp(2 * logits, -1) - 2 * logsumexp(logits, -1)
aligned_logp = per_token_logps - log_denom
return exp(sum(masked aligned_logp) / completion_token_count)

Omitting logits retains the unaligned geometric-mean weighting for backward compatibility. The result is detached: it contributes no gradient and serves only as a per-example multiplicative reweighting of the downstream loss.

Parameters:

Name	Type	Description	Default
per_token_logps	Tensor	Per-token log-probabilities of the completion under the current policy. Shape (..., T); the leading dims may be empty (per-example under vmap) or a batch axis.	required
completion_mask	Tensor	Tensor of shape (..., T); non-zero where a token belongs to the completion span. Cast to the logp dtype before multiplying.	required
logits	Tensor | None	Optional policy logits of shape (..., T, V) corresponding to per_token_logps. When present, the WPO weight-alignment term is computed over the vocabulary dimension.	None

Returns:

Type	Description
Tensor	Detached per-example weight tensor of shape (...) (one weight per
Tensor	sequence). All-zero mask rows are protected by clamp(min=1) so
Tensor	there is no division by zero.

Source code in packages/opake-alignment/src/opake/api/alignment/dpo/loss/_wpo.py

def wpo_weights(
    per_token_logps: torch.Tensor,
    completion_mask: torch.Tensor,
    logits: torch.Tensor | None = None,
) -> torch.Tensor:
    """Per-example WPO weight (arXiv:2406.11827).

    When ``logits`` are provided, computes the weight-aligned policy probability
    from Equation 2 of WPO and returns its detached geometric mean::

        log_denom = logsumexp(2 * logits, -1) - 2 * logsumexp(logits, -1)
        aligned_logp = per_token_logps - log_denom
        return exp(sum(masked aligned_logp) / completion_token_count)

    Omitting ``logits`` retains the unaligned geometric-mean weighting for
    backward compatibility. The result is **detached**: it contributes no
    gradient and serves only as a per-example multiplicative reweighting of the
    downstream loss.

    Args:
        per_token_logps: Per-token log-probabilities of the completion under
            the current policy. Shape ``(..., T)``; the leading dims may be
            empty (per-example under ``vmap``) or a batch axis.
        completion_mask: Tensor of shape ``(..., T)``; non-zero where a token
            belongs to the completion span. Cast to the logp dtype before
            multiplying.
        logits: Optional policy logits of shape ``(..., T, V)`` corresponding
            to ``per_token_logps``. When present, the WPO weight-alignment term
            is computed over the vocabulary dimension.

    Returns:
        Detached per-example weight tensor of shape ``(...)`` (one weight per
        sequence). All-zero mask rows are protected by ``clamp(min=1)`` so
        there is no division by zero.
    """
    acc_dtype = _compute_dtype(per_token_logps)
    if logits is not None:
        acc_dtype = torch.promote_types(acc_dtype, _compute_dtype(logits))

    mask = completion_mask.to(torch.bool)
    token_count = completion_mask.to(torch.bool).sum(dim=-1).clamp(min=1)
    with torch.no_grad():
        aligned_logps = per_token_logps.to(acc_dtype)
        if logits is not None:
            compute_logits = logits.to(acc_dtype)
            log_denom = torch.logsumexp(
                2.0 * compute_logits, dim=-1
            ) - 2.0 * torch.logsumexp(compute_logits, dim=-1)
            aligned_logps = aligned_logps - log_denom
        masked_logps = torch.where(mask, aligned_logps, 0.0)
        mean_logps = masked_logps.sum(dim=-1) / token_count.to(acc_dtype)
        weights = mean_logps.exp()
    return weights

fused_sequence_logp(hidden_states, lm_head_weight, input_ids, completion_mask, *, length_normalized=False)

Memory-efficient :func:sequence_logp over hidden states (fused, per-example).

Mathematically sequence_logp(hidden_states @ lm_head_weight.T, input_ids, completion_mask), but on CUDA + half precision with opake-alignment[patches] installed it computes the completion log-prob via the patches fused linear-CE kernel — no (T, V) logits are materialised and the LSE is recomputed in the backward. Otherwise it falls back to the eager form.

Since sequence_logp = Σ_completion log p = −Σ_completion CE, the kernel is the preference-logp kernel: encode the completion span as the kept labels (everything else → -100), take the kernel's CE sum, and negate.

Per-example. Pass one sequence (T, H) and drive with vmap(grad(...)) (how train_dpo already computes logp inside clipped_grad); the kernel's merged vmap rules then make the chosen / rejected microbatch one forward + one backward kernel. Unlike :func:sequence_logp, it is not meant to be called directly on a batch axis. The LD-DPO ld_alpha split is not supported on the fused path; use the eager :func:sequence_logp for that.

Parameters:

Name	Type	Description	Default
hidden_states	Tensor	Last-layer hidden states (T, H) for one sequence (batched to (B, T, H) only by an outer vmap).	required
lm_head_weight	Tensor	LM-head weight (V, H); logits are hidden_states @ lm_head_weight.T.	required
input_ids	Tensor	Token ids (T,).	required
completion_mask	Tensor	(T,); non-zero on completion-span tokens.	required
length_normalized	bool	Divide by the completion-token count for the per-token mean reward (SimPO / ORPO); matches sequence_logp.	False

Returns:

Name	Type	Description
	Tensor	Scalar completion log-prob (per-token mean when length_normalized);
matches	Tensor	func:sequence_logp.

Source code in packages/opake-alignment/src/opake/api/alignment/logprob/_sequence.py

def fused_sequence_logp(
    hidden_states: torch.Tensor,
    lm_head_weight: torch.Tensor,
    input_ids: torch.Tensor,
    completion_mask: torch.Tensor,
    *,
    length_normalized: bool = False,
) -> torch.Tensor:
    """Memory-efficient :func:`sequence_logp` over hidden states (fused, per-example).

    Mathematically ``sequence_logp(hidden_states @ lm_head_weight.T, input_ids,
    completion_mask)``, but on CUDA + half precision with
    ``opake-alignment[patches]`` installed it computes the completion log-prob
    via the patches fused linear-CE kernel — no ``(T, V)`` logits are
    materialised and the LSE is recomputed in the backward. Otherwise it falls
    back to the eager form.

    Since ``sequence_logp = Σ_completion log p = −Σ_completion CE``, the kernel is
    the preference-logp kernel: encode the completion span as the kept labels
    (everything else → ``-100``), take the kernel's CE sum, and negate.

    **Per-example.** Pass one sequence ``(T, H)`` and drive with
    ``vmap(grad(...))`` (how ``train_dpo`` already computes logp inside
    ``clipped_grad``); the kernel's merged vmap rules then make the chosen /
    rejected microbatch one forward + one backward kernel. Unlike
    :func:`sequence_logp`, it is *not* meant to be called directly on a batch
    axis. The LD-DPO ``ld_alpha`` split is not supported on the fused path; use
    the eager :func:`sequence_logp` for that.

    Args:
        hidden_states: Last-layer hidden states ``(T, H)`` for one sequence
            (batched to ``(B, T, H)`` only by an outer ``vmap``).
        lm_head_weight: LM-head weight ``(V, H)``; logits are
            ``hidden_states @ lm_head_weight.T``.
        input_ids: Token ids ``(T,)``.
        completion_mask: ``(T,)``; non-zero on completion-span tokens.
        length_normalized: Divide by the completion-token count for the per-token
            mean reward (SimPO / ORPO); matches ``sequence_logp``.

    Returns:
        Scalar completion log-prob (per-token mean when ``length_normalized``);
        matches :func:`sequence_logp`.
    """
    if lce_available(hidden_states):
        # Completion tokens keep their id; everything else → ignore_index, so the
        # kernel's CE sum is exactly −Σ_completion logp. The kernel applies the
        # next-token shift internally, matching sequence_logp's shifted mask.
        masked_labels = torch.where(
            completion_mask.to(torch.bool),
            input_ids,
            torch.full_like(input_ids, -100),
        )
        seq_logp = -linear_nll_sum(hidden_states, lm_head_weight, masked_labels)
        if length_normalized:
            count = completion_mask[..., 1:].sum(-1).clamp(min=1)
            seq_logp = seq_logp / count.to(seq_logp.dtype)
        return seq_logp
    logits = hidden_states @ lm_head_weight.transpose(-2, -1)
    return sequence_logp(
        logits, input_ids, completion_mask, length_normalized=length_normalized
    )

sequence_logp(logits, input_ids, completion_mask, *, ld_alpha=None, shared_prefix_len=None, length_normalized=False)

Sum (or per-token mean) of completion-token log-probabilities per sequence.

Applies the causal-LM shift, computes per-token log-probs, multiplies by the (shifted) completion mask, and sums over the sequence (last remaining) axis.

Parameters:

Name	Type	Description	Default
logits	Tensor	Float tensor of shape (..., T, V). The (...) leading dims may be empty (per-example under vmap) or a batch axis.	required
input_ids	Tensor	Integer tensor of shape (..., T) (token ids).	required
completion_mask	Tensor	Tensor of shape (..., T); non-zero where a token belongs to the completion span and should contribute to the logp. It is cast to the logits dtype before multiplying.	required
ld_alpha	float | None	LD-DPO (arXiv:2409.06411) length-desensitisation coefficient in [0, 1]. When set, completion tokens beyond shared_prefix_len are weighted by ld_alpha instead of 1.0, damping the verbose tail's contribution; ld_alpha=1.0 recovers the plain masked sum. The fused path does not support it — LD-DPO needs the per-token log-probs, so use this eager sequence_logp.	None
shared_prefix_len	int | Tensor | None	LD-DPO public length (the shorter completion length), retained under this compatibility name. An int or a per-example tensor broadcasting against the shifted position axis. Required when ld_alpha is set; ignored otherwise. This single-sequence helper cannot derive it: callers claiming LD-DPO must pass the minimum of the paired completion lengths.	None
length_normalized	bool	Divide the per-sequence logp by its completion-token count, giving the per-token mean reward (1/|y|)·log π(y) used by the reference-free heads (SimPO / ORPO). Default False (plain sum).	False

Returns:

Type	Description
Tensor	Float tensor of shape (...) (one logp per sequence; a per-token mean
Tensor	when length_normalized).

Raises:

Type	Description
ValueError	If ld_alpha is outside [0, 1] or is set without shared_prefix_len.

Source code in packages/opake-alignment/src/opake/api/alignment/logprob/_sequence.py

def sequence_logp(
    logits: torch.Tensor,
    input_ids: torch.Tensor,
    completion_mask: torch.Tensor,
    *,
    ld_alpha: float | None = None,
    shared_prefix_len: int | torch.Tensor | None = None,
    length_normalized: bool = False,
) -> torch.Tensor:
    """Sum (or per-token mean) of completion-token log-probabilities per sequence.

    Applies the causal-LM shift, computes per-token log-probs, multiplies by
    the (shifted) completion mask, and sums over the sequence (last remaining)
    axis.

    Args:
        logits: Float tensor of shape ``(..., T, V)``. The ``(...)`` leading
            dims may be empty (per-example under ``vmap``) or a batch axis.
        input_ids: Integer tensor of shape ``(..., T)`` (token ids).
        completion_mask: Tensor of shape ``(..., T)``; non-zero where a token
            belongs to the completion span and should contribute to the logp.
            It is cast to the logits dtype before multiplying.
        ld_alpha: LD-DPO (arXiv:2409.06411) length-desensitisation coefficient
            in ``[0, 1]``.
            When set, completion tokens beyond ``shared_prefix_len`` are weighted
            by ``ld_alpha`` instead of ``1.0``, damping
            the verbose tail's contribution; ``ld_alpha=1.0`` recovers the plain
            masked sum. The fused path does not support it — LD-DPO needs the
            per-token log-probs, so use this eager ``sequence_logp``.
        shared_prefix_len: LD-DPO public length (the shorter completion length),
            retained under this compatibility name. An ``int`` or a per-example
            tensor broadcasting against the shifted position axis. Required when
            ``ld_alpha`` is set; ignored otherwise. This single-sequence helper
            cannot derive it: callers claiming LD-DPO must pass the minimum of
            the paired completion lengths.
        length_normalized: Divide the per-sequence logp by its completion-token
            count, giving the per-token mean reward ``(1/|y|)·log π(y)`` used by
            the reference-free heads (SimPO / ORPO). Default ``False`` (plain sum).

    Returns:
        Float tensor of shape ``(...)`` (one logp per sequence; a per-token mean
        when ``length_normalized``).

    Raises:
        ValueError: If ``ld_alpha`` is outside ``[0, 1]`` or is set without
            ``shared_prefix_len``.
    """
    # Causal-LM shift: predict token t+1 from the logits at position t.
    shifted_logits = logits[..., :-1, :]
    target_ids = input_ids[..., 1:]
    target_mask = completion_mask[..., 1:]
    target_weight = target_mask.to(shifted_logits.dtype)
    per_token_logp = selective_log_softmax(shifted_logits, target_ids)

    weight = target_weight
    if ld_alpha is not None:
        if not math.isfinite(ld_alpha) or not 0.0 <= ld_alpha <= 1.0:
            raise ConfigurationError(
                *(f"ld_alpha (LD-DPO) must be finite and in [0, 1], got {ld_alpha}",)
            )
        if shared_prefix_len is None:
            raise ConfigurationError(*("ld_alpha (LD-DPO) requires shared_prefix_len",))
        completion_pos = target_mask.to(torch.bool).cumsum(dim=-1)
        prefix_len = shared_prefix_len
        if isinstance(prefix_len, torch.Tensor):
            while prefix_len.ndim < completion_pos.ndim:
                prefix_len = prefix_len.unsqueeze(-1)
        is_prefix = (completion_pos >= 1) & (completion_pos <= prefix_len)
        ld_weight = torch.where(
            is_prefix,
            torch.ones((), dtype=per_token_logp.dtype, device=per_token_logp.device),
            torch.full(
                (), ld_alpha, dtype=per_token_logp.dtype, device=per_token_logp.device
            ),
        )
        weight = target_weight * ld_weight

    acc_dtype = _compute_dtype(per_token_logp)
    logp = (per_token_logp * weight).to(acc_dtype).sum(dim=-1)
    if length_normalized:
        token_count = target_mask.to(torch.bool).sum(-1).clamp(min=1)
        logp = logp / token_count.to(acc_dtype)
    return logp

collator

opake.alignment.dpo.collator façade — the preference (DPO) collator.

The language-modeling (SFT) collator lives under :mod:opake.alignment.sft.collator.

preference_collator(pad_token_id, max_length, *, pad_to_multiple_of=None)

DPO preference collator — factory returning a callable.

.. note:: Layout choice

TRL uses a (2B, L) concatenated layout; opake-alignment uses a (B, ...) layout because per-example DP-SGD clips per preference pair — the chosen/rejected forward passes share one per-example gradient.

Parameters:

Name	Type	Description	Default
pad_token_id	int	Token id used for right-padding.	required
max_length	int	Maximum sequence length. Sequences longer than this are truncated from the right (keep-start).	required
pad_to_multiple_of	int | None	When given, each side's padded sequence length is independently rounded up to the nearest multiple.	None

Returns:

Type	Description
Callable[[list[PreferenceExample]], PreferenceBatch]	A callable collate(batch: list[PreferenceExample]) -> PreferenceBatch
Callable[[list[PreferenceExample]], PreferenceBatch]	that accepts a list of example dicts and returns:
Callable[[list[PreferenceExample]], PreferenceBatch]	Mandatory keys (always present):
Callable[[list[PreferenceExample]], PreferenceBatch]	chosen_input_ids: (B, Lc) long tensor, right-padded with pad_token_id.
Callable[[list[PreferenceExample]], PreferenceBatch]	chosen_attention_mask: (B, Lc) long tensor; 1 for real tokens, 0 for pad.
Callable[[list[PreferenceExample]], PreferenceBatch]	chosen_completion_mask: (B, Lc) long tensor; 1 in the completion span, 0 else.
Callable[[list[PreferenceExample]], PreferenceBatch]	rejected_input_ids: (B, Lr) long tensor.
Callable[[list[PreferenceExample]], PreferenceBatch]	rejected_attention_mask: (B, Lr) long tensor.
Callable[[list[PreferenceExample]], PreferenceBatch]	rejected_completion_mask: (B, Lr) long tensor.
Callable[[list[PreferenceExample]], PreferenceBatch]	Lc and Lr are the per-side batch maxima (may differ from each
Callable[[list[PreferenceExample]], PreferenceBatch]	other) after truncation, rounded up to pad_to_multiple_of if given.
Callable[[list[PreferenceExample]], PreferenceBatch]	Optional keys (present iff all examples in the batch carry them):
Callable[[list[PreferenceExample]], PreferenceBatch]	ref_chosen_logps: (B,) float tensor.
Callable[[list[PreferenceExample]], PreferenceBatch]	ref_rejected_logps: (B,) float tensor.

Each input example dict must have:

• chosen_input_ids: list[int]
• rejected_input_ids: list[int]
• chosen_completion_mask: list[int]
• rejected_completion_mask: list[int]

And may optionally carry ref_chosen_logps: float and ref_rejected_logps: float.

Source code in packages/opake-alignment/src/opake/api/alignment/dpo/collator/_preference.py

def preference_collator(
    pad_token_id: int,
    max_length: int,
    *,
    pad_to_multiple_of: int | None = None,
) -> Callable[[list[PreferenceExample]], PreferenceBatch]:
    """DPO preference collator — factory returning a callable.

    .. note:: Layout choice

       TRL uses a ``(2B, L)`` concatenated layout; opake-alignment uses a
       ``(B, ...)`` layout because per-example DP-SGD clips per preference
       pair — the chosen/rejected forward passes share one per-example
       gradient.

    Args:
        pad_token_id: Token id used for right-padding.
        max_length: Maximum sequence length.  Sequences longer than this are
            truncated from the right (keep-start).
        pad_to_multiple_of: When given, each side's padded sequence length is
            independently rounded up to the nearest multiple.

    Returns:
        A callable ``collate(batch: list[PreferenceExample]) -> PreferenceBatch``
        that accepts a list of example dicts and returns:

        **Mandatory keys** (always present):

        - ``chosen_input_ids``: ``(B, Lc)`` long tensor, right-padded with ``pad_token_id``.
        - ``chosen_attention_mask``: ``(B, Lc)`` long tensor; 1 for real tokens, 0 for pad.
        - ``chosen_completion_mask``: ``(B, Lc)`` long tensor; 1 in the completion span, 0 else.
        - ``rejected_input_ids``: ``(B, Lr)`` long tensor.
        - ``rejected_attention_mask``: ``(B, Lr)`` long tensor.
        - ``rejected_completion_mask``: ``(B, Lr)`` long tensor.

        ``Lc`` and ``Lr`` are the per-side batch maxima (may differ from each
        other) after truncation, rounded up to ``pad_to_multiple_of`` if given.

        **Optional keys** (present iff *all* examples in the batch carry them):

        - ``ref_chosen_logps``: ``(B,)`` float tensor.
        - ``ref_rejected_logps``: ``(B,)`` float tensor.

    Each input example dict must have:

    - ``chosen_input_ids``: ``list[int]``
    - ``rejected_input_ids``: ``list[int]``
    - ``chosen_completion_mask``: ``list[int]``
    - ``rejected_completion_mask``: ``list[int]``

    And may optionally carry ``ref_chosen_logps: float`` and
    ``ref_rejected_logps: float``.
    """

    def collate(batch: list[PreferenceExample]) -> PreferenceBatch:
        # ------------------------------------------------------------------ #
        # Extract per-side sequences.
        # ------------------------------------------------------------------ #
        chosen_ids_raw: list[list[int]] = [ex["chosen_input_ids"] for ex in batch]
        rejected_ids_raw: list[list[int]] = [ex["rejected_input_ids"] for ex in batch]
        chosen_cmask_raw: list[list[int]] = [
            ex["chosen_completion_mask"] for ex in batch
        ]
        rejected_cmask_raw: list[list[int]] = [
            ex["rejected_completion_mask"] for ex in batch
        ]

        # Truncate completion masks in sync with input ids (keep-start).
        chosen_cmask_trunc: list[list[int]] = [m[:max_length] for m in chosen_cmask_raw]
        rejected_cmask_trunc: list[list[int]] = [
            m[:max_length] for m in rejected_cmask_raw
        ]

        # ------------------------------------------------------------------ #
        # Pad chosen side.
        # ------------------------------------------------------------------ #
        chosen_input_ids, chosen_attention_mask = _pad_side(
            chosen_ids_raw,
            pad_token_id=pad_token_id,
            max_length=max_length,
            pad_to_multiple_of=pad_to_multiple_of,
        )
        chosen_completion_mask = _pad_completion_mask(
            chosen_cmask_trunc,
            target_length=chosen_input_ids.shape[1],
        )

        # ------------------------------------------------------------------ #
        # Pad rejected side (independently — Lc may differ from Lr).
        # ------------------------------------------------------------------ #
        rejected_input_ids, rejected_attention_mask = _pad_side(
            rejected_ids_raw,
            pad_token_id=pad_token_id,
            max_length=max_length,
            pad_to_multiple_of=pad_to_multiple_of,
        )
        rejected_completion_mask = _pad_completion_mask(
            rejected_cmask_trunc,
            target_length=rejected_input_ids.shape[1],
        )

        # ------------------------------------------------------------------ #
        # Assemble output dict.
        # ------------------------------------------------------------------ #
        out: PreferenceBatch = {
            "chosen_input_ids": chosen_input_ids,
            "chosen_attention_mask": chosen_attention_mask,
            "chosen_completion_mask": chosen_completion_mask,
            "rejected_input_ids": rejected_input_ids,
            "rejected_attention_mask": rejected_attention_mask,
            "rejected_completion_mask": rejected_completion_mask,
        }

        # Optional ref logps — included only if ALL examples carry them.
        if all("ref_chosen_logps" in ex for ex in batch):
            out["ref_chosen_logps"] = torch.tensor(
                [ex["ref_chosen_logps"] for ex in batch],
                dtype=torch.float32,
            )
        if all("ref_rejected_logps" in ex for ex in batch):
            out["ref_rejected_logps"] = torch.tensor(
                [ex["ref_rejected_logps"] for ex in batch],
                dtype=torch.float32,
            )

        return out

    return collate

data

opake.alignment.dpo.data façade — preference prompt extraction.

extract_prompt(example)

Extract the shared prompt prefix from a preference example.

Given a single dataset row, returns a new dict with the same keys plus a "prompt" key holding the longest common prefix of "chosen" and "rejected". The values for "chosen" and "rejected" are replaced by their respective suffixes (the parts after the common prefix).

If the example already contains a "prompt" key the dict is returned unchanged (idempotent). Unpaired examples ("prompt" + "completion", no chosen/rejected) are also returned unchanged.

All other keys in example are preserved without modification.

Parameters:

Name	Type	Description	Default
example	dict	A single dataset record — a plain :class:dict. For preference examples the values of "chosen" and "rejected" must be :class:list objects. Elements are compared with == and need not be JSON-serialisable.	required

Returns:

Type	Description
dict	A new :class:dict with the same keys as example. For preference
dict	examples without a pre-existing "prompt", "prompt",
dict	"chosen", and "rejected" are rewritten as described above.
dict	All other cases return the original dict unchanged.

Examples:

Conversational preference pair::

>>> example = {
...     "chosen":   [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "good"}],
...     "rejected": [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "bad"}],
... }
>>> extract_prompt(example)
{
    "prompt":   [{"role": "user", "content": "hi"}],
    "chosen":   [{"role": "assistant", "content": "good"}],
    "rejected": [{"role": "assistant", "content": "bad"}],
}

Idempotent when "prompt" is already present::

>>> extract_prompt({"prompt": [...], "chosen": [...], "rejected": [...]})
# returned unchanged

Source code in packages/opake-alignment/src/opake/api/alignment/dpo/data/_prompt.py

def extract_prompt(example: dict) -> dict:
    """Extract the shared prompt prefix from a preference example.

    Given a single dataset row, returns a new dict with the same keys plus
    a ``"prompt"`` key holding the longest common prefix of ``"chosen"`` and
    ``"rejected"``.  The values for ``"chosen"`` and ``"rejected"`` are
    replaced by their respective suffixes (the parts after the common prefix).

    If the example already contains a ``"prompt"`` key the dict is returned
    **unchanged** (idempotent).  Unpaired examples (``"prompt"`` +
    ``"completion"``, no chosen/rejected) are also returned unchanged.

    All other keys in *example* are preserved without modification.

    Args:
        example: A single dataset record — a plain :class:`dict`.  For
            preference examples the values of ``"chosen"`` and ``"rejected"``
            must be :class:`list` objects.  Elements are compared with ``==``
            and need not be JSON-serialisable.

    Returns:
        A new :class:`dict` with the same keys as *example*.  For preference
        examples without a pre-existing ``"prompt"``, ``"prompt"``,
        ``"chosen"``, and ``"rejected"`` are rewritten as described above.
        All other cases return the original dict unchanged.

    Examples:
        Conversational preference pair::

            >>> example = {
            ...     "chosen":   [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "good"}],
            ...     "rejected": [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "bad"}],
            ... }
            >>> extract_prompt(example)
            {
                "prompt":   [{"role": "user", "content": "hi"}],
                "chosen":   [{"role": "assistant", "content": "good"}],
                "rejected": [{"role": "assistant", "content": "bad"}],
            }

        Idempotent when ``"prompt"`` is already present::

            >>> extract_prompt({"prompt": [...], "chosen": [...], "rejected": [...]})
            # returned unchanged
    """
    # Idempotency: if "prompt" already exists, return unchanged.
    if "prompt" in example:
        return example

    # Only operate on preference examples (both chosen and rejected present).
    if "chosen" not in example or "rejected" not in example:
        return example

    chosen: list = example["chosen"]
    rejected: list = example["rejected"]

    # Walk element-wise; stop at first mismatch.
    prefix_len = 0
    for c_elem, r_elem in zip(chosen, rejected, strict=False):
        if c_elem == r_elem:
            prefix_len += 1
        else:
            break

    prompt = chosen[:prefix_len]
    chosen_suffix = chosen[prefix_len:]
    rejected_suffix = rejected[prefix_len:]

    return {
        **example,
        "prompt": prompt,
        "chosen": chosen_suffix,
        "rejected": rejected_suffix,
    }

reference

opake.alignment.dpo.reference — reference-model handling for DPO.

DPO scores each response against a frozen reference policy π_ref. These helpers produce π_ref's log-probabilities and keep that work outside the per-example vmap(grad(...)) region (it runs a separate forward and/or mutates nn.Module adapter state):

• compute_ref_logprobs_for_dataset(dataset, ref, ...) — the common path: run the reference once over the dataset (outer loop), caching per-example chosen/rejected sequence logps to a content-addressed .safetensors cache keyed by dataset identity, a structured model/tokenizer/preprocessing cache_identity, and output_columns. Feed the cached ref_chosen_logps / ref_rejected_logps into the per-pair head as the log-ratio baseline policy_logp - ref_logp.
• null_ref_context(model) / with_disabled_adapter(model) — when the policy is a PEFT/LoRA adapter, use the base model as the reference: enter the context (which disables the adapter), run the reference forward, no second model needed.
• ema_update_reference(...) — TR-DPO: periodically move the reference toward the policy by an EMA step between training steps.

End-to-end wiring is in examples/train_dpo.py.

compute_ref_logprobs_for_dataset(dataset, ref, collator, output_columns, *, cache_identity, batch_size=8, cache_dir=None, use_cache=True, shard=None)

Compute per-example reference logprobs, optionally using a cache.

Runs a single pass over dataset calling ref(batch) for each collated batch, collecting one (B,) tensor per name in output_columns, and attaches the concatenated per-example results to dataset as new columns. The forward runs under torch.no_grad(). When use_cache is true, results are cached to a content-addressed .safetensors file keyed on (dataset._fingerprint, cache_identity, output_columns). The cache directory and archive are owner-only because they contain private per-example values. The caller is responsible for removing selected cache directories when their contents are no longer needed.

Under a live process group each rank runs the reference over its own shard and the results are gathered back into dataset order, so the returned dataset is the same on every rank and matches a single-process run. Pass the whole dataset on every rank; shard for training afterwards.

Outside vmap only — see the module docstring.

Parameters:

Name	Type	Description	Default
dataset	Any	A :class:datasets.Dataset, identical on every rank. Iterated via a DataLoader with collator — over this rank's shard when a process group is live — and returned with output_columns added through :meth:datasets.Dataset.add_column.	required
ref	Callable[[dict[str, Tensor]], dict[str, Tensor]]	A callable mapping a collated batch dict[str, Tensor] to a dict containing one (B,) tensor per name in output_columns. A :class:PreTrainedModel is wrapped into such a callable by the caller, keeping this mechanism-agnostic.	required
collator	Callable[[list[dict]], dict[str, Tensor]]	Callable[[list[dict]], dict[str, Tensor]] producing a batched dict from a list of raw rows (the same collator later used for training).	required
output_columns	Sequence[str]	The keys ref returns, e.g. ("ref_chosen_logps", "ref_rejected_logps"). Also folded into the cache fingerprint.	required
cache_identity	Any	Deterministic JSON-like identity for ref. The DPO trainer supplies its effective reference-model state; dataset preparation changes are represented by dataset._fingerprint. Mapping order does not affect the fingerprint. Unsupported values raise TypeError.	required
batch_size	int	DataLoader batch size for the forward pass. Default 8.	8
cache_dir	str | None	Directory for the .safetensors cache. Defaults to <tempdir>/opake_ref_cache (created on first miss).	None
use_cache	bool	Whether to read or write the persistent cache. Set to False when results cannot be reused, such as TR-DPO seed values that are overwritten before each training and evaluation step. Must agree across ranks.	True
shard	bool | None	Whether each rank scores only its own slice of dataset. None (default) shards when a process group is live. False makes every rank score the whole dataset. True requires an initialised process group and raises without one. Must agree across ranks.	None

Returns:

Type	Description
Any	dataset with one new column per name in output_columns, each of
Any	length len(dataset).

Source code in packages/opake-alignment/src/opake/api/alignment/dpo/reference/_precompute.py

def compute_ref_logprobs_for_dataset(
    dataset: Any,
    ref: Callable[[dict[str, torch.Tensor]], dict[str, torch.Tensor]],
    collator: Callable[[list[dict]], dict[str, torch.Tensor]],
    output_columns: Sequence[str],
    *,
    cache_identity: Any,
    batch_size: int = 8,
    cache_dir: str | None = None,
    use_cache: bool = True,
    shard: bool | None = None,
) -> Any:
    """Compute per-example reference logprobs, optionally using a cache.

    Runs a single pass over ``dataset`` calling ``ref(batch)`` for each
    collated batch, collecting one ``(B,)`` tensor per name in
    ``output_columns``, and attaches the concatenated per-example results to
    ``dataset`` as new columns. The forward runs under ``torch.no_grad()``.
    When ``use_cache`` is true, results are cached to a content-addressed
    ``.safetensors`` file keyed on ``(dataset._fingerprint, cache_identity,
    output_columns)``. The cache directory and archive are owner-only because
    they contain private per-example values. The caller is responsible for
    removing selected cache directories when their contents are no longer
    needed.

    Under a live process group each rank runs the reference over its own shard
    and the results are gathered back into dataset order, so the returned
    dataset is the same on every rank and matches a single-process run. Pass
    the whole dataset on every rank; shard for training afterwards.

    **Outside vmap only** — see the module docstring.

    Args:
        dataset: A :class:`datasets.Dataset`, identical on every rank. Iterated
            via a ``DataLoader`` with ``collator`` — over this rank's shard when
            a process group is live — and returned with ``output_columns`` added
            through :meth:`datasets.Dataset.add_column`.
        ref: A callable mapping a collated batch ``dict[str, Tensor]`` to a dict
            containing one ``(B,)`` tensor per name in ``output_columns``. A
            :class:`PreTrainedModel` is wrapped into such a callable by the
            caller, keeping this mechanism-agnostic.
        collator: ``Callable[[list[dict]], dict[str, Tensor]]`` producing a
            batched dict from a list of raw rows (the same collator later used
            for training).
        output_columns: The keys ``ref`` returns, e.g.
            ``("ref_chosen_logps", "ref_rejected_logps")``. Also folded into the
            cache fingerprint.
        cache_identity: Deterministic JSON-like identity for ``ref``. The DPO
            trainer supplies its effective reference-model state; dataset
            preparation changes are represented by ``dataset._fingerprint``.
            Mapping order does not affect the fingerprint. Unsupported values
            raise ``TypeError``.
        batch_size: ``DataLoader`` batch size for the forward pass. Default 8.
        cache_dir: Directory for the ``.safetensors`` cache. Defaults to
            ``<tempdir>/opake_ref_cache`` (created on first miss).
        use_cache: Whether to read or write the persistent cache. Set to
            ``False`` when results cannot be reused, such as TR-DPO seed
            values that are overwritten before each training and evaluation
            step. Must agree across ranks.
        shard: Whether each rank scores only its own slice of ``dataset``.
            ``None`` (default) shards when a process group is live. ``False``
            makes every rank score the whole dataset. ``True`` requires an
            initialised process group and raises without one. Must agree
            across ranks.

    Returns:
        ``dataset`` with one new column per name in ``output_columns``, each of
        length ``len(dataset)``.
    """
    columns = tuple(output_columns)
    n_examples = len(dataset)
    sharded = _resolve_sharding(shard)
    if sharded:
        # Sharding is only well defined when every rank holds the same dataset.
        assert_scalar_equal(n_examples, name="reference precompute dataset size")
        dataset_id = _dataset_fingerprint(dataset)
        assert_string_equal(dataset_id, name="reference precompute dataset fingerprint")
        if dataset_id is None:
            raise ConfigurationError(
                *(
                    "dataset must expose a deterministic `_fingerprint` for "
                    "distributed reference precomputation",
                )
            )

    path: str | None = None
    cached: dict[str, torch.Tensor] | None = None
    if use_cache:
        fingerprint = _cache_fingerprint(dataset, cache_identity, columns)
        path = _cache_path(cache_dir, fingerprint)
        _secure_cache_path(path)
        cached = _load_cache(path, columns, expected_rows=n_examples)

    # Reduced unconditionally: the collective sequence must not depend on the
    # rank-local ``use_cache``.
    if _cache_hit_on_every_rank(cached is not None):
        # HIT on every rank: attach cached columns WITHOUT calling ``ref``.
        # PyArrow has no bf16 column type, so demote to a Python ``float``
        # list at the storage boundary (explicit, single conversion).
        assert cached is not None
        result = dataset
        for name in columns:
            result = result.add_column(name, cached[name].float().tolist())
        return result

    # MISS: run the reference forward once over this rank's shard.
    local = _reference_forward(
        dataset,
        ref,
        collator,
        columns,
        batch_size=batch_size,
        rank=get_rank() if sharded else 0,
        world_size=get_world_size() if sharded else 1,
    )
    gathered = _gather_columns(
        local, columns, expected_rows=n_examples, sharded=sharded
    )

    # Only the main process writes the cache; all ranks sync before returning.
    if use_cache and is_main_process():
        cpu_columns = {
            name: gathered[name].detach().cpu().contiguous() for name in columns
        }
        assert path is not None
        _save_cache(path, cpu_columns)
    wait_for_everyone()

    result = dataset
    for name in columns:
        # Same explicit fp32 demotion at the HF/PyArrow boundary as on the
        # HIT path above. Native dtype is preserved on disk; only the column
        # exposed to ``datasets`` is downcast.
        values = gathered[name].detach().cpu().float().tolist()
        result = result.add_column(name, values)
    return result

ema_update_reference(ref_params, policy_params, alpha)

Functional EMA update of the reference parameters toward the policy.

Implements the TR-DPO core update (arXiv:2404.09656):

ref ← (1 − alpha) · ref + alpha · policy

Applied leafwise over a parameter pytree. Returns a new pytree of the same structure and does not mutate either input.

Parameters:

Name	Type	Description	Default
ref_params	PyTree	Current reference-model parameters as a pytree of :class:torch.Tensor leaves (e.g. the dict returned by make_functional(partition_trainable=True)).	required
policy_params	PyTree	Current policy-model parameters as a pytree with the same structure as ref_params.	required
alpha	float	EMA interpolation coefficient in [0, 1]. alpha = 0 — keep ref_params unchanged (output equals ref_params values). alpha = 1 — copy policy_params exactly (output equals policy_params values).	required

Returns:

Type	Description
PyTree	A new pytree of tensors with the same structure as ref_params,
PyTree	where each leaf r is replaced by
PyTree	(1 - alpha) * r + alpha * p for the corresponding policy leaf
PyTree	p.

Note

This function is pure (no in-place mutation) and is intended to be called outside the per-example gradient loop. It is safe to call under torch.no_grad() to avoid accumulating an autograd graph over the reference parameters.

Source code in packages/opake-alignment/src/opake/api/alignment/dpo/reference/_sync.py

def ema_update_reference(
    ref_params: PyTree,
    policy_params: PyTree,
    alpha: float,
) -> PyTree:
    """Functional EMA update of the reference parameters toward the policy.

    Implements the TR-DPO core update (arXiv:2404.09656):

        ref ← (1 − alpha) · ref + alpha · policy

    Applied leafwise over a parameter pytree.  Returns a **new** pytree of
    the same structure and does **not** mutate either input.

    Args:
        ref_params: Current reference-model parameters as a pytree of
            :class:`torch.Tensor` leaves (e.g. the dict returned by
            ``make_functional(partition_trainable=True)``).
        policy_params: Current policy-model parameters as a pytree with the
            same structure as ``ref_params``.
        alpha: EMA interpolation coefficient in ``[0, 1]``.

            - ``alpha = 0`` — keep ``ref_params`` unchanged (output equals
              ``ref_params`` values).
            - ``alpha = 1`` — copy ``policy_params`` exactly (output equals
              ``policy_params`` values).

    Returns:
        A new pytree of tensors with the same structure as ``ref_params``,
        where each leaf ``r`` is replaced by
        ``(1 - alpha) * r + alpha * p`` for the corresponding policy leaf
        ``p``.

    Note:
        This function is **pure** (no in-place mutation) and is intended to
        be called **outside** the per-example gradient loop. It is safe to
        call under ``torch.no_grad()`` to avoid accumulating an autograd graph
        over the reference parameters.
    """
    return optree.tree_map(
        lambda r, p: (1 - alpha) * r + alpha * p,
        ref_params,
        policy_params,
    )

null_ref_context(model, ref_model=None)

Context manager that turns model into its own reference inside the block.

Dispatches per the ref-model design table:

======================== ================================================ Configuration Behaviour ======================== ================================================ Separate model ref_model is not None → no-op (caller uses ref_model directly; no policy adapter to toggle). LoRA with "ref" is_peft_model(model) and "ref" in adapter clone model.peft_config → call model.set_adapter("ref") on enter; restore the previously-active adapter(s) on exit via a finally block. LoRA without separate is_peft_model(model) and ref_model is ref None → call model.disable_adapter() (base model serves as reference); re-enable on exit via finally. Explicit callable / model is not a PEFT model and ref_model is non-PEFT no-ref None → no-op (caller-supplied callable or frozen base model, nothing to toggle). ======================== ================================================

Parameters:

Name	Type	Description	Default
model	Any	The policy model. May be a PEFT :class:~peft.PeftModel or an ordinary :class:torch.nn.Module.	required
ref_model	Any | None	Optional separate reference model. When not None, the caller is responsible for forwarding through ref_model directly; this context manager is a no-op. It must not be model itself.	None

.. warning:: Outside vmap only.

This helper mutates ``nn.Module`` adapter-selection state. Call it in
the outer loop; vmap the pure forward *inside* the activated context.

Note

Modelled on TRL utils.use_adapter (trl/trainer/utils.py).

Source code in packages/opake-alignment/src/opake/api/alignment/dpo/reference/_adapter.py

@contextmanager
def null_ref_context(
    model: Any,
    ref_model: Any | None = None,
) -> Generator[None, None, None]:
    """Context manager that turns *model* into its own reference inside the block.

    Dispatches per the ref-model design table:

    ========================  ================================================
    Configuration             Behaviour
    ========================  ================================================
    Separate model            *ref_model* is not ``None`` → **no-op**
                              (caller uses *ref_model* directly; no policy
                              adapter to toggle).
    LoRA with ``"ref"``       ``is_peft_model(model)`` and ``"ref" in
    adapter clone             model.peft_config`` → call
                 ``model.set_adapter("ref")`` on enter; restore
                              the previously-active adapter(s) on exit via a
                              ``finally`` block.
    LoRA without separate     ``is_peft_model(model)`` and *ref_model* is
    ref                       ``None`` → call ``model.disable_adapter()``
                 (base model serves as reference); re-enable on
                              exit via ``finally``.
    Explicit callable /       *model* is not a PEFT model and *ref_model* is
    non-PEFT no-ref           ``None`` → **no-op** (caller-supplied callable
                 or frozen base model, nothing to toggle).
    ========================  ================================================

    Args:
        model: The policy model.  May be a PEFT :class:`~peft.PeftModel` or an
            ordinary :class:`torch.nn.Module`.
        ref_model: Optional separate reference model. When not ``None``, the
            caller is responsible for forwarding through *ref_model* directly;
            this context manager is a no-op. It must not be ``model`` itself.

    .. warning:: **Outside vmap only**.

        This helper mutates ``nn.Module`` adapter-selection state. Call it in
        the outer loop; vmap the pure forward *inside* the activated context.

    Note:
        Modelled on TRL ``utils.use_adapter`` (``trl/trainer/utils.py``).
    """
    if ref_model is model:
        raise ConfigurationError(
            *(
                "ref_model must be a separate model; omit ref_model to use the "
                "policy's PEFT reference adapter or adapter-disabled base model.",
            )
        )

    # ── Row 1: separate ref_model ───────────────────────────────────────────
    if ref_model is not None:
        # Caller owns ref_model; nothing to toggle on the policy model.
        yield
        return

    # ── Remaining rows require model to be a PEFT model. ────────────────────
    if not _is_peft_model(model):
        # Row 4: non-PEFT model with no ref → no-op.
        yield
        return

    # ── Row 2: LoRA model with a "ref" adapter clone. ───────────────────────
    if "ref" in model.peft_config:
        saved = _get_active_adapters(model)
        model.set_adapter("ref")
        try:
            yield
        finally:
            _restore_adapters(model, saved)
        return

    # ── Row 3: LoRA model with no separate ref → disable adapter. ───────────
    # (ref_model is None here, established by the guard at Row 1.)
    with with_disabled_adapter(model):
        yield

with_disabled_adapter(model)

Thin context manager that disables adapters on model for the block.

When model is a PEFT model, enters model.disable_adapter() so the base model weights serve as the forward pass. The adapter is re-enabled on exit (finally). When model is not a PEFT model, this is a no-op.

Used internally by the LoRA-disable reference path in :func:null_ref_context and exposed publicly for callers that want to forward through the base model directly.

Parameters:

Name	Type	Description	Default
model	Any	The model whose adapters should be disabled. May be a PEFT :class:~peft.PeftModel or an ordinary :class:torch.nn.Module.	required

.. warning:: Outside vmap only.

This helper mutates ``nn.Module`` adapter-selection state. Call it in
the outer loop; vmap the pure forward *inside* the activated context.

Note

Modelled on TRL utils.use_adapter (trl/trainer/utils.py).

Source code in packages/opake-alignment/src/opake/api/alignment/dpo/reference/_adapter.py

@contextmanager
def with_disabled_adapter(model: Any) -> Generator[None, None, None]:
    """Thin context manager that disables adapters on *model* for the block.

    When *model* is a PEFT model, enters ``model.disable_adapter()`` so the
    base model weights serve as the forward pass.  The adapter is re-enabled on
    exit (``finally``).  When *model* is not a PEFT model, this is a **no-op**.

    Used internally by the LoRA-disable reference path in
    :func:`null_ref_context` and exposed publicly for callers
    that want to forward through the base model directly.

    Args:
        model: The model whose adapters should be disabled.  May be a PEFT
            :class:`~peft.PeftModel` or an ordinary :class:`torch.nn.Module`.

    .. warning:: **Outside vmap only**.

        This helper mutates ``nn.Module`` adapter-selection state. Call it in
        the outer loop; vmap the pure forward *inside* the activated context.

    Note:
        Modelled on TRL ``utils.use_adapter`` (``trl/trainer/utils.py``).
    """
    if not _is_peft_model(model):
        yield
        return

    with model.disable_adapter():
        yield

metric

opake.alignment.dpo.metric façade — preference reward telemetry.

General token-level metrics live in the shared impl :mod:opake.api.alignment.metric.

reward_metrics(chosen_logratio, rejected_logratio, *, beta)

Compute rewards/* telemetry from per-example preference log-ratios.

Parameters:

Name	Type	Description	Default
chosen_logratio	Tensor	Per-example log-ratio for the chosen response, log pi(y_w) - log pi_ref(y_w).	required
rejected_logratio	Tensor	Per-example log-ratio for the rejected response, log pi(y_l) - log pi_ref(y_l).	required
beta	float	DPO temperature; scales the log-ratios into implicit rewards.	required

Returns:

Type	Description
dict[str, Tensor]	A dict of detached scalar tensors:
dict[str, Tensor]	"rewards/chosen": mean implicit reward for chosen responses.
dict[str, Tensor]	"rewards/rejected": mean implicit reward for rejected responses.
dict[str, Tensor]	"rewards/accuracies": fraction of examples where the chosen log-ratio exceeds the rejected one.
dict[str, Tensor]	"rewards/margins": mean implicit reward margin (chosen minus rejected).

Source code in packages/opake-alignment/src/opake/api/alignment/dpo/metric/_reward.py

def reward_metrics(
    chosen_logratio: torch.Tensor,
    rejected_logratio: torch.Tensor,
    *,
    beta: float,
) -> dict[str, torch.Tensor]:
    """Compute ``rewards/*`` telemetry from per-example preference log-ratios.

    Args:
        chosen_logratio: Per-example log-ratio for the chosen response,
            ``log pi(y_w) - log pi_ref(y_w)``.
        rejected_logratio: Per-example log-ratio for the rejected response,
            ``log pi(y_l) - log pi_ref(y_l)``.
        beta: DPO temperature; scales the log-ratios into implicit rewards.

    Returns:
        A dict of detached scalar tensors:

        - ``"rewards/chosen"``: mean implicit reward for chosen responses.
        - ``"rewards/rejected"``: mean implicit reward for rejected responses.
        - ``"rewards/accuracies"``: fraction of examples where the chosen
            log-ratio exceeds the rejected one.
        - ``"rewards/margins"``: mean implicit reward margin (chosen minus
            rejected).
    """
    chosen_reward = beta * chosen_logratio
    rejected_reward = beta * rejected_logratio
    accuracies = (chosen_logratio > rejected_logratio).float()
    margins = beta * (chosen_logratio - rejected_logratio)
    return {
        "rewards/chosen": chosen_reward.mean().detach(),
        "rewards/rejected": rejected_reward.mean().detach(),
        "rewards/accuracies": accuracies.mean().detach(),
        "rewards/margins": margins.mean().detach(),
    }

metric

opake.alignment.metric — shared token-level metrics.

Detached, un-noised token telemetry; release is outside Opake's DP accounting. Preference metrics live in :mod:opake.alignment.dpo.metric.

entropy_from_logits(logits, mask)

Mean per-position Shannon entropy over masked (shifted) positions.

Logits and mask are shifted so that position t scores the predictive distribution over the token at t + 1 (standard causal-LM next-token alignment), matching :func:mean_token_accuracy. Callers therefore pass the full-length logits and mask; the shift happens here.

Parameters:

Name	Type	Description	Default
logits	Tensor	Logits of shape (..., seq, vocab); the last axis is the categorical distribution per position.	required
mask	Tensor	Boolean/float mask of shape (..., seq) selecting which (shifted) target positions contribute to the mean.	required

Returns:

Type	Description
Tensor	A detached scalar tensor: the entropy averaged over positions where the
Tensor	shifted mask is truthy (0 when no position is unmasked).

Source code in packages/opake-alignment/src/opake/api/alignment/metric/_token.py

def entropy_from_logits(logits: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
    """Mean per-position Shannon entropy over masked (shifted) positions.

    Logits and mask are shifted so that position ``t`` scores the predictive
    distribution over the token at ``t + 1`` (standard causal-LM next-token
    alignment), matching :func:`mean_token_accuracy`. Callers therefore pass the
    full-length logits and mask; the shift happens here.

    Args:
        logits: Logits of shape ``(..., seq, vocab)``; the last axis is the
            categorical distribution per position.
        mask: Boolean/float mask of shape ``(..., seq)`` selecting which
            (shifted) target positions contribute to the mean.

    Returns:
        A detached scalar tensor: the entropy averaged over positions where the
        shifted ``mask`` is truthy (``0`` when no position is unmasked).
    """
    logits = logits[..., :-1, :]
    mask = mask[..., 1:]
    probs = logits.softmax(dim=-1)
    log_probs = logits.log_softmax(dim=-1)
    entropy = -(probs * log_probs).sum(dim=-1)
    mask = mask.to(entropy.dtype)
    total = (entropy * mask).sum()
    return (total / mask.sum().clamp(min=1)).detach()

mean_token_accuracy(logits, labels, mask)

Mean next-token prediction accuracy over masked (shifted) positions.

Logits, labels, and mask are shifted so that position t predicts the token at t + 1 (standard causal-LM next-token alignment).

Parameters:

Name	Type	Description	Default
logits	Tensor	Logits of shape (..., seq, vocab).	required
labels	Tensor	Target token ids of shape (..., seq).	required
mask	Tensor	Boolean/float mask of shape (..., seq) selecting which (shifted) target positions count toward accuracy.	required

Returns:

Type	Description
Tensor	A detached scalar tensor: the fraction of unmasked positions whose
Tensor	argmax prediction matches the label (0 when none are unmasked).

Source code in packages/opake-alignment/src/opake/api/alignment/metric/_token.py

def mean_token_accuracy(
    logits: torch.Tensor,
    labels: torch.Tensor,
    mask: torch.Tensor,
) -> torch.Tensor:
    """Mean next-token prediction accuracy over masked (shifted) positions.

    Logits, labels, and mask are shifted so that position ``t`` predicts the
    token at ``t + 1`` (standard causal-LM next-token alignment).

    Args:
        logits: Logits of shape ``(..., seq, vocab)``.
        labels: Target token ids of shape ``(..., seq)``.
        mask: Boolean/float mask of shape ``(..., seq)`` selecting which
            (shifted) target positions count toward accuracy.

    Returns:
        A detached scalar tensor: the fraction of unmasked positions whose
        argmax prediction matches the label (``0`` when none are unmasked).
    """
    logits = logits[..., :-1, :]
    labels = labels[..., 1:]
    mask = mask[..., 1:]
    preds = logits.argmax(dim=-1)
    correct = (preds == labels) & mask.bool()
    mask = mask.to(logits.dtype)
    return (correct.sum() / mask.sum().clamp(min=1)).detach()