`dockyard_rl.algorithms.loss.loss_functions`¶

Loss functions for Project Dockyard RL training.

Contains: ClippedPGLossFn — GRPO / PPO / REINFORCE / DAPO / GSPO clipped PG loss NLLLossFn — Negative log-likelihood (SFT auxiliary loss) PreferenceLossFn — Preference-based loss base class DPOLossFn — Direct Preference Optimisation DistillationLossFn — Forward / reverse / mixed KL distillation DraftCrossEntropyLossFn — Speculative-decoding draft model training loss

DistributedCrossEntropy is an optional dependency imported lazily.

Module Contents¶

Classes¶

`DraftCrossEntropyLossConfig`
`DraftCrossEntropyLossDataDict`
`DraftCrossEntropyLossFn`	Auxiliary soft-target cross-entropy for draft-model training.
`ClippedPGLossConfig`
`ClippedPGLossDataDict`
`ClippedPGLossFn`	Generalised Clipped Policy Gradient loss.
`NLLLossFn`	Negative Log-Likelihood loss (SFT auxiliary).
`PreferenceLossDataDict`
`PreferenceLossFn`	Base class for preference-based losses (DPO, IPO, etc.).
`DPOLossConfig`
`DPOLossDataDict`
`DPOLossFn`	Direct Preference Optimisation loss.
`CDPOLossFn`	Conservative DPO (cDPO) — label-smoothed preference loss.
`DPOPLossFn`	DPO-positive (DPOP) — adds a positive penalty that resists the chosen log-probability dropping below the reference.
`RDPOLossFn`	Length-regularised DPO (R-DPO) — subtracts an explicit length penalty from the implicit reward to disentangle preference quality from response length.
`IPOLossFn`	Identity Preference Optimisation (IPO) — replaces the logsigmoid with a squared loss around a target margin, regularising toward the reference rather than over-fitting deterministic preferences.
`KTOLossFn`	Kahneman-Tversky Optimisation (KTO) — unpaired preference loss.
`SimPOLossFn`	SimPO — reference-free, length-normalised preference loss with a target margin.
`ORPOLossFn`	ORPO — reference-free monolithic objective: SFT NLL plus a log-odds-ratio term.
`DistillationLossConfig`
`DistillationLossDataDict`
`DistillationLossFn`	Forward / reverse / mixed KL distillation loss.
`CrossTokenizerDistillationLossConfig`	Config for (multi-teacher) cross-tokenizer distillation loss.
`CrossTokenizerDistillationLossDataDict`	Student-side keys are fixed; teacher-side keys are teacher-indexed.
`CrossTokenizerDistillationLossFn`	Multi-teacher cross-tokenizer distillation loss.

Functions¶

`build_preference_loss`	Construct the preference loss for `dpo.loss_variant`.
`is_reference_free`	Whether a variant needs no frozen-reference logprobs (SimPO/ORPO, D4).

Data¶

PREFERENCE_LOSS_REGISTRY

API¶

class dockyard_rl.algorithms.loss.loss_functions.DraftCrossEntropyLossConfig¶

Bases: typing.TypedDict

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

class dockyard_rl.algorithms.loss.loss_functions.DraftCrossEntropyLossDataDict¶

Bases: typing.TypedDict

teacher_logits: torch.Tensor¶: None

student_logits: torch.Tensor¶: None

token_mask: torch.Tensor¶: None

sample_mask: torch.Tensor¶: None

student_vocab_indices: NotRequired[torch.Tensor]¶: None

class dockyard_rl.algorithms.loss.loss_functions.DraftCrossEntropyLossFn(vocab_parallel_group: Optional[torch.distributed.ProcessGroup] = None)¶

Bases: dockyard_rl.algorithms.loss.interfaces.LossFunction

Auxiliary soft-target cross-entropy for draft-model training.

Initialization

loss_type¶: None

input_type¶: None

class dockyard_rl.algorithms.loss.loss_functions.ClippedPGLossConfig¶

Bases: pydantic.BaseModel

disable_ppo_ratio: bool¶: False

token_level_loss: bool¶: True

sequence_level_importance_ratios: bool¶: False

ratio_clip_min: float¶: 0.2

ratio_clip_max: float¶: 0.2

ratio_clip_c: Optional[float]¶: None

reference_policy_kl_penalty: float¶: 0.01

reference_policy_kl_type: str¶: ‘k3’

kl_input_clamp_value: Optional[float]¶: 20.0

