bridge.models.bagel.data.external

Module Contents

Classes

BagelRNGIterator	Run the official packer with data-only process RNG state.
BagelExternalLoader	Expose packed BAGEL batches through Megatron’s external-loader contract.
BagelMegatronMIMODatasetProvider	Pass one BAGEL packed-batch iterator into MegatronMIMO unchanged.

Functions

_reject_collate

Reject collation because Megatron external loaders bypass it.

API

bridge.models.bagel.data.external._reject_collate(_: object) → object

Reject collation because Megatron external loaders bypass it.

class bridge.models.bagel.data.external.BagelRNGIterator(

iterator: collections.abc.Iterator[dict[str, object]],

rank_seed: int,

)

Run the official packer with data-only process RNG state.

Initialization

Initialize the independent RNG states of BAGEL’s one DataLoader worker.

__iter__() → Self

Return this stateful RNG-isolated iterator.

__next__() → dict[str, object]

Advance the packer without changing the model process RNG.

state_dict() → dict[str, object]

Capture the wrapped packer and isolated data RNG.

load_state_dict(state: collections.abc.Mapping[str, object]) → None

Restore the wrapped packer without leaking its RNG to the model.

class bridge.models.bagel.data.external.BagelExternalLoader(

batches: collections.abc.Iterable[dict[str, object]],

*,

length: int,

stateful_loaders: collections.abc.Sequence[object] = (),

)

Expose packed BAGEL batches through Megatron’s external-loader contract.

Initialization

Wrap a packed batch stream with the length required by loader setup.

__iter__() → Self

Return this next-compatible external loader.

__next__() → dict[str, object]

Return the next already-packed BAGEL batch.

__len__() → int

Return the configured number of training batches.

save_state() → dict[str, object]

Capture the Energon readers, packer, and external-loader position.

restore_state(state: collections.abc.Mapping[str, object]) → None

Restore the Energon readers, packer, and external-loader position.

class bridge.models.bagel.data.external.BagelMegatronMIMODatasetProvider

Bases: megatron.bridge.data.megatron_mimo.base_provider.MegatronMIMODatasetProvider

Pass one BAGEL packed-batch iterator into MegatronMIMO unchanged.

train_loader: bridge.models.bagel.data.external.BagelExternalLoader

None

dataloader_type: Literal[external]

‘external’

num_workers: int

0

persistent_workers: bool

False

build_datasets(

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

) → tuple[bridge.models.bagel.data.external.BagelExternalLoader | None, None, None]

Return only the external training loader requested by the build context.

get_collate_fn() → collections.abc.Callable

Return a guard callable; external loader setup never invokes it.