nemo_automodel.components.training.step_scheduler

Module Contents

Classes

Name	Description
`StepScheduler`	Scheduler for managing gradient accumulation and checkpointing steps.
`StepSchedulerConfig`	User-facing step scheduler configuration.

Functions

Name	Description
`_calculate_max_steps`	Calculate the maximum number of steps.
`_calculate_num_epochs`	Calculate the number of epochs out of maximum number of steps.

Data

logger

API

class nemo_automodel.components.training.step_scheduler.StepScheduler(
    global_batch_size: int,
    local_batch_size: int,
    dp_size: int,
    dataloader: torch.utils.data.DataLoader | None,
    ckpt_every_steps: int | None = None,
    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,
    num_epochs: int | None = None,
    max_steps: int | None = None,
    preemption_signal: nemo_automodel.components.training.signal_handler.SignalLike | list[nemo_automodel.components.training.signal_handler.SignalLike] | None = signal.SIGTERM,
    process_group: torch.distributed.ProcessGroup | None = None,
    synchronized_stream: bool = False
)

Bases: Stateful

Scheduler for managing gradient accumulation and checkpointing steps.

_sig_polled_step

int | None = None

epoch_len

= ceil(len(dataloader) / self.grad_acc_steps)

epochs

Epoch iterator.

grad_acc_steps

= global_batch_size // (local_batch_size * dp_size)

is_ckpt_step

Returns whether this step needs to call the checkpoint saving.

is_gc_step

Returns whether this step needs to run manual garbage collection.

is_last_batch

Returns whether this is the last batch for this epoch.

is_last_step

Returns whether the current step is the final training step.

Training stops at whichever comes first: reaching max_steps or exhausting the configured number of epochs (see __iter__ and epochs). max_steps alone is therefore not enough to detect the end — a small dataset can run out of epochs long before max_steps is hit (e.g. max_steps=100 with only 60 steps’ worth of data). In that case the last batch of the last epoch is the final step. Detect it so the final checkpoint and consolidated export — which key off this flag (see is_ckpt_step and the recipes’ is_final_checkpoint) — are still written.

is_remote_logging_step

Returns whether this step should log to remote services (WandB, MLflow, etc.).

is_val_step

Returns whether this step needs to call the validation.

sigterm_received

Returns whether SIGTERM was received.

stream_exhausted

bool = False

nemo_automodel.components.training.step_scheduler.StepScheduler.__iter__()

Iterates over dataloader while keeping track of counters.

Raises:

StopIteration: If the dataloader was exhausted or max_steps was reached.

nemo_automodel.components.training.step_scheduler.StepScheduler._iter_stream() -> collections.abc.Iterator[list[collections.abc.Mapping[str, object]]]

Yield complete accumulation windows after collective availability checks.

The only buffered data is the current optimizer step, as in the native iterator. There is no next-step lookahead that could advance checkpointed data progress.

Returns: Iterator[list[Mapping[str, object]]]

Iterator of microbatch mappings passed through without mutation. Tensor fields

nemo_automodel.components.training.step_scheduler.StepScheduler.load_state_dict(
    s
)

Load the scheduler state from a dictionary.

Parameters:

dict

Dictionary containing ‘step’ and ‘epoch’.

nemo_automodel.components.training.step_scheduler.StepScheduler.set_epoch(
    epoch: int
)

Set the epoch for the sampler and the collate function.

The collate function is notified as well as the sampler because a collator may derive per-epoch augmentation from the epoch. Only this scheduler’s dataloader is touched, so a separately built validation loader keeps a fixed epoch and therefore a fixed augmentation mix.

nemo_automodel.components.training.step_scheduler.StepScheduler.state_dict()

Get the current state of the scheduler.

Returns:

Current state with ‘step’ and ‘epoch’ keys.

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.

nemo_automodel.components.training.step_scheduler._calculate_max_steps(
    num_epochs: int,
    epoch_len: int | None,
    default_max_steps: int = 9223372036854775807
) -> int

Calculate the maximum number of steps.

nemo_automodel.components.training.step_scheduler._calculate_num_epochs(
    max_steps: int | None,
    epoch_len: int | None,
    default_num_epochs: int = 10
) -> int

Calculate the number of epochs out of maximum number of steps.

nemo_automodel.components.training.step_scheduler.logger = logging.getLogger(__name__)