kl_output_clamp_value: Optional[float]¶: 10.0

use_kl_in_reward: bool¶: False

use_importance_sampling_correction: bool¶: False

truncated_importance_sampling_type: Optional[str]¶: None

truncated_importance_sampling_ratio: Optional[float]¶: None

truncated_importance_sampling_ratio_min: Optional[float]¶: None

use_on_policy_kl_approximation: bool¶: False

force_on_policy_ratio: bool¶: False

use_cispo: bool¶: False

class dockyard_rl.algorithms.loss.loss_functions.ClippedPGLossDataDict¶

Bases: typing.TypedDict

input_ids: torch.Tensor¶: None

advantages: torch.Tensor¶: None

prev_logprobs: torch.Tensor¶: None

generation_logprobs: torch.Tensor¶: None

reference_policy_logprobs: torch.Tensor¶: None

token_mask: torch.Tensor¶: None

sample_mask: torch.Tensor¶: None

class dockyard_rl.algorithms.loss.loss_functions.ClippedPGLossFn(cfg: dockyard_rl.algorithms.loss.loss_functions.ClippedPGLossConfig)¶

Bases: dockyard_rl.algorithms.loss.interfaces.LossFunction

Generalised Clipped Policy Gradient loss.

Implements PPO, GRPO, REINFORCE/RLOO, DAPO, GSPO, and dual-clipping from a single configurable loss function.

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

where r_t(θ) = π_θ(a_t|s_t) / π_θ_old(a_t|s_t).

Initialization

input_type¶: None

class dockyard_rl.algorithms.loss.loss_functions.NLLLossFn(use_linear_ce_fusion: bool = False)¶

Bases: dockyard_rl.algorithms.loss.interfaces.LossFunction

Negative Log-Likelihood loss (SFT auxiliary).

Initialization

loss_type¶: None

input_type¶: None

class dockyard_rl.algorithms.loss.loss_functions.PreferenceLossDataDict¶

Bases: typing.TypedDict

input_ids: torch.Tensor¶: None

token_mask: torch.Tensor¶: None

sample_mask: torch.Tensor¶: None

class dockyard_rl.algorithms.loss.loss_functions.PreferenceLossFn¶

Bases: dockyard_rl.algorithms.loss.interfaces.LossFunction

Base class for preference-based losses (DPO, IPO, etc.).

loss_type¶: None

input_type¶: None

split_output_tensor(tensor: torch.Tensor) → tuple[torch.Tensor, torch.Tensor]¶

class dockyard_rl.algorithms.loss.loss_functions.DPOLossConfig¶

Bases: typing.TypedDict

reference_policy_kl_penalty: float¶: None

preference_loss_weight: float¶: None

sft_loss_weight: float¶: None

preference_average_log_probs: bool¶: None

sft_average_log_probs: bool¶: None

loss_variant: NotRequired[str]¶: None

label_smoothing: NotRequired[float]¶: None

dpop_lambda: NotRequired[float]¶: None

length_penalty: NotRequired[float]¶: None

desirable_weight: NotRequired[float]¶: None

undesirable_weight: NotRequired[float]¶: None

simpo_gamma: NotRequired[float]¶: None

orpo_lambda: NotRequired[float]¶: None

class dockyard_rl.algorithms.loss.loss_functions.DPOLossDataDict¶

Bases: typing.TypedDict

input_ids: torch.Tensor¶: None

reference_policy_logprobs: torch.Tensor¶: None

token_mask: torch.Tensor¶: None

sample_mask: torch.Tensor¶: None

class dockyard_rl.algorithms.loss.loss_functions.DPOLossFn(cfg: dockyard_rl.algorithms.loss.loss_functions.DPOLossConfig, use_linear_ce_fusion: bool = False)¶

Bases: dockyard_rl.algorithms.loss.loss_functions.PreferenceLossFn

Direct Preference Optimisation loss.

Initialization

loss_type¶: None

input_type¶: None

reference_free¶: False

sequences_per_datum¶: 2

class dockyard_rl.algorithms.loss.loss_functions.CDPOLossFn(cfg: dockyard_rl.algorithms.loss.loss_functions.DPOLossConfig, use_linear_ce_fusion: bool = False)¶

Bases: dockyard_rl.algorithms.loss.loss_functions.DPOLossFn

Conservative DPO (cDPO) — label-smoothed preference loss.

