bridge.data.base

Shared dataset configuration and runtime build contracts.

Module Contents

Classes

DataloaderConfig	Base configuration for data loading.
DatasetBuildContext	Runtime metadata supplied to dataset builders by training setup.
DatasetProvider	Deprecated custom dataset-provider compatibility contract.

Functions

validate_declarative_mapping

Reject runtime objects from mappings stored in serializable configs.

API

class bridge.data.base.DataloaderConfig

Base configuration for data loading.

dataloader_type: Literal[single, cyclic, batch, external] | None

None

Dataloader type used by the training setup.

num_workers: int

2

Number of dataloader workers.

data_sharding: bool

True

Whether data is sharded across data-parallel ranks.

pin_memory: bool

True

Whether dataloaders pin host memory.

drop_last: bool

True

Whether dataloaders drop the last incomplete batch.

persistent_workers: bool

True

Whether dataloader workers persist between iterations.

trust_remote_code: bool | None

None

Whether remote code is trusted for a configured Hugging Face path.

finalize() → None

Finalize dataloader field constraints.

class bridge.data.base.DatasetBuildContext

Runtime metadata supplied to dataset builders by training setup.

train_samples: int

None

valid_samples: int

None

test_samples: int

None

tokenizer: megatron.bridge.training.tokenizers.tokenizer.MegatronTokenizer | None

None

pg_collection: megatron.core.process_groups_config.ProcessGroupCollection | None

None

class bridge.data.base.DatasetProvider

Bases: bridge.data.base.DataloaderConfig, abc.ABC

Deprecated custom dataset-provider compatibility contract.

abstractmethod build_datasets(

context: bridge.data.base.DatasetBuildContext,

) → tuple[Any | None, Any | None, Any | None]

Build train, validation, and test datasets.

Parameters:

context – Runtime sample counts, tokenizer, and process groups.

Returns:

Train, validation, and test datasets. Optional splits may be None.

bridge.data.base.validate_declarative_mapping(

value: dict[str, Any] | None,

*,

field_name: str,

) → None

Reject runtime objects from mappings stored in serializable configs.