nemo_automodel.components.training.rng

Compatibility imports for shared training and data-construction RNG utilities.

Keep this module path available for existing callers and legacy RNGState pickles.

Module Contents

Classes

Name	Description
`RNGState`	Legacy RNG state kept for trusted pickle-based checkpoint restores.
`ScopedRNG`	Context manager for reproducible RNG states across random, NumPy, and PyTorch.
`StatefulRNG`	RNG manager for reproducible RNG states across random, NumPy, and PyTorch.

Functions

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

API

class nemo_automodel.shared.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.shared.rng.ScopedRNG(
    seed: int = 95050,
    ranked: bool = False
)

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

_saved_state

_RNGState | None = None

nemo_automodel.shared.rng.ScopedRNG.__enter__() -> nemo_automodel.shared.rng.ScopedRNG

Save current RNG states.

nemo_automodel.shared.rng.ScopedRNG.__exit__(
    exc_type: type[BaseException] | None,
    exc_value: BaseException | None,
    traceback: types.TracebackType | None
) -> None

Restore RNG states on context exit.

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

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

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

Restore RNG states from a saved state.

Parameters:

state

dict

RNG states as returned by state_dict().

nemo_automodel.shared.rng.StatefulRNG.state_dict() -> nemo_automodel.shared.rng._RNGState

Get current RNG states.

Returns: _RNGState

RNG states for random, NumPy, and PyTorch.

nemo_automodel.shared.rng.init_all_rng(
    seed: int,
    ranked: bool = False
) -> None

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.