Mixes the flipped target with weight ε to tolerate noisy / mislabelled preferences: L = −[(1−ε)·logσ(βΔ) + ε·logσ(−βΔ)]. Reduces to DPO at ε=0.

Initialization

class dockyard_rl.algorithms.loss.loss_functions.DPOPLossFn(cfg: dockyard_rl.algorithms.loss.loss_functions.DPOLossConfig, use_linear_ce_fusion: bool = False)¶

Bases: dockyard_rl.algorithms.loss.loss_functions.DPOLossFn

DPO-positive (DPOP) — adds a positive penalty that resists the chosen log-probability dropping below the reference.

L = −logσ(β·(Δ − λ·max(0, s_ref_chosen − s_policy_chosen)))

Mitigates the failure where both chosen and rejected log-probs fall (a capability regression). Reduces to DPO at λ=0.

Initialization

class dockyard_rl.algorithms.loss.loss_functions.RDPOLossFn(cfg: dockyard_rl.algorithms.loss.loss_functions.DPOLossConfig, use_linear_ce_fusion: bool = False)¶

Bases: dockyard_rl.algorithms.loss.loss_functions.DPOLossFn

Length-regularised DPO (R-DPO) — subtracts an explicit length penalty from the implicit reward to disentangle preference quality from response length.

L = −logσ(β·Δ − α·(|y_chosen| − |y_rejected|))

The α penalty is outside the β scaling (per Park et al. 2024). Reduces to DPO at α=0.

Initialization

class dockyard_rl.algorithms.loss.loss_functions.IPOLossFn(cfg: dockyard_rl.algorithms.loss.loss_functions.DPOLossConfig, use_linear_ce_fusion: bool = False)¶

Bases: dockyard_rl.algorithms.loss.loss_functions.DPOLossFn

Identity Preference Optimisation (IPO) — replaces the logsigmoid with a squared loss around a target margin, regularising toward the reference rather than over-fitting deterministic preferences.

L = (Δ − 1/(2τ))²,   τ = reference_policy_kl_penalty

Distinct loss family — does not reduce to DPO at any hyperparameter.

Initialization

class dockyard_rl.algorithms.loss.loss_functions.KTOLossFn(cfg: dockyard_rl.algorithms.loss.loss_functions.DPOLossConfig, use_linear_ce_fusion: bool = False)¶

Bases: dockyard_rl.algorithms.loss.loss_functions.DPOLossFn

Kahneman-Tversky Optimisation (KTO) — unpaired preference loss.

Operates on a batch of independently-labelled examples (desirable or undesirable) rather than chosen/rejected pairs. Per example, with implicit reward (logratio) r = Σ_t (logπ − logπ_ref)·mask and a detached KL reference point z:

desirable:    λ_D · (1 − σ(β·(r − z)))
undesirable:  λ_U · (1 − σ(β·(z − r)))

z is the policy↔reference KL estimated from mismatched completions through the current policy; it is supplied at train time via data["kto_reference_kl"] (worker-computed, detached, clamped ≥ 0). When absent it defaults to 0 (no KL anchor) — the pre-worker behaviour. Consumes preference_label (1 desirable / 0 undesirable) from the unpaired kto_collate_fn.

Initialization

sequences_per_datum¶: 1

class dockyard_rl.algorithms.loss.loss_functions.SimPOLossFn(cfg: dockyard_rl.algorithms.loss.loss_functions.DPOLossConfig, use_linear_ce_fusion: bool = False)¶

Bases: dockyard_rl.algorithms.loss.loss_functions.DPOLossFn

SimPO — reference-free, length-normalised preference loss with a target margin.

Reward is the average log-probability of the response (no reference model): s = (1/|y|) Σ_t logπ(y_t). The loss is

L = −logσ(β·(s_chosen − s_rejected) − γ)

where γ = simpo_gamma is the target reward margin. Reference-free: needs no frozen-reference logprobs.

Initialization

reference_free¶: True

class dockyard_rl.algorithms.loss.loss_functions.ORPOLossFn(cfg: dockyard_rl.algorithms.loss.loss_functions.DPOLossConfig, use_linear_ce_fusion: bool = False)¶

Bases: dockyard_rl.algorithms.loss.loss_functions.DPOLossFn

ORPO — reference-free monolithic objective: SFT NLL plus a log-odds-ratio term.

With length-normalised average log-probabilities s = (1/|y|) Σ logπ:

