bridge.data.builders.dpo

ConfigContainer.dataset entry for DPO preference training.

DPODatasetConfig + DPODatasetBuilder + the registered dpo_train_valid_test_datasets_provider follow the config-and-builder dataset architecture (registered in data/utils.py); loader construction is the DPO branch of data/loaders.py.

Module Contents

Classes

DPODatasetConfig	Preference dataset plus its offline-scored ref-logprob artifact.
DPODatasetBuilder	Build the runtime preference-pair dataset with reference logprobs attached.

Functions

hf_chat_tokenizer	The HF tokenizer behind the run tokenizer; pairs are rendered with its chat template.
build_preference_split	Load one split’s rows and wrap them as a pair dataset.
dpo_train_valid_test_datasets_provider	Build DPO preference datasets through the canonical runtime builder.

API

class bridge.data.builders.dpo.DPODatasetConfig

Bases: megatron.bridge.data.base.DataloaderConfig

Preference dataset plus its offline-scored ref-logprob artifact.

source is an HFDatasetSourceConfig: a Hub dataset, or local or msc:// files through a file loader (path_or_dataset="json" with the files in load_kwargs["data_files"]). The source identity, the run tokenizer, and seq_length must equal the scoring run’s values, which dpo_train checks against scoring_metadata.json (where the length is recorded under its dataset-level name, max_seq_length).

A validation split is optional and mirrors the train shape: one validation_source plus its own scored validation_ref_artifact (the tokenizer and seq_length are shared with training, so the validation split must be scored with the same values). Like the SFT dataset configs, a configured validation split is built only when the run actually validates (validation.eval_iters/eval_interval > 0).

Batch sizes stay row-denominated; pair halving happens only in the loader builder.

seq_length: int

None

Truncation ceiling for tokenized pairs.

source: megatron.bridge.data.sources.hf.HFDatasetSourceConfig

None

ref_artifact: str | None

None

Scorer artifact directory or URL. Required to build datasets; only source-only helpers (load_source) may leave it unset.

validation_ref_artifact: str | None

None

Scorer artifact for the validation split. Required when a validation source is set.

validation_source: megatron.bridge.data.sources.hf.HFDatasetSourceConfig | None

None

validation_num_pairs: int

0

chosen_key: str

‘chosen’

rejected_key: str

‘rejected’

prompt_key: str | None

None

tools_key: str | None

None

Field holding the row’s tool schemas (list or JSON string), rendered into the prompt by the chat template.

chat_template_kwargs: dict[str, Any] | None

None

Extra chat-template kwargs for every row, e.g. {"truncate_history_thinking": false} so multi-turn trajectories keep the reasoning of turns before the last user message (Nemotron 3 / Qwen3 strip it).

num_pairs: int

0

dataloader_type: Literal[batch] | None

‘batch’

Pinned: the batch sampler’s cyclic wrap provides epochs beyond one dataset pass.

pad_seq_length_to_mult: int

1

Round each batch’s padded row length up to this multiple. Sequence parallelism needs TP | seq_len; ConfigContainer.validate derives it from the model config, the scorer from the recipe’s TP and sequence parallelism.

shuffle: bool

True

static resolve_source(

source: megatron.bridge.data.sources.hf.HFDatasetSourceConfig,

*,

field: str,

) → megatron.bridge.data.sources.hf.HFDatasetSourceConfig

Rebuild a flattened source if needed, then validate it.

validate() → None

Resolve and validate the sources; pin the dataloader type.

finalize() → None

Finalize dataloader settings and validate. ref_artifact may be local or msc://.

property has_validation_split: bool

Whether a validation source is configured.

static source_identity(

source: megatron.bridge.data.sources.hf.HFDatasetSourceConfig,

) → tuple[str, str | None]

(dataset, split) as recorded in and checked against scoring_metadata.json.

File-backed loaders (json, parquet, …) are told apart by their data_files.

load_source(split: Literal[train, validation] = 'train')

Load the raw preference rows of split, truncated to that split’s num_pairs.

The offline scorer calls this too: both jobs must see the same rows in the same order for pair_id to line up between the artifact and training.

static _load_rows(

source: megatron.bridge.data.sources.hf.HFDatasetSourceConfig,

num_pairs: int,

)

bridge.data.builders.dpo.hf_chat_tokenizer(tokenizer: Any) → Any

The HF tokenizer behind the run tokenizer; pairs are rendered with its chat template.

bridge.data.builders.dpo.build_preference_split(

config: bridge.data.builders.dpo.DPODatasetConfig,

split: Literal[train, validation],

tokenizer,

ref_artifact: str | None = None,

) → megatron.bridge.data.datasets.preference_pair.PreferencePairDataset

Load one split’s rows and wrap them as a pair dataset.

The scorer builds without ref_artifact; the trainer supplies the split’s artifact, which must cover every loaded pair.

class bridge.data.builders.dpo.DPODatasetBuilder(config: bridge.data.builders.dpo.DPODatasetConfig)

Build the runtime preference-pair dataset with reference logprobs attached.

Initialization

build(

context: megatron.bridge.data.base.DatasetBuildContext,

) → tuple[megatron.bridge.data.datasets.preference_pair.PreferencePairDataset, megatron.bridge.data.datasets.preference_pair.PreferencePairDataset | None, None]

bridge.data.builders.dpo.dpo_train_valid_test_datasets_provider(

train_val_test_num_samples: list[int],

dataset_config: bridge.data.builders.dpo.DPODatasetConfig,

tokenizer: MegatronTokenizer | None = None,

pg_collection: ProcessGroupCollection | None = None,

) → tuple[megatron.bridge.data.datasets.preference_pair.PreferencePairDataset, megatron.bridge.data.datasets.preference_pair.PreferencePairDataset | None, None]

Build DPO preference datasets through the canonical runtime builder.