nemo_automodel.components.datasets.multimodal

BAGEL-style multimodal data pipeline for packed three-group training.

Ports the subset of upstream BAGEL data needed to feed BAGEL training from fully AM-native code: VLM SFT, T2I pretrain, and unified image editing. The packed batch schema is shared by Stage 1 and Stage 2; whether VAE / flow-matching tensors are consumed is controlled by the model stage.

Submodules

Package Contents

Classes

Name	Description
`BagelDatasetConfig`	Declarative configuration for the packed BAGEL multimodal dataset.
`DistributedIterableDataset`	Base class for rank/worker-aware iterable datasets.
`FrameSampler`	Callable that returns a list of PIL frames for a given file path.
`ImageTransform`	Full BAGEL image transform: resize + to_tensor + normalize.
`InterleavedBaseIterableDataset`	Builder mixin for interleaved image/text/video sequence plans.
`MaxLongEdgeMinShortEdgeResize`	Resize so longest/shortest edges stay within bounds and both edges are stride-divisible.
`PackedDataset`	Greedy pack of samples drawn from weighted groups into token-budgeted batches.
`ParquetStandardIterableDataset`	Base class: iterate per-(file, row_group) across a list of parquet shards.
`SftJSONLIterableDataset`	Iterable over conversation JSONL rows where each row may reference images.
`SimpleCustomBatch`	Pass-through wrapper around one packed batch from `PackedDataset`.
`T2IIterableDataset`	Iterable over a parquet-sharded (image, captions) text-to-image dataset.
`UnifiedEditIterableDataset`	Image-editing dataset: `(input, instruction, output)` chains over parquet.

Functions

Name	Description
`add_special_tokens`	Add BAGEL’s four special tokens to a Qwen2 tokenizer and return the ids.
`bagel_packed_collate_fn`	Canonical name in AM’s collate-fn registry.
`collate_wrapper`	Return the BAGEL-style identity collate (wraps a single packed dict).
`get_flattened_position_ids_extrapolate`	Return flattened 2D patch position IDs by direct grid extrapolation.
`get_flattened_position_ids_interpolate`	Return flattened 2D patch position IDs by interpolating to the max grid.
`get_parquet_data_paths`	Return a flat list of parquet file paths sharded across ranks.
`init_arrow_pf_fs`	Return a pyarrow filesystem for `parquet_file_path`.
`len2weight`	Convert a sequence length into BAGEL’s per-sample loss weight.
`make_bagel_multimodal_dataset`	Build a BAGEL packed dataset for Stage 1 or Stage 2 training.
`patchify`	Patchify a CxHxW image tensor into (H/p * W/p, ppC).
`pil_img2rgb`	Convert a PIL image to RGB, compositing transparent pixels on white.
`prepare_attention_mask_per_sample`	Build a per-sample additive float mask honoring causal / full / noise splits.

Data

DATASET_REGISTRY

DEFAULT_DATASET_INFO

DataConfig

API

class nemo_automodel.components.datasets.multimodal.datasets.BagelDatasetConfig(
    grouped_datasets: collections.abc.Mapping[str, object] | None = None,
    grouped_datasets_path: str | None = None,
    dataset_info: collections.abc.Mapping[str, object] | None = None,
    dataset_info_path: str | None = None,
    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,
    interpolate_pos: bool = False,
    use_flex: bool = False,
    data_seed: int = 42,
    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
)

Dataclass

Declarative configuration for the packed BAGEL multimodal dataset.

data_seed

int = 42

dataset_info

Mapping[str, object] | None = None

dataset_info_path

str | None = None

expected_num_tokens

int = 32768

grouped_datasets

Mapping[str, object] | None = None

grouped_datasets_path

str | None = None

interpolate_pos

bool = False

max_buffer_size

int = 50

max_latent_size

int = 32

max_num_patch_per_side

int = 70

max_num_tokens

int = 36864

max_num_tokens_per_sample

int = 16384

prefer_buffer_before

int = 16384

text_cond_dropout_prob

float = 0.1

use_flex

bool = False

vae_cond_dropout_prob

float = 0.1

vae_image_downsample

int = 16

vit_cond_dropout_prob

float = 0.4

vit_patch_size

int = 14

nemo_automodel.components.datasets.multimodal.datasets.BagelDatasetConfig._resolve_mapping(
    inline: collections.abc.Mapping[str, object] | None,
    path: str | None,
    name: str
) -> dict[str, object]

staticmethod

Resolve one inline or YAML/JSON mapping without mutating config state.

