nemo_automodel.components.datasets.loader

Typed, data-parallel-aware dataloader construction.

:class:ParallelAwareDataloader is a stateful, data-parallel-aware DataLoader: it shards an IterableDataset (or attaches a distributed / length-grouped sampler to a map-style dataset) for (dp_rank, dp_world_size) and inherits StatefulDataLoader for per-rank checkpoint resume.
:class:DataloaderConfig owns dataset, sampler, packing, and dataloader construction. Runtime-only objects such as tokenizers and the rank-ordering context are explicit :meth:DataloaderConfig.build arguments.

Module Contents

Classes

Name	Description
`AllRanksDatasetConfig`	Dataset config whose build implementation performs collectives across all ranks.
`BatchSamplerConfig`	Typed construction contract for a dataset-specific batch sampler.
`CollatorConfig`	Construction-time configuration for a tokenizer-aware collator class.
`DataloaderConfig`	A typed dataset config + loader settings; :meth:`build` produces a :class:`ParallelAwareDataloader`.
`DatasetBuildSchedule`	Training schedule values required while materializing schedule-aware datasets.
`NeatPackingConfig`	NEAT (bin-packed) packing paired with `neat_packed_collater`.
`PackingConfig`	Base config for sequence packing; `None` (no config) means no packing.
`ParallelAwareDataloader`	Stateful, data-parallel-aware `DataLoader`.
`PlainDatasetConfig`	Dataset config whose build contract has no runtime dependencies.
`ScheduledDatasetConfig`	Dataset config that additionally consumes the recipe training schedule.
`ThdPackingConfig`	THD (flattened, `seq_lens`-based) packing; pairs with `packed_sequence_thd_collater`.
`TokenizerDatasetConfig`	Dataset config whose build contract accepts a runtime tokenizer or processor.
`_LegacyDatasetConfig`	Fallback shim for a dataset `_target_` that has no typed `<Name>Config`.

Functions

Name	Description
`_make_sampler`	Build the default map-style sampler (distributed, or length-grouped).
`_resolve_target`	Resolve `target` via `registry`, a dotted import path, or pass it through unchanged.
`_set_spawn_start_method`	Set the multiprocessing start method to `spawn` if not already set.
`_shard_iterable_dataset`	Shard an `IterableDataset` across data-parallel ranks for unique samples.
`make_collate_fn`	Resolve a collate target into a batch callable or tokenizer-aware config.
`make_dataset_config`	Resolve a dataset `_target_` to an object exposing a typed `build` method.
`make_packing_config`	Resolve a packing-config `target` and construct it from `kwargs` (`target=None` → no packing).

Data

_LEGACY_PACKING_FIELDS

_PACKING_CONFIGS

__all__

logger

API

class nemo_automodel.components.datasets.loader.AllRanksDatasetConfig()

Protocol

Dataset config whose build implementation performs collectives across all ranks.

builds_on_all_ranks

bool

class nemo_automodel.components.datasets.loader.BatchSamplerConfig()

Protocol

Typed construction contract for a dataset-specific batch sampler.

nemo_automodel.components.datasets.loader.BatchSamplerConfig.build(
    dataset_len: int,
    rank: int,
    world_size: int
) -> torch.utils.data.sampler.Sampler[list[int]]

Build a per-rank batch sampler for a materialized dataset.

class nemo_automodel.components.datasets.loader.CollatorConfig(
    factory: collections.abc.Callable[..., nemo_automodel.components.datasets.loader.CollateFn],
    kwargs: dict[str, object] = dict()
)

Dataclass

Construction-time configuration for a tokenizer-aware collator class.

factory

Callable[..., CollateFn]

kwargs

dict[str, object] = field(default_factory=dict)

nemo_automodel.components.datasets.loader.CollatorConfig.build(
    tokenizer: 'PreTrainedTokenizerBase | ProcessorMixin'
) -> nemo_automodel.components.datasets.loader.CollateFn

