nemo_automodel.components.datasets.llm.megatron_dataset

Module Contents

Classes

Name	Description
`MegatronPretraining`	Build Megatron pretraining datasets and dataloaders.
`MegatronPretrainingConfig`	Construction-time configuration for :class:`MegatronPretraining` (tokenizer is a build arg).

Functions

Name	Description
`get_list_of_files`	Get the list of unique dataset prefixes (full paths without extension) from a glob pattern.
`is_number_tryexcept`	Returns True if string is a number.
`is_zipped_list`	Check if the paths are zipped.
`try_load_blend_from_json`	Load a data blend configuration from a JSON file.
`validate_dataset_asset_accessibility`	Validate the accessibility of the dataset assets.

Data

logger

API

class nemo_automodel.components.datasets.llm.megatron_dataset.MegatronPretraining(
    paths: pathlib.Path | typing.List | typing.Dict[str, typing.List],
    seq_length: int = 2048,
    tokenizer: typing.Optional[transformers.tokenization_utils_base.PreTrainedTokenizerBase] = None,
    micro_batch_size: int = 4,
    global_batch_size: int = 8,
    create_attention_mask: bool = False,
    seed: int = 1234,
    split: str = '900,50,50',
    index_mapping_dir: typing.Optional[str] = None,
    num_dataset_builder_threads: int = 1,
    num_train_samples: typing.Optional[int] = None,
    num_val_samples: typing.Optional[int] = None,
    num_test_samples: typing.Optional[int] = None,
    trainer_max_steps: typing.Optional[int] = None,
    trainer_val_check_interval: int = 1000,
    trainer_limit_val_batches: typing.Union[int, float] = 1,
    trainer_limit_test_batches: typing.Union[int, float] = 1,
    mmap_bin_files: bool = True,
    splits_to_build: typing.Optional[typing.Union[str, typing.List[str]]] = None,
    object_storage_config: typing.Optional[typing.Union[typing.Dict, 'ObjectStorageConfig']] = None
)

Build Megatron pretraining datasets and dataloaders.

gpt_dataset_config

'GPTDatasetConfig'

Get the GPT dataset configuration.

nemo_automodel.components.datasets.llm.megatron_dataset.MegatronPretraining.build()

Build the datasets using the trainer parameters provided during initialization.

nemo_automodel.components.datasets.llm.megatron_dataset.MegatronPretraining.get_dataset(
    split: str
)

Get the dataset for a given split.

class nemo_automodel.components.datasets.llm.megatron_dataset.MegatronPretrainingConfig(
    paths: pathlib.Path | list[str] | dict[str, list[str]],
    seq_length: int = 2048,
    create_attention_mask: bool = False,
    seed: int = 1234,
    split: str = '900,50,50',
    index_mapping_dir: str | None = None,
    num_dataset_builder_threads: int = 1,
    num_train_samples: int | None = None,
    num_val_samples: int | None = None,
    num_test_samples: int | None = None,
    trainer_limit_val_batches: int | float = 1,
    trainer_limit_test_batches: int | float = 1,
    mmap_bin_files: bool = True,
    splits_to_build: str | list[str] | None = None,
    object_storage_config: dict[str, object] | nemo_automodel.components.datasets.llm.megatron.indexed_dataset.ObjectStorageConfig | None = None
)

Dataclass

Construction-time configuration for :class:MegatronPretraining (tokenizer is a build arg).

accepts_tokenizer

bool = True

builds_on_all_ranks

bool = True

create_attention_mask

bool = False

Whether to generate attention masks (not supported with fused/flash attention).

index_mapping_dir

str | None = None

Directory to write index mapping files.

mmap_bin_files

bool = True

Whether to memory-map .bin files.

num_dataset_builder_threads

int = 1

Number of threads to use for dataset building.

num_test_samples

int | None = None

Number of test samples.

num_train_samples

int | None = None

Number of training samples (defaults to total train steps x global batch size).

num_val_samples

int | None = None

Number of validation samples.

object_storage_config

dict[str, object] | ObjectStorageConfig | None = None

Configuration for reading .bin/.idx files from S3/MSC.

paths

Path | list[str] | dict[str, list[str]]

Paths of the data distributions (single path, list, dict, or path to a JSON blend file).

requires_training_schedule

bool = True

seed

int = 1234

Seed for generating the GPT dataset.

seq_length

int = 2048

Sequence length.

split

str = '900,50,50'

Comma-separated train/validation/test ratios (unused if paths is a dict).

splits_to_build

str | list[str] | None = None

Splits to build (None builds all splits).

trainer_limit_test_batches

int | float = 1

Limit for test batches.

trainer_limit_val_batches

int | float = 1

Limit for validation batches.

nemo_automodel.components.datasets.llm.megatron_dataset.MegatronPretrainingConfig.build(
    tokenizer: transformers.tokenization_utils_base.PreTrainedTokenizerBase | None,
    training_schedule: nemo_automodel.components.datasets.loader.DatasetBuildSchedule
) -> object

Build the Megatron pretraining torch Dataset (for DataloaderConfig.dataset_config).

Constructs the :class:MegatronPretraining builder, runs its build(), and returns the requested split via get_dataset. Batch sizes and scheduler limits are runtime values owned by the recipe’s step scheduler, so they arrive through training_schedule rather than being duplicated in YAML.

Parameters:

tokenizer

PreTrainedTokenizerBase | None

Runtime tokenizer used by the Megatron GPT dataset.

training_schedule

DatasetBuildSchedule

Local/global batch sizes and train/validation cadence from the recipe.

Returns: object

Requested Megatron dataset split.

nemo_automodel.components.datasets.llm.megatron_dataset.get_list_of_files(
    path: str
)

Get the list of unique dataset prefixes (full paths without extension) from a glob pattern.

nemo_automodel.components.datasets.llm.megatron_dataset.is_number_tryexcept(
    s
)

Returns True if string is a number.

nemo_automodel.components.datasets.llm.megatron_dataset.is_zipped_list(
    paths
)

Check if the paths are zipped.

nemo_automodel.components.datasets.llm.megatron_dataset.try_load_blend_from_json(
    path: typing.Union[str, pathlib.Path]
) -> typing.Optional[typing.Union[typing.Dict[str, typing.List], typing.List]]

Load a data blend configuration from a JSON file.

Two top-level JSON shapes are accepted:

Dict-of-splits (Automodel native form): keys are split names (‘train’, ‘valid’, ‘test’); values are path lists. Common aliases ‘valid’ / ‘val’ / ‘dev’ are normalized to ‘validation’.
Flat list (Megatron-LM canonical form): a single zipped list of alternating weights and dataset prefixes. The caller uses the split= parameter to allocate this blend across train / validation / test splits.

Example flat-list JSON (Megatron-LM convention, paired with split=): [“30”, “path/to/dataset1”, “70”, “path/to/dataset2”]

Parameters:

path

Union[str, Path]

Path to a JSON file containing the blend configuration.

Returns: Optional[Union[Dict[str, List], List]]

Dictionary or list containing the blend configuration if path is

Raises:

FileNotFoundError: If the JSON file does not exist.
PermissionError: If the JSON file cannot be read.
ValueError: If the JSON is invalid or is neither a list nor a dict.

nemo_automodel.components.datasets.llm.megatron_dataset.validate_dataset_asset_accessibility(
    paths,
    object_storage_config = None
)

Validate the accessibility of the dataset assets. Skips local-filesystem checks for S3/MSC paths when object_storage_config is provided.

nemo_automodel.components.datasets.llm.megatron_dataset.logger = logging.getLogger(__name__)