> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

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

### Functions

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

### Data

[`CollateFn`](#nemo_automodel-components-datasets-loader-CollateFn)

[`DatasetConfig`](#nemo_automodel-components-datasets-loader-DatasetConfig)

[`_COLLATE_FNS`](#nemo_automodel-components-datasets-loader-_COLLATE_FNS)

[`_DATASETS`](#nemo_automodel-components-datasets-loader-_DATASETS)

[`_DATASET_CONFIGS`](#nemo_automodel-components-datasets-loader-_DATASET_CONFIGS)

[`_LEGACY_PACKING_FIELDS`](#nemo_automodel-components-datasets-loader-_LEGACY_PACKING_FIELDS)

[`_PACKING_CONFIGS`](#nemo_automodel-components-datasets-loader-_PACKING_CONFIGS)

[`__all__`](#nemo_automodel-components-datasets-loader-__all__)

[`logger`](#nemo_automodel-components-datasets-loader-logger)

### API

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

Protocol

Dataset config whose build implementation performs collectives across all ranks.

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

Protocol

Typed construction contract for a dataset-specific batch sampler.

```python
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.

```python
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.

```python
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:**

Runtime tokenizer or multimodal processor used by the collator.

**Returns:** `CollateFn`

Callable collator passed to the dataloader.

```python
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.

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

```python
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.

```python
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:**

Rank within the data-parallel group.

Size of the data-parallel group.

Whether pipeline parallelism requires dropping incomplete batches.

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

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

Whether the model forward contract accepts THD `seq_lens` metadata.

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

Attention backend used by NEAT packing.

Optional recipe-owned wrapper around the resolved collator.

**Returns:** `DataLoader`

Stateful, data-parallel-aware dataloader.

```python
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.

```python
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](#nemo_automodel-components-datasets-loader-PackingConfig)

NEAT (bin-packed) packing paired with `neat_packed_collater`.

```python
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.

```python
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.

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

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

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

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

```python
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)`.

```python
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.

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

Protocol

Dataset config whose build contract has no runtime dependencies.

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

Build the configured dataset.

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

Protocol

**Bases:** [TokenizerDatasetConfig](#nemo_automodel-components-datasets-loader-TokenizerDatasetConfig)

Dataset config that additionally consumes the recipe training schedule.

```python
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.

```python
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](#nemo_automodel-components-datasets-loader-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.

```python
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`.

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

Protocol

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

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

Build the configured dataset with a runtime tokenizer or processor.

```python
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.

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

Call the wrapped dataset target with its declared YAML arguments.

**Parameters:**

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.

```python
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).

```python
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.

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

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

```python
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`.

```python
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.

```python
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`.

```python
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.

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

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

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

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

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

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

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

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

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