ReferenceFull Library ReferenceNemo AutomodelNemo AutomodelComponentsTraining

nemo_automodel.components.training.rng

View as Markdown

Module Contents

Classes

NameDescription
RNGStateLegacy RNG state kept for trusted pickle-based checkpoint restores.
ScopedRNGContext manager for reproducible RNG states across random, NumPy, and PyTorch.
StatefulRNGRNG manager for reproducible RNG states across random, NumPy, and PyTorch.
_RNGStateWeights-only-safe snapshot of Python, NumPy, Torch, and CUDA RNG states.

Functions

NameDescription
_get_rng_stateGet current RNG states.
_restore_rng_stateRestore RNG states from a saved state.
init_all_rngInitialize RNGs for Python, NumPy, and PyTorch (incl. CUDA) with a seed.

API

class nemo_automodel.components.training.rng.RNGState(
    random_rng_state: tuple[int, tuple[int, ...], float | None],
    np_rng_state: tuple[str, numpy.ndarray, int, int, float],
    torch_rng_state: torch.Tensor,
    cuda_rng_state: list[torch.Tensor]
)
Dataclass

Legacy RNG state kept for trusted pickle-based checkpoint restores.

cuda_rng_state
list[Tensor]
np_rng_state
tuple[str, ndarray, int, int, float]
random_rng_state
tuple[int, tuple[int, ...], float | None]
torch_rng_state
Tensor
class nemo_automodel.components.training.rng.ScopedRNG(
    seed: int = 95050,
    ranked: bool = False
)

Context manager for reproducible RNG states across random, NumPy, and PyTorch.

nemo_automodel.components.training.rng.ScopedRNG.__enter__()

Save current RNG states.

nemo_automodel.components.training.rng.ScopedRNG.__exit__(
    exc_type,
    exc_value,
    traceback
)

Restore RNG states on context exit.

class nemo_automodel.components.training.rng.StatefulRNG(
    seed: int,
    ranked: bool = False
)

RNG manager for reproducible RNG states across random, NumPy, and PyTorch.

nemo_automodel.components.training.rng.StatefulRNG.load_state_dict(
    state: nemo_automodel.components.training.rng._RNGState | nemo_automodel.components.training.rng.RNGState
) -> None

Restore RNG states from a saved state.

Parameters:

state
dict

RNG states as returned by state_dict().

nemo_automodel.components.training.rng.StatefulRNG.state_dict() -> nemo_automodel.components.training.rng._RNGState

Get current RNG states.

Returns: _RNGState

RNG states for random, NumPy, and PyTorch.

class nemo_automodel.components.training.rng._RNGState

Bases: typing.TypedDict

Weights-only-safe snapshot of Python, NumPy, Torch, and CUDA RNG states.

cuda_rng_state
list[Tensor]
np_bit_generator
str
np_cached_gaussian
float
np_has_gauss
int
np_keys
Tensor
np_position
int
random_rng_state
tuple[int, tuple[int, ...], float | None]
torch_rng_state
Tensor
nemo_automodel.components.training.rng._get_rng_state() -> nemo_automodel.components.training.rng._RNGState

Get current RNG states.

Returns: _RNGState

RNG states represented only by primitives and tensors so the state can

nemo_automodel.components.training.rng._restore_rng_state(
    state: nemo_automodel.components.training.rng._RNGState | nemo_automodel.components.training.rng.RNGState
) -> None

Restore RNG states from a saved state.

Parameters:

state
_RNGState | RNGState

Current weights-only-safe RNG state or legacy RNG state loaded from a trusted pickle-based checkpoint.

nemo_automodel.components.training.rng.init_all_rng(
    seed: int,
    ranked: bool = False
)

Initialize RNGs for Python, NumPy, and PyTorch (incl. CUDA) with a seed.

Parameters:

seed
int

Base seed value.

ranked
boolDefaults to False

Adjust seed by process rank if True.