nemo_automodel.components.datasets.multimodal.datasets.BagelDatasetConfig.build(
    tokenizer: 'PreTrainedTokenizerBase',
    special_tokens: collections.abc.Mapping[str, object],
    rank: int,
    world_size: int,
    num_workers: int,
    global_seed: int,
    data_status: object | None = None
) -> 'PackedDataset'

Build and initialize the packed BAGEL dataset.

Parameters:

tokenizer

'PreTrainedTokenizerBase'

Runtime tokenizer used by the packed dataset.

special_tokens

Mapping[str, object]

Runtime BAGEL special-token mapping.

rank

int

Global distributed rank used for source-data sharding.

world_size

int

Global distributed world size.

num_workers

int

Dataloader worker count used to shard source files.

global_seed

int

Ranked-training base seed used by worker-start reseeding.

data_status

object | NoneDefaults to None

Optional dataset resume state.

Returns: 'PackedDataset'

Initialized packed BAGEL iterable dataset.

class nemo_automodel.components.datasets.multimodal.distributed_iterable.DistributedIterableDataset(
    dataset_name,
    local_rank = 0,
    world_size = 1,
    num_workers = 8
)

Bases: IterableDataset

Base class for rank/worker-aware iterable datasets.

Owns a private rng used only to shuffle file paths deterministically in set_epoch — NOT used for per-sample randomness. Per-sample randomness still goes through the Python global random module (see packing for the reseed hook).

_drop_counters

= {}

data_paths_per_rank

= []

num_files_per_rank

= 0

rng

= random.Random()

nemo_automodel.components.datasets.multimodal.distributed_iterable.DistributedIterableDataset.__iter__()

nemo_automodel.components.datasets.multimodal.distributed_iterable.DistributedIterableDataset._get_worker_data_status(
    worker_id
)

nemo_automodel.components.datasets.multimodal.distributed_iterable.DistributedIterableDataset._log_drop(
    reason,
    message,
    args = (),
    every = 100,
    exc_info = False
)

nemo_automodel.components.datasets.multimodal.distributed_iterable.DistributedIterableDataset._set_worker_resume_data_status(
    worker_id,
    status
)

nemo_automodel.components.datasets.multimodal.distributed_iterable.DistributedIterableDataset.get_data_paths(
    args = (),
    kwargs = {}
)

nemo_automodel.components.datasets.multimodal.distributed_iterable.DistributedIterableDataset.get_data_paths_per_worker()

nemo_automodel.components.datasets.multimodal.distributed_iterable.DistributedIterableDataset.load_state_dict(
    state_dict
)

nemo_automodel.components.datasets.multimodal.distributed_iterable.DistributedIterableDataset.set_data_status(
    data_status
)

nemo_automodel.components.datasets.multimodal.distributed_iterable.DistributedIterableDataset.set_epoch(
    seed = 42
)

nemo_automodel.components.datasets.multimodal.distributed_iterable.DistributedIterableDataset.state_dict()

class nemo_automodel.components.datasets.multimodal.video.FrameSampler(
    max_num_frames = -1,
    min_num_frames = 8,
    sample = 'rand'
)

Callable that returns a list of PIL frames for a given file path.

nemo_automodel.components.datasets.multimodal.video.FrameSampler.__call__(
    file_name
)

class nemo_automodel.components.datasets.multimodal.transforms.ImageTransform(
    max_image_size,
    min_image_size,
    image_stride,
    max_pixels = 14 * 14 * 9 * 1024,
    image_mean = (0.5, 0.5, 0.5),
    image_std = (0.5, 0.5, 0.5)
)

Full BAGEL image transform: resize + to_tensor + normalize.

Used for both ViT input (stride=14) and VAE input (stride=16, via separate instances). stride is exposed as an attribute so the dataset can compute patch counts without knowing the transform class.

normalize_transform

resize_transform

to_tensor_transform

= tv_transforms.ToTensor()

nemo_automodel.components.datasets.multimodal.transforms.ImageTransform.__call__(
    img,
    img_num = 1
)

class nemo_automodel.components.datasets.multimodal.interleave.InterleavedBaseIterableDataset()

Bases: DistributedIterableDataset

Builder mixin for interleaved image/text/video sequence plans.

Subclasses still provide __init__ + __iter__ + parse_row (via ParquetStandardIterableDataset); this class only holds the per-item append helpers used inside parse_row.

nemo_automodel.components.datasets.multimodal.interleave.InterleavedBaseIterableDataset._add_image(
    data,
    image,
    need_loss,
    need_vae,
    need_vit,
    enable_cfg = True
)

