> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.components.training.rng

## Module Contents

### Classes

| Name                                                                 | Description                                                                    |
| -------------------------------------------------------------------- | ------------------------------------------------------------------------------ |
| [`RNGState`](#nemo_automodel-components-training-rng-RNGState)       | Legacy RNG state kept for trusted pickle-based checkpoint restores.            |
| [`ScopedRNG`](#nemo_automodel-components-training-rng-ScopedRNG)     | Context manager for reproducible RNG states across random, NumPy, and PyTorch. |
| [`StatefulRNG`](#nemo_automodel-components-training-rng-StatefulRNG) | RNG manager for reproducible RNG states across random, NumPy, and PyTorch.     |
| [`_RNGState`](#nemo_automodel-components-training-rng-_RNGState)     | Weights-only-safe snapshot of Python, NumPy, Torch, and CUDA RNG states.       |

### Functions

| Name                                                                               | Description                                                              |
| ---------------------------------------------------------------------------------- | ------------------------------------------------------------------------ |
| [`_get_rng_state`](#nemo_automodel-components-training-rng-_get_rng_state)         | Get current RNG states.                                                  |
| [`_restore_rng_state`](#nemo_automodel-components-training-rng-_restore_rng_state) | Restore RNG states from a saved state.                                   |
| [`init_all_rng`](#nemo_automodel-components-training-rng-init_all_rng)             | Initialize RNGs for Python, NumPy, and PyTorch (incl. CUDA) with a seed. |

### API

```python
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`

---

```python
class nemo_automodel.components.training.rng.ScopedRNG(
    seed: int = 95050,
    ranked: bool = False
)
```

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

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

Save current RNG states.

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

Restore RNG states on context exit.

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

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

```python
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().

---

```python
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.

```python
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`

---

```python
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

```python
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.

---

```python
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`** `bool` — default: False

Adjust seed by process rank if True.

---