nemo_automodel.training.rng#

Module Contents#

Classes#

StatefulRNG

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

Functions#

init_all_rng

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

API#

nemo_automodel.training.rng.init_all_rng(seed: int, ranked: bool = False)[source]#

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

Parameters:

  • seed (int) – Base seed value.

  • ranked (bool) – Adjust seed by process rank if True.

class nemo_automodel.training.rng.StatefulRNG(seed: int, ranked: bool = False)[source]#

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

Initialization

Initialize and optionally rank-adjust RNGs with a given seed.

Parameters:

  • seed (int) – Base seed for RNGs.

  • ranked (bool) – Adjust seed based on process rank.

__del__()[source]#
state_dict()[source]#

Get current RNG states.

Returns:

RNG states for random, NumPy, and PyTorch.

Return type:

dict

load_state_dict(state)[source]#

Restore RNG states from a saved state.

Parameters:

state (dict) – RNG states as returned by state_dict().

__enter__()[source]#

Save current RNG states.

__exit__(exc_type, exc_value, traceback)[source]#

Restore RNG states on context exit.