nemo_automodel.components.datasets.multimodal.interleave.InterleavedBaseIterableDataset._add_text(
    data,
    text,
    need_loss,
    enable_cfg = True
)

nemo_automodel.components.datasets.multimodal.interleave.InterleavedBaseIterableDataset._add_video(
    data,
    frames,
    frame_indexes,
    need_loss,
    need_vae,
    enable_cfg = True
)

nemo_automodel.components.datasets.multimodal.interleave.InterleavedBaseIterableDataset._init_data()

class nemo_automodel.components.datasets.multimodal.transforms.MaxLongEdgeMinShortEdgeResize(
    max_size: int,
    min_size: int,
    stride: int,
    max_pixels: int,
    interpolation = None,
    antialias = True
)

Bases: Module

Resize so longest/shortest edges stay within bounds and both edges are stride-divisible.

Parameters:

max_size

int

Maximum size for the longest edge.

min_size

int

Minimum size for the shortest edge.

stride

int

Value both edges must be divisible by (ViT patch size).

max_pixels

int

Maximum total pixels for the full image.

interpolation

Defaults to None

Torchvision interpolation mode (default bicubic).

antialias

Defaults to True

Whether to apply antialiasing.

nemo_automodel.components.datasets.multimodal.transforms.MaxLongEdgeMinShortEdgeResize._apply_scale(
    width,
    height,
    scale
)

nemo_automodel.components.datasets.multimodal.transforms.MaxLongEdgeMinShortEdgeResize._make_divisible(
    value,
    stride
)

nemo_automodel.components.datasets.multimodal.transforms.MaxLongEdgeMinShortEdgeResize.forward(
    img,
    img_num = 1
)

class nemo_automodel.components.datasets.multimodal.packing.PackedDataset(
    data_config: 'BagelDatasetConfig',
    tokenizer: 'PreTrainedTokenizerBase',
    special_tokens: collections.abc.Mapping[str, object],
    local_rank: int,
    world_size: int,
    num_workers: 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,
    interpolate_pos: bool = False,
    use_flex: bool = False,
    data_status: object | None = None,
    dataset_info: collections.abc.Mapping[str, object] | None = None,
    global_seed: int | None = None
)

Bases: IterableDataset

Greedy pack of samples drawn from weighted groups into token-budgeted batches.

The dataset reseeds at iterator start so AM sees a deterministic BAGEL-compatible packed-data stream regardless of earlier RNG consumption during model construction or checkpoint loading.

_drop_counters

= {}

_resume_buffer

= []

_resume_sequence_status

= self.set_sequence_status()

_yielded_batches

= 0

nemo_automodel.components.datasets.multimodal.packing.PackedDataset.__iter__()

nemo_automodel.components.datasets.multimodal.packing.PackedDataset._grouped_dataset_state_dicts()

nemo_automodel.components.datasets.multimodal.packing.PackedDataset._load_grouped_dataset_state_dicts(
    states
)

nemo_automodel.components.datasets.multimodal.packing.PackedDataset._load_rng_state_dict(
    state
)

nemo_automodel.components.datasets.multimodal.packing.PackedDataset._log_drop(
    reason,
    message,
    args = (),
    every = 100
)

nemo_automodel.components.datasets.multimodal.packing.PackedDataset._rng_state_dict()

nemo_automodel.components.datasets.multimodal.packing.PackedDataset._set_resume_point(
    buffer,
    yielded_batches
)

nemo_automodel.components.datasets.multimodal.packing.PackedDataset.build_datasets(
    datasets_metainfo,
    data_status
)

nemo_automodel.components.datasets.multimodal.packing.PackedDataset.load_state_dict(
    state_dict
)

nemo_automodel.components.datasets.multimodal.packing.PackedDataset.pack_sequence(
    sample,
    sequence_status
)

nemo_automodel.components.datasets.multimodal.packing.PackedDataset.set_epoch(
    seed
)

nemo_automodel.components.datasets.multimodal.packing.PackedDataset.set_sequence_status()

nemo_automodel.components.datasets.multimodal.packing.PackedDataset.state_dict()

nemo_automodel.components.datasets.multimodal.packing.PackedDataset.to_tensor(
    sequence_status
)

class nemo_automodel.components.datasets.multimodal.interleave.ParquetStandardIterableDataset(
    dataset_name,
    transform,
    tokenizer,
    vit_transform,
    data_dir_list,
    num_used_data,
    parquet_info,
    local_rank = 0,
    world_size = 1,
    num_workers = 8,
    data_status = None
)