Instantiate the collator once with its runtime tokenizer or processor.

Parameters:

tokenizer

'PreTrainedTokenizerBase | ProcessorMixin'

Runtime tokenizer or multimodal processor used by the collator.

Returns: CollateFn

Callable collator passed to the dataloader.

class nemo_automodel.components.datasets.loader.DataloaderConfig(
    dataset_config: nemo_automodel.components.datasets.loader.DatasetConfig,
    packing: nemo_automodel.components.datasets.loader.PackingConfig | None = None,
    batch_sampler_config: nemo_automodel.components.datasets.loader.BatchSamplerConfig | None = None,
    dataset_build_schedule: nemo_automodel.components.datasets.loader.DatasetBuildSchedule | None = None,
    shuffle: bool | None = None,
    group_by_length: bool = False,
    shuffle_buffer_size: int = 10000,
    batch_size: int | None = 1,
    seed: int = 42,
    collate_fn: nemo_automodel.components.datasets.loader.CollateFn | nemo_automodel.components.datasets.loader.CollatorConfig | None = None,
    num_workers: int = 0,
    pin_memory: bool = False,
    persistent_workers: bool = False,
    prefetch_factor: int | None = None,
    drop_last: bool = False
)

Dataclass

A typed dataset config + loader settings; :meth:build produces a :class:ParallelAwareDataloader.

dataset_config is a typed per-dataset config (for example ChatDatasetConfig or GLUE_MRPCConfig); the dataset’s declarative arguments live there. Runtime dependencies are explicit :meth:build arguments, while supported DataLoader settings are named fields on this config.

batch_sampler_config

BatchSamplerConfig | None = None

batch_size

int | None = 1

collate_fn

CollateFn | CollatorConfig | None = None

dataset_build_schedule

DatasetBuildSchedule | None = None

dataset_builds_on_all_ranks

bool

Whether dataset construction must bypass rank-zero-first ordering.

dataset_config

DatasetConfig

drop_last

bool = False

group_by_length

bool = False

num_workers

int = 0

packing

PackingConfig | None = None

persistent_workers

bool = False

pin_memory

bool = False

prefetch_factor

int | None = None

seed

int = 42

shuffle

bool | None = None

shuffle_buffer_size

int = 10000

nemo_automodel.components.datasets.loader.DataloaderConfig._build_dataset(
    tokenizer: 'PreTrainedTokenizerBase | ProcessorMixin | None',
    dataset_build_context: contextlib.AbstractContextManager[object] | None
) -> object

Materialize only the dataset inside the caller-provided ordering context.

nemo_automodel.components.datasets.loader.DataloaderConfig.build(
    dp_rank: int,
    dp_world_size: int,
    pp_enabled: bool = False,
    tokenizer: 'PreTrainedTokenizerBase | ProcessorMixin | None' = None,
    dataset_build_context: contextlib.AbstractContextManager[object] | None = None,
    supports_seq_lens: bool = True,
    cp_size: int = 1,
    attn_implementation: str | None = None,
    collate_wrapper: collections.abc.Callable[[CollateFn], nemo_automodel.components.datasets.loader.CollateFn] | None = None
) -> torch.utils.data.DataLoader

Build the configured dataset, packing, sampler, collator, and stateful dataloader.

Parameters:

dp_rank

int

Rank within the data-parallel group.

dp_world_size

int

Size of the data-parallel group.

pp_enabled

boolDefaults to False

Whether pipeline parallelism requires dropping incomplete batches.

tokenizer

'PreTrainedTokenizerBase | ProcessorMixin | None'Defaults to None

Runtime tokenizer or multimodal processor for tokenizer-aware datasets and collators.

dataset_build_context

AbstractContextManager[object] | NoneDefaults to None

Optional caller-owned rank-ordering context used only while materializing the dataset. Collective dataset builders should pass None.

supports_seq_lens

boolDefaults to True

Whether the model forward contract accepts THD seq_lens metadata.

