> This page is for version 0.5.0 · 26.06.
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt

> 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.training

Training utilities shared across recipes.

## Submodules

* **[`nemo_automodel.components.training._mamba_ssd_prewarm`](/nemo-automodel/nemo_automodel/components/training/_mamba_ssd_prewarm)**
* **[`nemo_automodel.components.training._prewarm_utils`](/nemo-automodel/nemo_automodel/components/training/_prewarm_utils)**
* **[`nemo_automodel.components.training.domain_mixture`](/nemo-automodel/nemo_automodel/components/training/domain_mixture)**
* **[`nemo_automodel.components.training.ema`](/nemo-automodel/nemo_automodel/components/training/ema)**
* **[`nemo_automodel.components.training.embedding_row_repair`](/nemo-automodel/nemo_automodel/components/training/embedding_row_repair)**
* **[`nemo_automodel.components.training.garbage_collection`](/nemo-automodel/nemo_automodel/components/training/garbage_collection)**
* **[`nemo_automodel.components.training.model_output_utils`](/nemo-automodel/nemo_automodel/components/training/model_output_utils)**
* **[`nemo_automodel.components.training.neftune`](/nemo-automodel/nemo_automodel/components/training/neftune)**
* **[`nemo_automodel.components.training.prewarm`](/nemo-automodel/nemo_automodel/components/training/prewarm)**
* **[`nemo_automodel.components.training.rng`](/nemo-automodel/nemo_automodel/components/training/rng)**
* **[`nemo_automodel.components.training.signal_handler`](/nemo-automodel/nemo_automodel/components/training/signal_handler)**
* **[`nemo_automodel.components.training.step_scheduler`](/nemo-automodel/nemo_automodel/components/training/step_scheduler)**
* **[`nemo_automodel.components.training.timers`](/nemo-automodel/nemo_automodel/components/training/timers)**
* **[`nemo_automodel.components.training.utils`](/nemo-automodel/nemo_automodel/components/training/utils)**

## Package Contents

### Classes

| Name                                                                                            | Description                               |
| ----------------------------------------------------------------------------------------------- | ----------------------------------------- |
| [`StepSchedulerConfig`](#nemo_automodel-components-training-step_scheduler-StepSchedulerConfig) | User-facing step scheduler configuration. |

### API

```python
class nemo_automodel.components.training.step_scheduler.StepSchedulerConfig(
    global_batch_size: int = 32,
    num_epochs: int | None = 10,
    max_steps: int | None = None,
    ckpt_every_steps: int | None = 100,
    save_checkpoint_every_epoch: bool = True,
    validate_on_checkpoint: bool = True,
    val_every_steps: int | None = None,
    log_remote_every_steps: int = 1,
    loss_average_window_steps: int = 50,
    gc_every_steps: int | None = None,
    start_step: int = 0,
    start_epoch: int = 0,
    preemption_signal: int | str | list[int | str] | None = 'SIGTERM'
)
```

Dataclass

User-facing step scheduler configuration.

These fields correspond to the YAML-configurable parameters of the
training loop.  Runtime-only values (`dataloader`, `dp_size`,
`local_batch_size`) are passed separately to `build_step_scheduler`.

**`ckpt_every_steps`** `int | None = 100`

Save a checkpoint every N optimizer steps.
`None` defaults to once per epoch.

---

**`gc_every_steps`** `int | None = None`

Force `gc.collect()` every N steps.
`None` disables manual GC.

---

**`global_batch_size`** `int = 32`

Total samples per optimizer step across all GPUs.

---

**`log_remote_every_steps`** `int = 1`

Log to WandB / MLflow every N steps.

---

**`loss_average_window_steps`** `int = 50`

Rolling window size for averaged training loss
metrics.

---

**`max_steps`** `int | None = None`

Hard cap on optimizer steps.  `None` means derive from
`num_epochs * epoch_len`.

---

**`num_epochs`** `int | None = 10`

Number of training epochs.  When `None` the builder
derives it from `max_steps`.  Default: 10.

---

**`preemption_signal`** `int | str | list[int | str] | None = 'SIGTERM'`

Signal(s) that trigger a graceful preemption checkpoint, each given as
a signal number, name (e.g. `"SIGTERM"`), or a list thereof.  `None` disables
preemption checkpointing.  Default: `"SIGTERM"`.

---

**`save_checkpoint_every_epoch`** `bool = True`

Also checkpoint at every epoch boundary.

---

**`start_epoch`** `int = 0`

Initial epoch (for checkpoint resume).

---

**`start_step`** `int = 0`

Initial global step (for checkpoint resume).

---

**`val_every_steps`** `int | None = None`

Run validation every N optimizer steps.
`None` disables periodic validation.

---

**`validate_on_checkpoint`** `bool = True`

Also run validation on checkpoint steps.

---

```python
nemo_automodel.components.training.step_scheduler.StepSchedulerConfig.build(
    dataloader: torch.utils.data.DataLoader,
    dp_group_size: int,
    local_batch_size: int,
    process_group: torch.distributed.ProcessGroup | None = None
) -> nemo_automodel.components.training.step_scheduler.StepScheduler
```

Build the step scheduler.

**Parameters:**

**`dataloader`** `DataLoader`

The training dataloader.

---

**`dp_group_size`** `int`

The size of the data parallel group.

---

**`local_batch_size`** `int`

The size of the local batch.

---

**`process_group`** `ProcessGroup | None` — default: None

Process group whose ranks participate in distributed signal handling.

---

**Returns:** `StepScheduler`

Configured StepScheduler.