> 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.prewarm

Setup-time prewarm utilities.

Several CUDA runtime components initialize lazily on first use:

* cuBLAS/cuBLASLt allocate their workspaces on the first backward matmul of
  each dtype;
* Triton autotuners benchmark every candidate config on a kernel's first
  launch (flash-linear-attention's gated-delta-rule backward kernels are the
  heavy case, as are Mamba SSD backward kernels);
* NCCL creates a process group's communicator (and cudaMallocs its scratch
  buffers outside the torch caching-allocator pool) on the group's first
  collective.

When that first use happens during step 1 -- at peak activation/gradient
memory -- the out-of-pool allocation can fail with `NCCL ... Cuda failure 2
'out of memory'` or `Triton Error [CUDA]: out of memory` even though the
run would otherwise fit. Running these warmups at setup time, while the
caching-allocator pool is still small, moves the one-time initialization
costs and their memory spikes out of the first optimization step.

Prewarms are opt-in from the recipe config::

prewarm:
cublas\_backward: true
fla\_gdn\_autotune: true
mamba\_ssd\_autotune: true
comm\_groups: true

All prewarms are best-effort: failures are logged and never abort setup.
Warmups that generate synthetic inputs preserve the PyTorch RNG state so
enabling them does not change the subsequent training trajectory.

## Module Contents

### Classes