cp_size

intDefaults to 1

Context-parallel world size used for packed-sequence divisibility.

attn_implementation

str | NoneDefaults to None

Attention backend used by NEAT packing.

collate_wrapper

Callable[[CollateFn], CollateFn] | NoneDefaults to None

Optional recipe-owned wrapper around the resolved collator.

Returns: DataLoader

Stateful, data-parallel-aware dataloader.

class nemo_automodel.components.datasets.loader.DatasetBuildSchedule(
    local_batch_size: int,
    global_batch_size: int,
    max_steps: int | None,
    val_check_interval: int | None
)

Dataclass

Training schedule values required while materializing schedule-aware datasets.

global_batch_size

int

local_batch_size

int

max_steps

int | None

val_check_interval

int | None

class nemo_automodel.components.datasets.loader.NeatPackingConfig(
    packed_sequence_size: int,
    max_packs: int | None = None,
    prepacked: bool = False,
    num_proc: int = 1,
    drop_long_samples: bool = True
)

Dataclass

Bases: PackingConfig

NEAT (bin-packed) packing paired with neat_packed_collater.

drop_long_samples

bool = True

nemo_automodel.components.datasets.loader.NeatPackingConfig.build(
    dataset: object,
    split: str | list[str] | None = None,
    seed: int = 42,
    supports_seq_lens: bool = True,
    pad_token_id: int = 0,
    cp_size: int = 1,
    attn_implementation: str | None = None
) -> tuple[object, nemo_automodel.components.datasets.loader.CollateFn]

Pack with NEAT and configure the collator for the selected attention implementation.

class nemo_automodel.components.datasets.loader.PackingConfig(
    packed_sequence_size: int,
    max_packs: int | None = None,
    prepacked: bool = False,
    num_proc: int = 1
)

Dataclass

Base config for sequence packing; None (no config) means no packing.

Subclasses (:class:ThdPackingConfig / :class:NeatPackingConfig) pick the packing strategy and the matching collater. :meth:build returns (dataset, collate_fn) — the packed dataset and the packing-specific collater. Construction-time knobs are the fields; runtime / model-derived values (split / seed / supports_seq_lens / pad_token_id / cp_size / attn_implementation) are :meth:build args.

max_packs

int | None = None

num_proc

int = 1

Number of processes used to pre-tokenize an indexable dataset before packing.

packed_sequence_size

int

prepacked

bool = False

Whether the dataset already contains packed samples and must not be repacked.

nemo_automodel.components.datasets.loader.PackingConfig._pretokenize(
    dataset: object
) -> object

Materialize lazy tokenization in parallel when packing can consume the full dataset.

nemo_automodel.components.datasets.loader.PackingConfig.build(
    dataset: object,
    split: str | list[str] | None = None,
    seed: int = 42,
    supports_seq_lens: bool = True,
    pad_token_id: int = 0,
    cp_size: int = 1,
    attn_implementation: str | None = None
) -> tuple[object, nemo_automodel.components.datasets.loader.CollateFn | None]

Pack dataset and return (dataset, collate_fn).

class nemo_automodel.components.datasets.loader.ParallelAwareDataloader(
    dataset: object,
    dp_rank: int,
    dp_world_size: int,
    batch_size: int | None = 1,
    collate_fn: nemo_automodel.components.datasets.loader.CollateFn | None = None,
    batch_sampler: torch.utils.data.sampler.Sampler[list[int]] | None = None,
    seed: int = 42,
    shuffle: bool | None = None,
    group_by_length: bool = False,
    pp_enabled: bool = False,
    shuffle_buffer_size: int = 10000,
    num_workers: int = 0,
    pin_memory: bool = False,
    persistent_workers: bool = False,
    prefetch_factor: int | None = None,
    drop_last: bool = False
)

Bases: StatefulDataLoader

Stateful, data-parallel-aware DataLoader.

