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.

dataloader_save: str | None

None

Directory to save dataloader stream-position state into during checkpointing (currently only Energon’s SavableDataLoader), so a resumed run continues over the same data instead of restarting from an arbitrary position. When None (the default) and the dataloader supports state saving, it is colocated under an energon subdirectory of checkpoint.save. Has no effect for dataloaders that do not support state saving.

dataloader_load: str | None

None

Directory to restore dataloader stream-position state from on resume. When None (the default), it is resolved to the energon subdirectory of whichever checkpoint is actually loaded – checkpoint.save for a non-persistent or local checkpoint, checkpoint.load for a persistent one, or the parent of a directly specified iteration directory. If that directory does not exist (e.g. a checkpoint saved before this feature) the dataloader starts fresh; if it exists but the current rank’s state file is missing, resume fails loudly rather than silently changing the data order.

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.