| Name                                                                         | Description                                            |
| ---------------------------------------------------------------------------- | ------------------------------------------------------ |
| [`PrewarmConfig`](#nemo_automodel-components-training-prewarm-PrewarmConfig) | Typed view over the `prewarm` recipe config section.   |
| [`_GDNAttention`](#nemo_automodel-components-training-prewarm-_GDNAttention) | Structural type of a gated-delta-net attention module. |

### Functions

| Name                                                                                                       | Description                                                                          |
| ---------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------ |
| [`_collect_gdn_autotune_shapes`](#nemo_automodel-components-training-prewarm-_collect_gdn_autotune_shapes) | Discover the gated-delta-net kernel shapes present in `model_parts`.                 |
| [`_prewarm_comm_groups`](#nemo_automodel-components-training-prewarm-_prewarm_comm_groups)                 | Eagerly create the NCCL communicators gradient-norm clipping will use.               |
| [`_prewarm_cublas_backward`](#nemo_automodel-components-training-prewarm-_prewarm_cublas_backward)         | Initialize cuBLAS/cuBLASLt backward-pass state before real activations exist.        |
| [`_prewarm_fla_gdn_autotune`](#nemo_automodel-components-training-prewarm-_prewarm_fla_gdn_autotune)       | Pre-populate fla gated-delta-net autotune caches before large activations exist.     |
| [`_prewarm_fla_gdn_cp_kernels`](#nemo_automodel-components-training-prewarm-_prewarm_fla_gdn_cp_kernels)   | Warm the fla context-parallel GDN backward kernels (best-effort).                    |
| [`_prewarm_fla_gdn_end_to_end`](#nemo_automodel-components-training-prewarm-_prewarm_fla_gdn_end_to_end)   | Run dense and packed fwd+bwd warmups through `chunk_gated_delta_rule`.               |
| [`_triton_kernel_accepts`](#nemo_automodel-components-training-prewarm-_triton_kernel_accepts)             | Check that a (possibly wrapped) Triton kernel accepts the expected launch arguments. |

### Data

[`_FLA_MERGE_FWD_BWD_KERNEL_ARGS`](#nemo_automodel-components-training-prewarm-_FLA_MERGE_FWD_BWD_KERNEL_ARGS)

[`_FLA_PRE_PROCESS_BWD_KERNEL_ARGS`](#nemo_automodel-components-training-prewarm-_FLA_PRE_PROCESS_BWD_KERNEL_ARGS)

[`logger`](#nemo_automodel-components-training-prewarm-logger)

### API

```python
class nemo_automodel.components.training.prewarm.PrewarmConfig(
    cublas_backward: bool = False,
    fla_gdn_autotune: bool = False,
    mamba_ssd_autotune: bool = False,
    comm_groups: bool = False
)
```

Dataclass

Typed view over the `prewarm` recipe config section.

**`comm_groups`** `bool = False`

---

**`cublas_backward`** `bool = False`

---

**`fla_gdn_autotune`** `bool = False`

---

**`mamba_ssd_autotune`** `bool = False`

---

```python
nemo_automodel.components.training.prewarm.PrewarmConfig.apply(
    model_parts: list[torch.nn.Module],
    device: torch.device | int | str | None,
    batch_size: int = 1,
    pp_mesh: torch.distributed.device_mesh.DeviceMesh | None = None
) -> None
```

Run the enabled prewarms.

**Parameters:**

**`model_parts`** `list[torch.nn.Module]`

The (already parallelized) model parts.

---

**`device`** `torch.device | int | str | None`

The device assigned to this rank, or None when no
accelerator is available.

---

**`batch_size`** `int` — default: 1

Per-rank training batch size. FLA keys its dense GDN
cumulative-sum autotuner on this value.

---

**`pp_mesh`** `DeviceMesh | None` — default: None

The pipeline-parallel submesh, if pipeline parallelism is
enabled (its process group is warmed for the grad-norm
all-reduce).

---

```python
class nemo_automodel.components.training.prewarm._GDNAttention()
```

Protocol

Structural type of a gated-delta-net attention module.

Matches modules (e.g. the qwen3\_next / qwen3\_5\_moe GDN attention layers)
that expose the head geometry needed to reconstruct the fla kernel shapes
plus the `chunk_gated_delta_rule` op backed by those kernels.

**`chunk_gated_delta_rule`** `Callable[..., Any]`

---

**`head_k_dim`** `int`

---

**`head_v_dim`** `int`

---

**`num_v_heads`** `int`

---

```python
nemo_automodel.components.training.prewarm._collect_gdn_autotune_shapes(
    model_parts: list[torch.nn.Module]
) -> dict[tuple[int, int, int, torch.dtype], str]
```

Discover the gated-delta-net kernel shapes present in `model_parts`.

A module counts as a GDN attention module when it structurally matches
:class:`_GDNAttention`. Most fla autotune caches are keyed on
`(H, K, V[, BT])` rather than sequence length. Dense cumulative-sum
kernels additionally key on batch size; that runtime value is handled by
the warmup rather than shape discovery.

**Parameters:**

**`model_parts`** `list[torch.nn.Module]`

Model parts to scan for GDN modules.

---

**Returns:** `dict[tuple[int, int, int, torch.dtype], str]`

Mapping of `(num_v_heads, head_k_dim, head_v_dim, dtype)` to the

```python
nemo_automodel.components.training.prewarm._prewarm_comm_groups(
    model_parts: list[torch.nn.Module],
    device: torch.device | int | str | None,
    pp_mesh: torch.distributed.device_mesh.DeviceMesh | None = None
) -> int
```

Eagerly create the NCCL communicators gradient-norm clipping will use.

torch creates a process group's NCCL communicator lazily on its first
collective. `clip_grad_norm` all-reduces once per Shard mesh dim of each
gradient DTensor group, and for some dims that first collective runs at
step-1 peak memory; NCCL then cudaMallocs its scratch buffers outside the
torch pool and can die with `Cuda failure 2 'out of memory'` after a
clean forward+backward. Warm exactly those groups here, while the torch
pool is still small, by replaying `clip_grad_norm`'s own (mesh,
shard-dim) enumeration and issuing a one-element all-reduce per group.

**Parameters:**

**`model_parts`** `list[torch.nn.Module]`

Model parts whose DTensor parameters define the groups
(every non-`Replicate` placement dim of each parameter's mesh).

---

**`device`** `torch.device | int | str | None`

Device for the scalar warmup all-reduce; falls back to the
current CUDA device (or CPU) when None.

---

**`pp_mesh`** `DeviceMesh | None` — default: None

Pipeline-parallel submesh, if enabled. `clip_grad_norm`
also all-reduces the total norm across the PP group, but
parameters are never sharded along pp, so the placement
enumeration alone cannot discover that group.

---

**Returns:** `int`

The number of process groups warmed.

```python
nemo_automodel.components.training.prewarm._prewarm_cublas_backward(
    device: torch.device | int | str | None,
    size: int = 16
) -> bool
```

Initialize cuBLAS/cuBLASLt backward-pass state before real activations exist.

Runs a tiny fwd+bwd matmul per dtype so the library handles and workspaces
are allocated while the allocator pool is small, instead of at step-1 peak.

**Parameters:**

**`device`** `torch.device | int | str | None`

Target CUDA device (skipped when None or not CUDA).

---

**`size`** `int` — default: 16

Side length of the square `[size, size]` warmup matmul operands.

---

**Returns:** `bool`

True if the prewarm ran, False if it was skipped.

```python
nemo_automodel.components.training.prewarm._prewarm_fla_gdn_autotune(
    model_parts: list[torch.nn.Module],
    device: torch.device | int | str | None,
    seq_len: int = 64,
    batch_size: int = 1
) -> bool
```

Pre-populate fla gated-delta-net autotune caches before large activations exist.

On its first launch each Triton-autotuned kernel benchmarks all candidate
configs; when that first launch is the step-1 backward at peak memory, the
benchmark allocations can fail with `Triton Error [CUDA]: out of
memory`. The autotune cache keys include shape, chunk size, and
variable-length mode, but not sequence length. Dense cumulative-sum
kernels additionally key on batch size, so the real per-rank batch size
is used for that path.

Three warmup paths run per unique GDN shape found in `model_parts`:

1. the context-parallel backward kernels (pre-process and state merge),
   which the end-to-end warmups below cannot reach,
2. a dense end-to-end fwd+bwd using the real per-rank batch size, and
3. a packed end-to-end fwd+bwd with `cu_seqlens`.

The end-to-end paths cover every autotuned kernel in the public
`chunk_gated_delta_rule` call graph. Per-kernel prewarms alone are
whack-a-mole: warming one kernel just moves the step-1 autotune OOM to the
next cold kernel.

**Parameters:**

**`model_parts`** `list[torch.nn.Module]`

Model parts to scan for GDN modules.

---

**`device`** `torch.device | int | str | None`

Target CUDA device (skipped when None or not CUDA).

---

**`seq_len`** `int` — default: 64

Warmup sequence length (not part of the autotune key).

---

**`batch_size`** `int` — default: 1

Per-rank dense training batch size.

---

**Returns:** `bool`

True if at least the end-to-end prewarm ran, False if skipped.

```python
nemo_automodel.components.training.prewarm._prewarm_fla_gdn_cp_kernels(
    shapes: dict[tuple[int, int, int, torch.dtype], str],
    device: torch.device,
    seq_len: int
) -> None
```

Warm the fla context-parallel GDN backward kernels (best-effort).

These kernels only fire on the context-parallel code path, which the
end-to-end warmup in :func:`_prewarm_fla_gdn_end_to_end` does not reach,
so they are launched once directly with zero-filled tiny tensors (the
autotuner only measures timing; values are irrelevant). fla builds without
CP kernels are skipped, as are kernels whose parameter list no longer
matches the launch contract here (the kernels are private fla API, so the
signature is validated before every launch attempt).

**Parameters:**

**`shapes`** `dict[tuple[int, int, int, torch.dtype], str]`

Mapping of `(num_v_heads, head_k_dim, head_v_dim, dtype)` to
a module name, as returned by :func:`_collect_gdn_autotune_shapes`.

---

**`device`** `torch.device`

Target CUDA device.

---

**`seq_len`** `int`

Warmup sequence length (not part of the autotune key).

---

```python
nemo_automodel.components.training.prewarm._prewarm_fla_gdn_end_to_end(
    shapes: dict[tuple[int, int, int, torch.dtype], str],
    device: torch.device,
    seq_len: int,
    batch_size: int = 1
) -> bool
```

Run dense and packed fwd+bwd warmups through `chunk_gated_delta_rule`.

This caches the autotune config of every kernel in the op's call graph
(chunk fwd, dqkwg backward, dhu pre-process, ...) for both
`IS_VARLEN=False` and `IS_VARLEN=True` while the allocator pool is
empty. Must run with grad enabled so the backward kernels fire. The dense
path uses the real per-rank batch size because FLA's cumulative-sum
autotuner includes it in its cache key; packed GDN always flattens to
batch size 1.

**Parameters:**

**`shapes`** `dict[tuple[int, int, int, torch.dtype], str]`

Mapping of `(num_v_heads, head_k_dim, head_v_dim, dtype)` to
a module name, as returned by :func:`_collect_gdn_autotune_shapes`.

---

**`device`** `torch.device`

Target CUDA device.

---

**`seq_len`** `int`

Warmup sequence length (not part of the autotune key).

---

**`batch_size`** `int` — default: 1

Per-rank batch size for the dense path.

---

**Returns:** `bool`

True if the warmup ran for at least one shape.

```python
nemo_automodel.components.training.prewarm._triton_kernel_accepts(
    kernel: object,
    expected_args: frozenset[str],
    kernel_name: str
) -> bool
```

Check that a (possibly wrapped) Triton kernel accepts the expected launch arguments.

Unwraps autotuner/heuristics layers via their `fn` attribute until an
object exposing the JITFunction's `arg_names` is found.

**Parameters:**

**`kernel`** `object`

The Triton kernel object (a JITFunction, or an Autotuner /
Heuristics wrapper around one).

---

**`expected_args`** `frozenset[str]`

Keyword-argument names the prewarm launch will pass.

---

**`kernel_name`** `str`

Kernel name used in log messages.

---

**Returns:** `bool`

True when every expected argument is in the kernel's parameter list.

```python
nemo_automodel.components.training.prewarm._FLA_MERGE_FWD_BWD_KERNEL_ARGS = frozenset(('h', 'ag_hm', 'pre_or_post_num_ranks', 'rank', 'seq_offsets', 'init_o...
```

```python
nemo_automodel.components.training.prewarm._FLA_PRE_PROCESS_BWD_KERNEL_ARGS = frozenset(('q', 'k', 'w', 'g', 'gk', 'do', 'dhm', 'dv', 'cu_seqlens', 'scale', '...
```

```python
nemo_automodel.components.training.prewarm.logger = logging.getLogger(__name__)
```