Bases: DistributedIterableDataset

Base class: iterate per-(file, row_group) across a list of parquet shards.

Subclasses override parse_row to turn one pandas row into the dict schema consumed by .packing.PackedDataset.

data_paths

nemo_automodel.components.datasets.multimodal.interleave.ParquetStandardIterableDataset.__iter__()

nemo_automodel.components.datasets.multimodal.interleave.ParquetStandardIterableDataset.get_data_paths(
    data_dir_list,
    num_used_data,
    parquet_info
)

nemo_automodel.components.datasets.multimodal.interleave.ParquetStandardIterableDataset.parse_row(
    row
)

class nemo_automodel.components.datasets.multimodal.datasets.SftJSONLIterableDataset(
    dataset_name,
    transform,
    tokenizer,
    frame_sampler,
    jsonl_path_list,
    data_dir_list,
    num_used_data,
    local_rank = 0,
    world_size = 1,
    num_workers = 8,
    data_status = None,
    shuffle_lines = False,
    shuffle_seed = 0
)

Bases: DistributedIterableDataset

Iterable over conversation JSONL rows where each row may reference images.

data_paths

nemo_automodel.components.datasets.multimodal.datasets.SftJSONLIterableDataset.__iter__()

nemo_automodel.components.datasets.multimodal.datasets.SftJSONLIterableDataset.change_format(
    data,
    num_images
)

nemo_automodel.components.datasets.multimodal.datasets.SftJSONLIterableDataset.get_data_paths(
    jsonl_path_list,
    data_dir_list,
    num_used_data,
    shuffle_lines,
    shuffle_seed
)

class nemo_automodel.components.datasets.multimodal.collate_fns.SimpleCustomBatch(
    batch
)

Pass-through wrapper around one packed batch from PackedDataset.

attn_modes

= data['attn_modes']

batch_data_indexes

= data['batch_data_indexes']

ce_loss_indexes

= data['ce_loss_indexes']

ce_loss_weights

= data['ce_loss_weights']

mse_loss_indexes

= data['mse_loss_indexes']

nested_attention_masks

= data['nested_attention_masks']

packed_label_ids

= data['packed_label_ids']

packed_latent_position_ids

= data['packed_latent_position_ids']

packed_position_ids

= data['packed_position_ids']

packed_text_ids

= data['packed_text_ids']

packed_text_indexes

= data['packed_text_indexes']

packed_timesteps

= data['packed_timesteps']

packed_vae_token_indexes

= data['packed_vae_token_indexes']

packed_vit_position_ids

= data['packed_vit_position_ids']

packed_vit_token_indexes

= data['packed_vit_token_indexes']

packed_vit_tokens

= data['packed_vit_tokens']

padded_images

= data['padded_images']

patchified_vae_latent_shapes

= data['patchified_vae_latent_shapes']

sample_lens

= data['sample_lens']

sequence_length

= data['sequence_length']

split_lens

= data['split_lens']

use_flex

= 'nested_attention_masks' not in data.keys()

vit_token_seqlens

= data['vit_token_seqlens']

nemo_automodel.components.datasets.multimodal.collate_fns.SimpleCustomBatch.cuda(
    device
)

nemo_automodel.components.datasets.multimodal.collate_fns.SimpleCustomBatch.pin_memory()

nemo_automodel.components.datasets.multimodal.collate_fns.SimpleCustomBatch.to_dict()

class nemo_automodel.components.datasets.multimodal.datasets.T2IIterableDataset(
    dataset_name,
    transform,
    tokenizer,
    data_dir_list,
    num_used_data,
    local_rank = 0,
    world_size = 1,
    num_workers = 8,
    data_status = None
)

Bases: DistributedIterableDataset

Iterable over a parquet-sharded (image, captions) text-to-image dataset.

Each parquet row is expected to carry:

image: raw image bytes.
captions: JSON-encoded {key: caption} dict; one caption is sampled uniformly per row.

The yielded dict carries only one image tensor (the VAE input). Stage 1 models ignore the VAE branch; Stage 2 consumes it via the PackedDataset vae_image branch + flow-matching loss. No VAE encoding happens in the data pipeline — just tensor preparation.

data_paths

= self.get_data_paths(data_dir_list, num_used_data)

nemo_automodel.components.datasets.multimodal.datasets.T2IIterableDataset.__iter__()

nemo_automodel.components.datasets.multimodal.datasets.T2IIterableDataset.get_data_paths(
    data_dir_list,
    num_used_data
)

