nemo_automodel.components.datasets.multimodal.packing

Packed-sequence iterable for BAGEL training.

Module Contents

Classes

Name	Description
`PackedDataset`	Greedy pack of samples drawn from weighted groups into token-budgeted batches.

Data

DataConfig

logger

API

class nemo_automodel.components.datasets.multimodal.packing.PackedDataset(
    data_config: 'BagelDatasetConfig',
    tokenizer: 'PreTrainedTokenizerBase',
    special_tokens: collections.abc.Mapping[str, int],
    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: collections.abc.Mapping[str, list[int | list[int] | None]] | 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
)

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

nemo_automodel.components.datasets.multimodal.packing.logger = logging.getLogger(__name__)