Routes a dataset to (dp_rank, dp_world_size): an IterableDataset is sharded (and optionally buffer-shuffled); a map-style dataset gets the default distributed (or length-grouped) sampler with batch_size and drop_last under PP. Inherits StatefulDataLoader for per-rank checkpoint resume.

class nemo_automodel.components.datasets.loader.PlainDatasetConfig()

Protocol

Dataset config whose build contract has no runtime dependencies.

nemo_automodel.components.datasets.loader.PlainDatasetConfig.build() -> object

Build the configured dataset.

class nemo_automodel.components.datasets.loader.ScheduledDatasetConfig()

Protocol

Bases: TokenizerDatasetConfig

Dataset config that additionally consumes the recipe training schedule.

requires_training_schedule

bool

nemo_automodel.components.datasets.loader.ScheduledDatasetConfig.build(
    tokenizer: 'PreTrainedTokenizerBase | ProcessorMixin | None',
    training_schedule: nemo_automodel.components.datasets.loader.DatasetBuildSchedule
) -> object

Build the configured dataset with tokenizer and schedule dependencies.

class nemo_automodel.components.datasets.loader.ThdPackingConfig(
    packed_sequence_size: int,
    max_packs: int | None = None,
    prepacked: bool = False,
    num_proc: int = 1
)

Dataclass

Bases: PackingConfig

THD (flattened, seq_lens-based) packing; pairs with packed_sequence_thd_collater.

Requires a model whose forward accepts seq_lens — packing is skipped (with a warning) otherwise.

nemo_automodel.components.datasets.loader.ThdPackingConfig.build(
    dataset: object,
    split: str | list[str] | None = None,
    seed: int = 42,
    supports_seq_lens: bool = True,
    pad_token_id: int = 0,
    cp_size: int = 1,
    attn_implementation: str | None = None
) -> tuple[object, nemo_automodel.components.datasets.loader.CollateFn | None]

Pack with THD; returns (dataset, None) if the model does not accept seq_lens.

class nemo_automodel.components.datasets.loader.TokenizerDatasetConfig()

Protocol

Dataset config whose build contract accepts a runtime tokenizer or processor.

accepts_tokenizer

bool

nemo_automodel.components.datasets.loader.TokenizerDatasetConfig.build(
    tokenizer: 'PreTrainedTokenizerBase | ProcessorMixin | None'
) -> object

Build the configured dataset with a runtime tokenizer or processor.

class nemo_automodel.components.datasets.loader._LegacyDatasetConfig(
    factory: collections.abc.Callable[..., object],
    kwargs: dict[str, object],
    accepts_tokenizer: bool = False
)

Dataclass

Fallback shim for a dataset _target_ that has no typed <Name>Config.

:func:make_dataset_config maps known datasets onto their typed config; the few targets with no config yet (a handful of VLM make_* factories) land here instead. Wraps the target so it satisfies the typed dataset build contract. Config resolution records whether the external target accepts a tokenizer, and :meth:build forwards only that explicit runtime dependency.

accepts_tokenizer

bool = False

factory

Callable[..., object]

kwargs

dict[str, object]

nemo_automodel.components.datasets.loader._LegacyDatasetConfig.build(
    tokenizer: 'PreTrainedTokenizerBase | ProcessorMixin | None' = None
) -> object

Call the wrapped dataset target with its declared YAML arguments.

Parameters:

tokenizer

'PreTrainedTokenizerBase | ProcessorMixin | None'Defaults to None

Runtime tokenizer or processor. It is forwarded only when the target explicitly declares a tokenizer parameter; that compatibility decision is recorded once at config resolution.

Returns: object

Dataset object returned by the configured target.

nemo_automodel.components.datasets.loader._make_sampler(
    dataset: typing.Any,
    dp_rank: int,
    dp_world_size: int,
    seed: int,
    shuffle: bool,
    group_by_length: bool,
    batch_size: int
) -> torch.utils.data.sampler.Sampler