class nemo_automodel.components.datasets.multimodal.interleave.UnifiedEditIterableDataset()

Bases: InterleavedBaseIterableDataset, ParquetStandardIterableDataset

Image-editing dataset: (input, instruction, output) chains over parquet.

Row schema (upstream BAGEL seedxedit_multi + compatibles): image_list: list of raw image bytes (at least 2). instruction_list: list of lists; instruction_list[i] is a set of equivalent phrasings for the edit that turns image_list[i] into image_list[i+1].

nemo_automodel.components.datasets.multimodal.interleave.UnifiedEditIterableDataset.parse_row(
    row
)

nemo_automodel.components.datasets.multimodal.utils.add_special_tokens(
    tokenizer
)

Add BAGEL’s four special tokens to a Qwen2 tokenizer and return the ids.

Returns:

tuple (tokenizer, new_token_ids, num_new_tokens) where

nemo_automodel.components.datasets.multimodal.collate_fns.bagel_packed_collate_fn(
    batch
)

Canonical name in AM’s collate-fn registry.

nemo_automodel.components.datasets.multimodal.collate_fns.collate_wrapper()

Return the BAGEL-style identity collate (wraps a single packed dict).

nemo_automodel.components.datasets.multimodal.utils.get_flattened_position_ids_extrapolate(
    img_h,
    img_w,
    patch_size,
    max_num_patches_per_side
)

Return flattened 2D patch position IDs by direct grid extrapolation.

nemo_automodel.components.datasets.multimodal.utils.get_flattened_position_ids_interpolate(
    img_h,
    img_w,
    patch_size,
    max_num_patches_per_side
)

Return flattened 2D patch position IDs by interpolating to the max grid.

nemo_automodel.components.datasets.multimodal.parquet_utils.get_parquet_data_paths(
    data_dir_list,
    num_sampled_data_paths,
    rank = 0,
    world_size = 1
)

Return a flat list of parquet file paths sharded across ranks.

Directories are split across ranks via chunk_size = ceil(num_dirs / world_size). Each rank lists its local directories, repeats the file list to reach num_sampled_data_paths per directory, then all-gathers across ranks so every rank ends up with the same combined list.

nemo_automodel.components.datasets.multimodal.parquet_utils.init_arrow_pf_fs(
    parquet_file_path
)

Return a pyarrow filesystem for parquet_file_path.

pyarrow is imported lazily because not every AM install carries it, and the import is only needed when an actually-parquet-backed dataset (T2I, UnifiedEdit) is constructed.

nemo_automodel.components.datasets.multimodal.utils.len2weight(
    x,
    loss_reduction = 'square'
)

Convert a sequence length into BAGEL’s per-sample loss weight.

nemo_automodel.components.datasets.multimodal.datasets.make_bagel_multimodal_dataset(
    tokenizer,
    special_tokens,
    grouped_datasets,
    local_rank: int = 0,
    world_size: int = 1,
    num_workers: int = 4,
    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,
    interpolate_pos: bool = False,
    use_flex: bool = False,
    data_status = None,
    data_seed: int = 42,
    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,
    dataset_info = None
)

Build a BAGEL packed dataset for Stage 1 or Stage 2 training.

grouped_datasets is the YAML dict produced by data/configs/*.yaml. dataset_info is required and contains the concrete local paths for each dataset named by grouped_datasets.

Returns the constructed PackedDataset instance. The caller is still responsible for invoking set_epoch immediately before DataLoader iteration (and for the pre-iter RNG reseed — see PackedDataset.set_epoch docstring).

nemo_automodel.components.datasets.multimodal.utils.patchify(
    image,
    patch_size
)

Patchify a CxHxW image tensor into (H/p * W/p, ppC).

nemo_automodel.components.datasets.multimodal.utils.pil_img2rgb(
    image
)

Convert a PIL image to RGB, compositing transparent pixels on white.

nemo_automodel.components.datasets.multimodal.utils.prepare_attention_mask_per_sample(
    split_lens,
    attn_modes,
    device = 'cpu'
)

Build a per-sample additive float mask honoring causal / full / noise splits.

nemo_automodel.components.datasets.multimodal.datasets.DATASET_REGISTRY = {'t2i_pretrain': T2IIterableDataset, 'vlm_sft': SftJSONLIterableDataset, 'unifie...

nemo_automodel.components.datasets.multimodal.datasets.DEFAULT_DATASET_INFO = {}

nemo_automodel.components.datasets.multimodal.packing.DataConfig = BagelDatasetConfig