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`	Typed view over the `prewarm` recipe config section.
`_GDNAttention`	Structural type of a gated-delta-net attention module.

Functions

Name	Description
`_collect_gdn_autotune_shapes`	Discover the gated-delta-net kernel shapes present in `model_parts`.
`_prewarm_comm_groups`	Eagerly create the NCCL communicators gradient-norm clipping will use.
`_prewarm_cublas_backward`	Initialize cuBLAS/cuBLASLt backward-pass state before real activations exist.
`_prewarm_fla_gdn_autotune`	Pre-populate fla gated-delta-net autotune caches before large activations exist.
`_prewarm_fla_gdn_cp_kernels`	Warm the fla context-parallel GDN backward kernels (best-effort).
`_prewarm_fla_gdn_end_to_end`	Run a tiny fwd+bwd through `chunk_gated_delta_rule` per unique shape.
`_resolve_cuda_device`	Normalize `device` and return it if it is a usable CUDA device, else None.
`_triton_kernel_accepts`	Check that a (possibly wrapped) Triton kernel accepts the expected launch arguments.

Data

_FLA_MERGE_FWD_BWD_KERNEL_ARGS

_FLA_PRE_PROCESS_BWD_KERNEL_ARGS

logger

API

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.

comm_groups

bool = False

cublas_backward

bool = False

fla_gdn_autotune

bool = False

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:

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.

pp_mesh

DeviceMesh | NoneDefaults to None

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

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

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

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

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 | NoneDefaults to 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.

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

intDefaults to 16

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

Returns: bool

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

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:

the context-parallel backward kernels (pre-process and state merge), which the end-to-end warmup below cannot reach, and
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

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

intDefaults to 64

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

Returns: bool

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

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

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:

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

Returns: bool

True if the warmup ran for at least one shape.

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.

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.

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

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

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