log_odds = (s_chosen − s_rejected) − (log1mexp(s_chosen) − log1mexp(s_rejected))
L        = −s_chosen  +  λ·(−logσ(log_odds))

where λ = orpo_lambda weights the odds-ratio term against the SFT loss. Reference-free; the SFT term is the length-normalised NLL on the chosen response (sequence-weighted, matching the repo’s DPO SFT-aux convention).

Initialization

reference_free¶: True

dockyard_rl.algorithms.loss.loss_functions.PREFERENCE_LOSS_REGISTRY: dict[str, type[dockyard_rl.algorithms.loss.loss_functions.DPOLossFn]]¶: None

dockyard_rl.algorithms.loss.loss_functions.build_preference_loss(variant: Optional[str], cfg: dockyard_rl.algorithms.loss.loss_functions.DPOLossConfig, use_linear_ce_fusion: bool = False) → dockyard_rl.algorithms.loss.loss_functions.DPOLossFn¶

Construct the preference loss for dpo.loss_variant.

A missing/None variant defaults to vanilla DPO. An unrecognised variant raises a clear error rather than silently falling back.

dockyard_rl.algorithms.loss.loss_functions.is_reference_free(variant: Optional[str]) → bool¶

Whether a variant needs no frozen-reference logprobs (SimPO/ORPO, D4).

Consulted by the DPO driver to skip the reference-logprob stage. Returns False for every variant currently registered, so the frozen-ref path is unchanged for dpo/dpop/cdpo/rdpo/ipo/kto.

class dockyard_rl.algorithms.loss.loss_functions.DistillationLossConfig¶

Bases: typing.TypedDict

kl_type: str¶: None

mixed_kl_weight: float¶: None

zero_outside_topk: bool¶: None

class dockyard_rl.algorithms.loss.loss_functions.DistillationLossDataDict¶

Bases: typing.TypedDict

input_ids: torch.Tensor¶: None

input_lengths: torch.Tensor¶: None

token_mask: torch.Tensor¶: None

sample_mask: torch.Tensor¶: None

teacher_topk_logits: torch.Tensor¶: None

teacher_topk_indices: torch.Tensor¶: None

class dockyard_rl.algorithms.loss.loss_functions.DistillationLossFn(cfg: dockyard_rl.algorithms.loss.loss_functions.DistillationLossConfig)¶

Bases: dockyard_rl.algorithms.loss.interfaces.LossFunction

Forward / reverse / mixed KL distillation loss.

Initialization

loss_type¶: None

input_type¶: None

class dockyard_rl.algorithms.loss.loss_functions.CrossTokenizerDistillationLossConfig¶

Bases: typing.TypedDict

Config for (multi-teacher) cross-tokenizer distillation loss.

gold_loss / xtoken_loss and the scalar knobs below are the global defaults shared across every teacher; the per-teacher runtime-injected lists (projection_matrix_paths etc.) carry one entry per teachers[i] and can override the gold/xtoken flags per teacher in kd_loss_mode="sum". The single-teacher path is just a length-1 list.

Attributes: gold_loss: If True, use the gold-loss formulation: split the vocab into an exact-token-mapped common set (KL) and an uncommon set (L1). xtoken_loss: Modifier inside the gold-loss path. If True, relaxes the exact-map threshold to >= 0.6 (vs == 1.0) and adds a collision-replacement rule. Requires gold_loss=True. temperature: Softmax temperature applied symmetrically before KL. vocab_topk: Microbatch-global top-k size for the P-KL path (gold_loss=False). Inert when gold_loss=True. uncommon_topk: Cap on the L1 uncommon-tail sort in the gold path. Inert when gold_loss=False. reverse_kl: If True, KL(student || teacher) instead of KL(teacher || student). exact_token_match_only: P-KL path only — if True, only ‘is_correct’ pairs contribute to KL. kl_loss_weight: Multiplier on the aggregated KD term in fixed-weight mode. ce_loss_scale: Multiplier on the next-token CE term in fixed-weight mode. dynamic_loss_scaling: If True, rescale the KD term each step to match the detached CE magnitude; kl_loss_weight / ce_loss_scale are ignored. kd_loss_mode: How the per-teacher KD terms combine — "sum" (weighted sum), "averaged_logits" (convex-average teacher logits then one KL), or "select_teacher" (use only the lowest-CE teacher). normalize_teacher_by_vocab: sum-mode only — scale each teacher’s KD by log(V_t_i)/log(min_j V_t_j) so larger-vocab teachers don’t dominate purely by vocab size. alpha: Softmax temperature on the dynamic teacher-weight scores (sum_weights_metric). Inert when weights are static. sum_weights_metric: sum-mode only — "ce" / "entropy" / "max_prob" drives dynamic per-teacher weights; None (default) uses the static teacher_weights. student_vocab_size: Full student tokenizer vocab size (sizes the projection matrices’ V_s axis). Runtime-injected from len(student_tokenizer); not a user YAML knob. teacher_vocab_sizes / projection_matrix_paths / teacher_weights / teacher_gold_loss / teacher_xtoken_loss: Parallel per-teacher lists (one entry per teachers[i]), runtime-injected by the driver from the teachers config. A None projection path marks a same-tokenizer teacher (direct KL, no projection / alignment). Not user loss_fn keys in YAML.

