> 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);
* 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
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 a tiny fwd+bwd through `chunk_gated_delta_rule` per unique shape.                |
| [`_resolve_cuda_device`](#nemo_automodel-components-training-prewarm-_resolve_cuda_device)                 | Normalize `device` and return it if it is a usable CUDA device, else None.           |
| [`_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,
    comm_groups: bool = False
)
```

Dataclass

Typed view over the `prewarm` recipe config section.

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

Run the enabled prewarms.

**Parameters:**

The (already parallelized) model parts.

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

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.

```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`. The fla autotune caches are keyed on (H, K, V\[,
BT]) -- never on sequence length or batch -- so one tiny warmup per unique
shape covers the real workload.

**Parameters:**

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 whose DTensor parameters define the groups
(every non-`Replicate` placement dim of each parameter's mesh).

Device for the scalar warmup all-reduce; falls back to the
current CUDA device (or CPU) when 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:**

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

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
) -> 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 are shape-based -- (H, K, V) and the
chunk size, never the sequence length -- so launching each kernel once
here with tiny tensors caches the winning configs for the real workload.

Two warmups 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 warmup below cannot reach, and
2. a tiny end-to-end fwd+bwd through the public `chunk_gated_delta_rule`
   op, which covers every autotuned kernel in its 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 to scan for GDN modules.

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

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

**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:**

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

Target CUDA device.

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
) -> bool
```

Run a tiny fwd+bwd through `chunk_gated_delta_rule` per unique shape.

This caches the autotune config of every kernel in the op's call graph
(chunk fwd, dqkwg backward, dhu pre-process, ...) while the allocator pool
is empty. Must run with grad enabled so the backward kernels fire.

**Parameters:**

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

Target CUDA device.

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

**Returns:** `bool`

True if the warmup ran for at least one shape.

```python
nemo_automodel.components.training.prewarm._resolve_cuda_device(
    device: torch.device | int | str | None,
    label: str
) -> torch.device | None
```

Normalize `device` and return it if it is a usable CUDA device, else None.

```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:**

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

Keyword-argument names the prewarm launch will pass.

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__)
```