> 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.llm.megatron.sampler

## Module Contents

### Classes

| Name                                                                                                                            | Description                                                       |
| ------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------- |
| [`BaseMegatronSampler`](#nemo_automodel-components-datasets-llm-megatron-sampler-BaseMegatronSampler)                           | Base class for Megatron batch samplers.                           |
| [`MegatronPretrainingRandomSampler`](#nemo_automodel-components-datasets-llm-megatron-sampler-MegatronPretrainingRandomSampler) | Randomized sampler with per-epoch shuffling and per-rank slicing. |
| [`MegatronPretrainingSampler`](#nemo_automodel-components-datasets-llm-megatron-sampler-MegatronPretrainingSampler)             | Deterministic sequential sampler with per-rank slicing.           |
| [`MegatronSamplerConfig`](#nemo_automodel-components-datasets-llm-megatron-sampler-MegatronSamplerConfig)                       | Declarative configuration for a Megatron micro-batch sampler.     |

### Functions

| Name                                                                                                          | Description                    |
| ------------------------------------------------------------------------------------------------------------- | ------------------------------ |
| [`create_megatron_sampler`](#nemo_automodel-components-datasets-llm-megatron-sampler-create_megatron_sampler) | Factory for Megatron samplers. |

### API

```python
class nemo_automodel.components.datasets.llm.megatron.sampler.BaseMegatronSampler(
    total_samples: int,
    micro_batch_size: int,
    data_parallel_rank: int,
    data_parallel_size: int,
    drop_last: bool = True,
    global_batch_size: int | None = None,
    pad_samples_to_global_batch_size: bool | None = False
)
```

Base class for Megatron batch samplers.

Provides common validation and shared behavior for Megatron samplers.
Implementations must yield lists of dataset indices that correspond to
one micro-batch for a single data-parallel rank.

**Parameters:**

**`total_samples`** `int`

Total available samples in the dataset.

---

**`micro_batch_size`** `int`

Number of samples per micro-batch on each data-parallel
rank.

---

**`data_parallel_rank`** `int`

Rank id in the data-parallel group that this sampler
will serve.

---

**`data_parallel_size`** `int`

World size of the data-parallel group.

---

**`drop_last`** `bool` — default: True

If True, drop incomplete batches. If False, implementations
may yield a final partial micro-batch (subject to their constraints).

---

**`global_batch_size`** `int | None` — default: None

Effective global batch size across all data-parallel
ranks; when provided, length is computed in global-batch units and
converted to micro-batches.

---

**`pad_samples_to_global_batch_size`** `bool | None` — default: False

If True and supported by the sampler,
the last incomplete global batch will be padded to `global_batch_size`
when `drop_last` is False.

---

**`micro_batch_times_data_parallel_size`** `= self.micro_batch_size * data_parallel_size`

---

```python
nemo_automodel.components.datasets.llm.megatron.sampler.BaseMegatronSampler.__iter__()
```

abstract

```python
nemo_automodel.components.datasets.llm.megatron.sampler.BaseMegatronSampler.__len__()
```

Return the number of micro-batches this sampler will yield.

If `global_batch_size` is provided, the length is computed in terms of
global batches and converted to micro-batches to align with training
loops that iterate by micro-batch.

```python
class nemo_automodel.components.datasets.llm.megatron.sampler.MegatronPretrainingRandomSampler(
    total_samples: int,
    micro_batch_size: int,
    data_parallel_rank: int,
    data_parallel_size: int,
    drop_last: bool = True,
    global_batch_size: int | None = None,
    pad_samples_to_global_batch_size: bool | None = False,
    seed: int = 0
)
```

**Bases:** [BaseMegatronSampler](#nemo_automodel-components-datasets-llm-megatron-sampler-BaseMegatronSampler)

Randomized sampler with per-epoch shuffling and per-rank slicing.

Uses a deterministic seed schedule `seed + epoch` to randomize indices
within each data-parallel shard (bucket). Notably, this sampler:

* Does not support padding the last global batch.
* Requires `drop_last=True` when the product `micro_batch_size *
  data_parallel_size > 1`.

**`consumed_samples`** `= 0`

---

**`last_batch_size`**

---

```python
nemo_automodel.components.datasets.llm.megatron.sampler.MegatronPretrainingRandomSampler.__iter__()
```

Yield randomized micro-batches for this rank.

Each epoch shuffles indices within the per-rank bucket using
`torch.randperm` seeded by `seed + epoch`. The sampler then emits
contiguous micro-batches of size `micro_batch_size` for this rank.

```python
nemo_automodel.components.datasets.llm.megatron.sampler.MegatronPretrainingRandomSampler.__len__()
```

Return the number of micro-batches that will be produced.

Accounts for `drop_last` by excluding a trailing incomplete global batch.
When `global_batch_size` is provided, converts global batches to
micro-batches.

```python
class nemo_automodel.components.datasets.llm.megatron.sampler.MegatronPretrainingSampler(
    total_samples: int,
    micro_batch_size: int,
    data_parallel_rank: int,
    data_parallel_size: int,
    drop_last: bool = True,
    global_batch_size: int | None = None,
    pad_samples_to_global_batch_size: bool | None = False
)
```

**Bases:** [BaseMegatronSampler](#nemo_automodel-components-datasets-llm-megatron-sampler-BaseMegatronSampler)

Deterministic sequential sampler with per-rank slicing.

Iterates deterministically over sample indices, splits each global batch
across data-parallel ranks, and yields per-rank micro-batches. When
`drop_last` is False and `pad_samples_to_global_batch_size` is True, the
final global batch is padded to a full size so that all ranks emit complete
micro-batches.

**Raises:**

* `RuntimeError`: If there are no samples left to consume.

```python
nemo_automodel.components.datasets.llm.megatron.sampler.MegatronPretrainingSampler.__iter__()
```

Yield lists of indices forming per-rank micro-batches.

Iterates up to `total_samples`. Optionally pads
the last global batch when `drop_last` is False and
`pad_samples_to_global_batch_size` is True.

```python
nemo_automodel.components.datasets.llm.megatron.sampler.MegatronPretrainingSampler.get_start_end_idx()
```

Return slice boundaries for this rank within a global batch.

**Returns:**

Tuple of `(start_idx, end_idx)` used to extract this rank's

```python
class nemo_automodel.components.datasets.llm.megatron.sampler.MegatronSamplerConfig(
    micro_batch_size: int,
    global_batch_size: int,
    dataloader_type: typing.Literal['single', 'cyclic'] = 'single',
    drop_last: bool = True,
    pad_samples_to_global_batch_size: bool = False
)
```

Dataclass

Declarative configuration for a Megatron micro-batch sampler.

**`dataloader_type`** `Literal['single', 'cyclic'] = 'single'`

---

**`drop_last`** `bool = True`

---

**`global_batch_size`** `int`

---

**`micro_batch_size`** `int`

---

**`pad_samples_to_global_batch_size`** `bool = False`

---

```python
nemo_automodel.components.datasets.llm.megatron.sampler.MegatronSamplerConfig.build(
    dataset_len: int,
    rank: int,
    world_size: int
) -> nemo_automodel.components.datasets.llm.megatron.sampler.BaseMegatronSampler
```

Build a Megatron sampler for one data-parallel rank.

**Parameters:**

**`dataset_len`** `int`

Number of samples in the materialized dataset.

---

**`rank`** `int`

Rank within the data-parallel group.

---

**`world_size`** `int`

Size of the data-parallel group.

---

**Returns:** `BaseMegatronSampler`

Sampler yielding one local micro-batch of dataset indices at a time.

```python
nemo_automodel.components.datasets.llm.megatron.sampler.create_megatron_sampler(
    dataset_len: int,
    micro_batch_size: int,
    global_batch_size: int,
    dataloader_type: typing.Literal['single', 'cyclic'] = 'single',
    drop_last: bool = True,
    pad_samples_to_global_batch_size: bool = False,
    rank: int = 0,
    world_size: int = 1
) -> nemo_automodel.components.datasets.llm.megatron.sampler.BaseMegatronSampler
```

Factory for Megatron samplers.

Constructs and returns a Megatron-compatible sampler for a dataset of a
given length and batch configuration. The returned sampler yields lists of
indices per micro-batch for a single data-parallel rank.

**Parameters:**

**`dataset_len`** `int`

Number of samples in the underlying dataset.

---

**`micro_batch_size`** `int`

Number of samples per micro-batch on each
data-parallel rank.

---

**`global_batch_size`** `int`

Effective global batch size across all
data-parallel ranks (`micro_batch_size * world_size * grad_accum`).

---

**`dataloader_type`** `Literal['single', 'cyclic']` — default: 'single'

Sampler type to construct. Supported values:

* "single": Deterministic sequential sampling
  (`MegatronPretrainingSampler`).
* "cyclic": Randomized per-epoch sampling
  (`MegatronPretrainingRandomSampler`).
  The value "batch" is not supported in this implementation.

---

**`drop_last`** `bool` — default: True

When True, drop a trailing incomplete batch.

---

**`pad_samples_to_global_batch_size`** `bool` — default: False

When True and supported by the sampler,
pad the final global batch to `global_batch_size` if `drop_last` is
False.

---

**`rank`** `int` — default: 0

Data-parallel rank id for this process.

---

**`world_size`** `int` — default: 1

Number of data-parallel ranks.

---

**Returns:** `BaseMegatronSampler`

Configured sampler instance for the requested type.

**Raises:**

* `Exception`: If an unsupported `dataloader_type` is provided.