Initialization

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

gold_loss: bool¶: None

xtoken_loss: bool¶: None

temperature: float¶: None

vocab_topk: int¶: None

uncommon_topk: int¶: None

reverse_kl: bool¶: None

exact_token_match_only: bool¶: None

kl_loss_weight: float¶: None

ce_loss_scale: float¶: None

dynamic_loss_scaling: bool¶: None

kd_loss_mode: str¶: None

normalize_teacher_by_vocab: bool¶: None

alpha: float¶: None

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

student_vocab_size: NotRequired[int]¶: None

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

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

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

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

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

class dockyard_rl.algorithms.loss.loss_functions.CrossTokenizerDistillationLossDataDict¶

Bases: typing.TypedDict

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

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

Every teacher: teacher_{i}_full_logits_ipc (node-local CUDA IPC) or teacher_{i}_full_logits_cross_cluster (data_plane) — rebuilt to full-vocab teacher logits by the loss-input keystone.
Cross-tokenizer teacher only: teacher_{i}_input_ids / teacher_{i}_token_mask [B, T_t] and alignment_{i}_* (pair_valid / pair_is_correct / chunk ids / partition masks / num_chunks).
Same-tokenizer teacher (projection_matrix_paths[i] is None): no teacher_{i}_input_ids / alignment_{i}_*; reuses the student tokenization (identity 1:1 alignment).

Initialization

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

input_ids: torch.Tensor¶: None

input_lengths: torch.Tensor¶: None

token_mask: torch.Tensor¶: None

sample_mask: torch.Tensor¶: None

class dockyard_rl.algorithms.loss.loss_functions.CrossTokenizerDistillationLossFn(cfg: dockyard_rl.algorithms.loss.loss_functions.CrossTokenizerDistillationLossConfig)¶

Bases: dockyard_rl.algorithms.loss.interfaces.LossFunction

Multi-teacher cross-tokenizer distillation loss.

Each teacher contributes a KD term; the per-teacher terms are aggregated per kd_loss_mode and combined with a single student next-token CE. The single-teacher path is just num_teachers == 1.

Per-teacher KD path is selected by that teacher’s tokenizer kind and its (gold_loss, xtoken_loss) flags:

Cross-tokenizer teacher (projection_matrix_paths[i] set):
- (False, False) -> P-KL: full-vocab projection KL (student logits mapped through the projection matrix M) over a microbatch-global top-k teacher subset.
- (True, False) -> gold-loss: KL on the exact-mapped common partition plus a sorted-L1 term on the uncommon tail.
- (True, True) -> gold-loss with the xtoken modifier: relaxed exact-map threshold (>= 0.6) + collision replacement.
Same-tokenizer teacher (projection_matrix_paths[i] is None): direct top-k per-position KL, no projection / no alignment.

(False, True) is rejected in __init__.

Aggregation (kd_loss_mode): sum (weighted sum; static teacher_weights or dynamic sum_weights_metric), averaged_logits (convex-average teacher logits then one KL; same-tokenizer teachers only), select_teacher (use only the lowest-CE teacher).

Inputs (via LossInputType.DISTILLATION_CROSS_TOKENIZER): logits (raw student logits), student_logits_contig (CP-relaid), the per-teacher teacher_full_logits_by_idx (rebuilt by the loss-input keystone) and aligns_by_idx (localized + next-token-shifted per teacher), and the TP/CP groups.

Initialization

loss_type¶: None

input_type¶: None

dockyard_rl.algorithms.loss.loss_functions¶

Module Contents¶

Classes¶

Functions¶

Data¶

API¶

`dockyard_rl.algorithms.loss.loss_functions`¶