Build the default map-style sampler (distributed, or length-grouped).

nemo_automodel.components.datasets.loader._resolve_target(
    target: typing.Any,
    registry: dict[str, typing.Any]
) -> typing.Any

Resolve target via registry, a dotted import path, or pass it through unchanged.

target may be a key in registry (whose value is the object itself or a dotted path to it), a dotted import path ("pkg.mod.attr"), or an already-resolved object (returned as-is). This backs the make_* resolvers below so packing configs, collate fns and sampler factories share one lookup.

nemo_automodel.components.datasets.loader._set_spawn_start_method() -> None

Set the multiprocessing start method to spawn if not already set.

nemo_automodel.components.datasets.loader._shard_iterable_dataset(
    dataset: typing.Any,
    dp_rank: int,
    dp_world_size: int
) -> typing.Any

Shard an IterableDataset across data-parallel ranks for unique samples.

Calls the dataset’s own shard or HuggingFace split_dataset_by_node.

nemo_automodel.components.datasets.loader.make_collate_fn(
    target: object,
    kwargs: dict[str, object] | None = None
) -> nemo_automodel.components.datasets.loader.CollateFn | nemo_automodel.components.datasets.loader.CollatorConfig | None

Resolve a collate target into a batch callable or tokenizer-aware config.

target is a built-in collator key (for example "default"), a dotted import path, or an already resolved callable. Collator classes become :class:CollatorConfig so :meth:DataloaderConfig.build instantiates them once with the runtime tokenizer. Function kwargs remain partial-bound per batch.

nemo_automodel.components.datasets.loader.make_dataset_config(
    target: object,
    kwargs: dict[str, object] | None = None
) -> nemo_automodel.components.datasets.loader.DatasetConfig

Resolve a dataset _target_ to an object exposing a typed build method.

Exact legacy registrations map a pre-config dataset class or make_* factory onto its typed <Name>Config without class-name dispatch. Misses fall back to the _target_ itself, so a target that already names a config imports directly. Typed config fields are validated before construction; external factories use :class:_LegacyDatasetConfig.

nemo_automodel.components.datasets.loader.make_packing_config(
    target: str | None,
    kwargs: dict[str, object] | None = None
) -> nemo_automodel.components.datasets.loader.PackingConfig | None

Resolve a packing-config target and construct it from kwargs (target=None → no packing).

target is either a built-in strategy key ("thd" / "neat") or a dotted import path to a :class:PackingConfig subclass (e.g. "my_pkg.MyPackingConfig"). Strategy-specific fields from the union-shaped packed_sequence YAML block are accepted and filtered for the selected strategy; unknown fields are rejected instead of silently disappearing.

nemo_automodel.components.datasets.loader.CollateFn = Callable[[list[object]], object]

nemo_automodel.components.datasets.loader.DatasetConfig = PlainDatasetConfig | TokenizerDatasetConfig | ScheduledDatasetConfig

nemo_automodel.components.datasets.loader._COLLATE_FNS: dict[str, str] = {'default': 'nemo_automodel.components.datasets.utils.default_collater'}

nemo_automodel.components.datasets.loader._DATASETS = 'nemo_automodel.components.datasets'

nemo_automodel.components.datasets.loader._DATASET_CONFIGS: dict[str, str] = {f'{_DATASETS}.llm.megatron_dataset.MegatronPretraining': f'{_DATASETS}.llm.mega...

nemo_automodel.components.datasets.loader._LEGACY_PACKING_FIELDS = {'split_across_pack'}

nemo_automodel.components.datasets.loader._PACKING_CONFIGS: dict[str, type[PackingConfig]] = {'thd': ThdPackingConfig, 'neat': NeatPackingConfig}

nemo_automodel.components.datasets.loader.__all__ = ['CollatorConfig', 'DatasetBuildSchedule', 'DataloaderConfig', 'NeatPackingConfi...

nemo_automodel.components.datasets.loader.logger = logging.getLogger(__name__)