bridge.models.bagel.data.packing

Module Contents

Classes

BagelPacker

Pack cooked Energon samples in BAGEL’s official streaming order.

Functions

_checkpoint_safe_value	Copy a BAGEL sample value into the restricted checkpoint type set.
_serialize_buffered_sample	Serialize only fields the packer needs after restoring a buffered sample.
_restore_buffered_sample	Reconstruct a trusted BAGEL sample after restricted checkpoint loading.
_patchify	Patchify an image with BAGEL’s channel-last patch ordering.
_position_ids	Build BAGEL’s extrapolated flattened image position IDs.
_attention_mask	Build one official-style nested attention mask.

Data

logger

API

bridge.models.bagel.data.packing.logger

‘getLogger(…)’

bridge.models.bagel.data.packing._checkpoint_safe_value(value: object) → object

Copy a BAGEL sample value into the restricted checkpoint type set.

bridge.models.bagel.data.packing._serialize_buffered_sample(

sample: megatron.bridge.models.bagel.data.energon.BagelSample,

) → dict[str, object]

Serialize only fields the packer needs after restoring a buffered sample.

Energon source metadata can contain EPath objects whose pickle restore hook resolves a storage client. Source metadata is diagnostic and is not consumed by :class:BagelPacker, so it is intentionally omitted instead of widening the shared restricted-unpickler allowlist.

bridge.models.bagel.data.packing._restore_buffered_sample(

state: object,

) → megatron.bridge.models.bagel.data.energon.BagelSample

Reconstruct a trusted BAGEL sample after restricted checkpoint loading.

bridge.models.bagel.data.packing._patchify(image: torch.Tensor, patch_size: int) → torch.Tensor

Patchify an image with BAGEL’s channel-last patch ordering.

bridge.models.bagel.data.packing._position_ids(

height: int,

width: int,

patch_size: int,

max_patches_per_side: int,

) → torch.Tensor

Build BAGEL’s extrapolated flattened image position IDs.

bridge.models.bagel.data.packing._attention_mask(

split_lens: list[int],

attn_modes: list[str],

) → torch.Tensor

Build one official-style nested attention mask.

class bridge.models.bagel.data.packing.BagelPacker(

group_iters: collections.abc.Sequence[collections.abc.Iterator[megatron.bridge.models.bagel.data.energon.BagelSample]],

group_weights: collections.abc.Sequence[float],

is_mandatory: collections.abc.Sequence[bool],

special_tokens: collections.abc.Mapping[str, int],

*,

expected_num_tokens: int = 32768,

max_num_tokens_per_sample: int = 16384,

max_num_tokens: int = 36864,

prefer_buffer_before: int = 16384,

max_buffer_size: int = 50,

text_cond_dropout_prob: float = 0.1,

vit_cond_dropout_prob: float = 0.4,

vae_cond_dropout_prob: float = 0.1,

vae_image_downsample: int = 16,

max_latent_size: int = 32,

vit_patch_size: int = 14,

max_num_patch_per_side: int = 70,

)

Pack cooked Energon samples in BAGEL’s official streaming order.

Initialization

Configure the official non-Flex BAGEL packing path.

static _new_status() → dict[str, Any]

Create one empty official packing accumulator.

static _source_id(

sample: megatron.bridge.models.bagel.data.energon.BagelSample,

) → dict[str, object]

Return the canonical source coordinate carried by a cooked sample.

__iter__() → Self

Return this stateful packed-batch iterator.

__next__() → dict[str, object]

Return the next batch with official mandatory and FIFO-buffer behavior.

state_dict() → dict[str, object]

Capture packing, buffering, and process RNG state at a batch boundary.

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

Restore packing, buffering, and process RNG state.

_pack_sequence(

sample: megatron.bridge.models.bagel.data.energon.BagelSample,

status: dict[str, Any],

) → None

Append one cooked sample using BAGEL’s dropout and timestep calls.

static _to_tensor(status: dict[str, Any]) → dict[str, object]

Convert one completed accumulator to official packed tensors.