bridge.training.dpo_step

DPO forward step: pair-batch fetch, the model call, and the checked loss closure.

Module Contents

Functions

validate_dpo_batch	Fail loudly if a batch is missing DPO keys (e.g. built without ref logprobs).
trim_dpo_batch_padding	Drop all-pad tail columns so the microbatch width matches a standalone collate of its own pairs.
get_dpo_batch	Fetch one row-interleaved pair batch, move it to device, and trim pad tails.
_loss_not_on_this_stage
dpo_loss_with_checks	dpo_loss followed by the rerun-state-machine loss checks.
forward_step	DPO forward training step.

Data

DPO_BATCH_KEYS
_SEQ_DIM_KEYS

API

bridge.training.dpo_step.DPO_BATCH_KEYS

(‘tokens’, ‘labels’, ‘loss_mask’, ‘position_ids’, ‘pair_id’, ‘loss_multiplier’, ‘ref_logprob_sum’, ‘…

bridge.training.dpo_step.validate_dpo_batch(

batch: dict[str, torch.Tensor],

) → dict[str, torch.Tensor]

Fail loudly if a batch is missing DPO keys (e.g. built without ref logprobs).

bridge.training.dpo_step._SEQ_DIM_KEYS

(‘tokens’, ‘labels’, ‘loss_mask’, ‘position_ids’)

bridge.training.dpo_step.trim_dpo_batch_padding(

batch: dict[str, torch.Tensor],

pad_to_multiple: int = 1,

) → dict[str, torch.Tensor]

Drop all-pad tail columns so the microbatch width matches a standalone collate of its own pairs.

pad_to_multiple must equal the collate’s pad_seq_length_to_mult.

bridge.training.dpo_step.get_dpo_batch(

data_iterator: collections.abc.Iterator[dict[str, torch.Tensor]],

pad_to_multiple: int = 1,

) → dict[str, torch.Tensor]

Fetch one row-interleaved pair batch, move it to device, and trim pad tails.

bridge.training.dpo_step._loss_not_on_this_stage(output_tensor: torch.Tensor) → None

bridge.training.dpo_step.dpo_loss_with_checks(

loss_config: megatron.bridge.training.dpo.DPOLossConfig,

batch: dict[str, torch.Tensor],

output_tensor: torch.Tensor,

*,

check_for_nan_in_loss: bool,

check_for_spiky_loss: bool,

) → tuple[torch.Tensor, torch.Tensor, dict[str, torch.Tensor]]

dpo_loss followed by the rerun-state-machine loss checks.

bridge.training.dpo_step.forward_step(

state: megatron.bridge.training.state.GlobalState,

data_iterator: collections.abc.Iterator[dict[str, torch.Tensor]],

model: megatron.core.models.gpt.GPTModel,

) → tuple[torch.Tensor, collections.abc.Callable]

DPO forward training step.

Parameters:

• state – Global state for the run; state.cfg.dpo holds the DPOLossConfig.

• data_iterator – Yields preference_collate_fn batches (rows: chosen@even/rejected@odd).

• model – The GPT model chunk.

Returns:

The [rows, seq] per-token NLL tensor (hidden states on a non-final pipeline stage) and the loss closure, which the schedule only invokes on the last stage.