nemo_automodel.components.training

Training utilities shared across recipes.

Submodules

Package Contents

Classes

Name	Description
`StepSchedulerConfig`	User-facing step scheduler configuration.

API

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.

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

Build the step scheduler.

Parameters:

dataloader

DataLoader | None

The training dataloader, or None when max_steps is explicit.

dp_group_size

int

The size of the data parallel group.

local_batch_size

int

The size of the local batch.

process_group

ProcessGroup | NoneDefaults to None

Process group whose ranks participate in distributed signal handling. All ranks participating in training must also agree on stream completion.

synchronized_stream

bool | NoneDefaults to None

Whether to coordinate complete accumulation windows and stop at final source exhaustion without restarting the loader. None uses the loader’s synchronized_stream capability, defaulting to False.

Returns: StepScheduler

Configured StepScheduler.