`bridge.training.dpo`#

DPO pair loss and the logprob reduction it shares with the offline reference scorer.

Both jobs reduce per-token NLL to per-sequence sums through the same function, so a step-0 margin can only come from the models or the tokens, never from the reduction.

Module Contents#

Classes#

DPOLossConfig

DPO loss hyperparameters.

Functions#

`sequence_logprob_sums`	Reduce per-token NLL to per-row completion logprob sums and token counts.
`split_pair_rows`	Split a row-interleaved tensor into its (chosen, rejected) halves.
`dpo_loss`	DPO pair loss over one row-interleaved microbatch.

API#

bridge.training.dpo.sequence_logprob_sums( per_token_nll: torch.Tensor, loss_mask: torch.Tensor, ) → tuple[torch.Tensor, torch.Tensor]#

Reduce per-token NLL to per-row completion logprob sums and token counts.

Parameters:

per_token_nll – [rows, seq] next-token negative log-likelihood, as returned by model(..., labels=...).
loss_mask – [rows, seq]; nonzero marks completion-label positions.

Returns:

float32 -sum(nll * mask) per row and the number of counted positions.

Return type:

(logprob_sums, num_tokens), both [rows]

bridge.training.dpo.split_pair_rows( rows: torch.Tensor, ) → tuple[torch.Tensor, torch.Tensor]#: Split a row-interleaved tensor into its (chosen, rejected) halves.

class bridge.training.dpo.DPOLossConfig#

DPO loss hyperparameters.

Field names and defaults deliberately mirror NeMo-RL’s DPOLossConfig (nemo_rl/algorithms/loss/loss_functions.py) so runs are 1:1 comparable.

.. attribute:: reference_policy_kl_penalty

The DPO beta: temperature on the margin inside the sigmoid, i.e. the strength of the implicit KL leash to the reference policy.

.. attribute:: preference_loss_weight

Weight of the preference (-logsigmoid) term.

.. attribute:: sft_loss_weight

Weight of the auxiliary NLL term on the chosen response; 0 disables it.

.. attribute:: preference_average_log_probs

Average (instead of sum) per-token logprobs over completion tokens in the preference margins.

.. attribute:: sft_average_log_probs

Same, for the auxiliary SFT term.

reference_policy_kl_penalty: float#: 0.05

preference_loss_weight: float#: 1.0

sft_loss_weight: float#: 0.0

preference_average_log_probs: bool#: False

sft_average_log_probs: bool#: False

bridge.training.dpo.dpo_loss( config: bridge.training.dpo.DPOLossConfig, batch: dict[str, torch.Tensor], output_tensor: torch.Tensor, ) → tuple[torch.Tensor, torch.Tensor, dict[str, torch.Tensor]]#

DPO pair loss over one row-interleaved microbatch.

loss = sum over pairs of loss_multiplier * (w_p * -logsigmoid(beta * margin) + w_s * sft_nll) where margin = (pi_chosen - pi_rejected) - (ref_chosen - ref_rejected).

Stub pairs (loss_multiplier == 0) contribute nothing to the loss, the count, or any metric.