> For clean Markdown of any page, append .md to the page URL. > For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/training/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. ## Docs - [nemo_automodel.components.training.domain_mixture](https://docs.nvidia.com/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/training/domain_mixture.md) - [nemo_automodel.components.training.ema](https://docs.nvidia.com/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/training/ema.md) - [nemo_automodel.components.training.embedding_row_repair](https://docs.nvidia.com/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/training/embedding_row_repair.md) - [nemo_automodel.components.training.garbage_collection](https://docs.nvidia.com/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/training/garbage_collection.md) - [nemo_automodel.components.training.model_output_utils](https://docs.nvidia.com/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/training/model_output_utils.md) - [nemo_automodel.components.training.neftune](https://docs.nvidia.com/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/training/neftune.md) - [nemo_automodel.components.training.prewarm](https://docs.nvidia.com/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/training/prewarm.md) - [nemo_automodel.components.training.rng](https://docs.nvidia.com/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/training/rng.md) - [nemo_automodel.components.training.signal_handler](https://docs.nvidia.com/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/training/signal_handler.md) - [nemo_automodel.components.training.step_scheduler](https://docs.nvidia.com/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/training/step_scheduler.md) - [nemo_automodel.components.training.timers](https://docs.nvidia.com/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/training/timers.md) - [nemo_automodel.components.training.utils](https://docs.nvidia.com/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/training/utils.md)