nemo_automodel.components.models.step5.layers

Trainable Step-5 decoder layers and context-parallel attention.

Module Contents

Classes

Name	Description
`Step5Attention`	GQA with windowed or CSA-selected keys and contiguous query-shard CP.
`Step5DecoderLayer`	Pre-norm decoder with FP32 residual accumulation and optional shared MoE.
`Step5MLP`	SwiGLU with separate gate/up projections and FP32 asymmetric clamping.
`Step5RMSNorm`	Fuse the main-stream norm while retaining FP32 additive scale parameters.

Functions

Name	Description
`_activation_rms_norm`	-
`_grouped_rms_norm`	Normalize grouped query rows without copying interleaved QKV gaps.
`_residual_rms_norm`	-
`_rms_norm`	Normalize […, width] in FP32 before rounding once to the output dtype.
`_rms_norm_with_rstd`	Normalize in FP32 and retain the row statistic needed by backward.
`clamped_swiglu`	Fuse FP32 SwiGLU and its derivative without FP32 activation temporaries.
`fused_attention`	Run the shared compiled FlexAttention kernel without a dense token mask.

Data

region_replay

API

class nemo_automodel.components.models.step5.layers.Step5Attention(
    config: nemo_automodel.components.models.step5.configuration.Step5Config,
    layer_idx: int,
    backend: nemo_automodel.components.models.common.BackendConfig,
    dtype: torch.dtype
)

Bases: Module

GQA with windowed or CSA-selected keys and contiguous query-shard CP.

cp_group

ProcessGroup | None = None

g_proj

head_dim

= config.head_dim

heads_per_group

= self.num_heads // self.num_kv_heads

indexer

k_norm

num_heads

= config.num_attention_heads

num_kv_heads

= config.num_attention_groups

o_proj

q_norm

qkv_proj

rotary_emb

= Step5RotaryEmbedding(config, layer_idx, span=span)

window

nemo_automodel.components.models.step5.layers.Step5Attention._select_regions(
    hidden: torch.Tensor,
    positions: torch.Tensor,
    key_positions: torch.Tensor,
    key_padding: torch.Tensor,
    query_start: int,
    sparse_metadata: nemo_automodel.components.models.step5.sparse.RegionMetadata | None
) -> nemo_automodel.components.models.step5.sparse.RegionSelection

Select once or replay integer routes during activation recomputation.

Parameters:

hidden

torch.Tensor

Local normalized states [batch, local_queries, hidden].

positions

torch.Tensor

Int64 local document positions [batch, local_queries].

key_positions

torch.Tensor

Int64 global document positions [batch, global_keys].

key_padding

torch.Tensor

Boolean global padding flags [batch, global_keys].

query_start

int

Physical start of the local query shard.

sparse_metadata

RegionMetadata | None

Layer-independent global/local region layout.

Returns: RegionSelection

Integer-only region selection; no projected K/V or activations are

nemo_automodel.components.models.step5.layers.Step5Attention.forward(
    hidden: torch.Tensor,
    positions: torch.Tensor,
    padding_mask: torch.Tensor,
    global_positions: torch.Tensor,
    global_padding: torch.Tensor,
    sparse_metadata: nemo_automodel.components.models.step5.sparse.RegionMetadata | None = None
) -> torch.Tensor

Attend local queries to valid global keys, preserving document boundaries.

Parameters:

hidden

torch.Tensor

Tensor of shape [batch, local_sequence, hidden].

positions

torch.Tensor

Tensor of shape [batch, local_sequence], document-local positions.

padding_mask

torch.Tensor

Boolean tensor of shape [batch, local_sequence], true for padding.

global_positions

torch.Tensor

Int tensor [batch, global_sequence], gathered once per model call.

global_padding

torch.Tensor

Boolean tensor [batch, global_sequence], true for padding.

Returns: torch.Tensor

Tensor of shape [batch, local_sequence, hidden].

nemo_automodel.components.models.step5.layers.Step5Attention.setup_cp_attention(
    cp_mesh: torch.distributed.device_mesh.DeviceMesh
) -> None

Attach the CP group through AutoModel’s model-owned attention hook.

class nemo_automodel.components.models.step5.layers.Step5DecoderLayer(
    config: nemo_automodel.components.models.step5.configuration.Step5Config,
    layer_idx: int,
    backend: nemo_automodel.components.models.common.BackendConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig,
    dtype: torch.dtype
)

Bases: Module

Pre-norm decoder with FP32 residual accumulation and optional shared MoE.

input_layernorm

moe

post_attention_layernorm

self_attn

share_expert

nemo_automodel.components.models.step5.layers.Step5DecoderLayer.forward(
    hidden: torch.Tensor,
    positions: torch.Tensor,
    padding_mask: torch.Tensor,
    global_positions: torch.Tensor,
    global_padding: torch.Tensor,
    compute_dtype: torch.dtype,
    sparse_metadata: nemo_automodel.components.models.step5.sparse.RegionMetadata | None = None
) -> torch.Tensor

Apply attention and feed-forward residual branches.

Parameters:

hidden

torch.Tensor

FP32 tensor of shape [batch, local_sequence, hidden].

positions

torch.Tensor

Tensor of shape [batch, local_sequence].

padding_mask

torch.Tensor

Boolean tensor of shape [batch, local_sequence], true for padding.

global_positions

torch.Tensor

Int tensor [batch, global_sequence], gathered once per model call.

global_padding

torch.Tensor

Boolean tensor [batch, global_sequence], true for padding.

compute_dtype

torch.dtype

Activation dtype from the completed embedding projection.

sparse_metadata

RegionMetadata | NoneDefaults to None

Shared global region layout and local query metadata.

Returns: torch.Tensor

FP32 tensor of shape [batch, local_sequence, hidden].

nemo_automodel.components.models.step5.layers.Step5DecoderLayer.nemo_checkpoint_context_fn(
    context_fn: collections.abc.Callable[[], tuple[contextlib.AbstractContextManager, contextlib.AbstractContextManager]] | None
) -> collections.abc.Callable[[], tuple[contextlib.AbstractContextManager, contextlib.AbstractContextManager]] | None

Compose the native AC context with one integer-route log per call.

Parameters:

context_fn

Callable[[], tuple[AbstractContextManager, AbstractContextManager]] | None

Existing plain/selective checkpoint context factory.

Returns: Callable[[], tuple[AbstractContextManager, AbstractContextManager]] | None

A factory with a distinct recorder for each live pipeline microbatch,

class nemo_automodel.components.models.step5.layers.Step5MLP(
    config: nemo_automodel.components.models.step5.configuration.Step5Config,
    backend: nemo_automodel.components.models.common.BackendConfig,
    width: int,
    limit: float,
    dtype: torch.dtype
)

Bases: Module

SwiGLU with separate gate/up projections and FP32 asymmetric clamping.

down_proj

gate_proj

up_proj

nemo_automodel.components.models.step5.layers.Step5MLP.forward(
    hidden: torch.Tensor
) -> torch.Tensor

Evaluate the dense or shared feed-forward block.

Parameters:

hidden

torch.Tensor

Tensor of shape [batch, sequence, hidden].

Returns: torch.Tensor

Tensor of shape [batch, sequence, hidden] in the input dtype.

class nemo_automodel.components.models.step5.layers.Step5RMSNorm(
    hidden_size: int,
    eps: float = 1e-05,
    compile_norm: bool = False
)

Bases: Step3p5RMSNorm

Fuse the main-stream norm while retaining FP32 additive scale parameters.

Indexer normalization deliberately uses its separate implementation: its checkpoint requires rounding the scale to the activation dtype first.

Parameters:

hidden_size

int

Normalized feature width.

eps

floatDefaults to 1e-05

Additive variance epsilon.

compile_norm

boolDefaults to False

Enable the compiled CUDA core. CPU always uses eager math.

nemo_automodel.components.models.step5.layers.Step5RMSNorm._compiled_forward(
    hidden: torch.Tensor,
    dtype: torch.dtype
) -> torch.Tensor

Run the compiled CUDA norm on […, hidden_size] and restore the input shape.

nemo_automodel.components.models.step5.layers.Step5RMSNorm.forward(
    hidden: torch.Tensor,
    output_dtype: torch.dtype | None = None
) -> torch.Tensor

Normalize […, hidden_size] before the optional projection-input cast.

Parameters:

hidden

torch.Tensor

Activations […, hidden_size] with arbitrary leading axes. Grouped queries use [batch, sequence, groups, heads_per_group, hidden_size] and may be a strided slice of interleaved QKV. Decoder residual inputs remain FP32 through normalization.

output_dtype

torch.dtype | NoneDefaults to None

Final output dtype; defaults to the input dtype.

Returns: torch.Tensor

Normalized activations with the input shape. The learned FP32 scale

nemo_automodel.components.models.step5.layers._activation_rms_norm(
    hidden: torch.Tensor,
    weight: torch.Tensor,
    eps: float,
    output_dtype: torch.dtype
) -> tuple[torch.Tensor, torch.Tensor]

nemo_automodel.components.models.step5.layers._grouped_rms_norm(
    hidden: torch.Tensor,
    weight: torch.Tensor,
    eps: float,
    output_dtype: torch.dtype
) -> tuple[torch.Tensor, torch.Tensor]

Normalize grouped query rows without copying interleaved QKV gaps.

Parameters:

hidden

torch.Tensor

Query tensor [token_groups, heads_per_group, channels], where token_groups flattens batch, sequence and KV group. The outer stride includes the interleaved K/V rows, and channels has unit stride.

weight

torch.Tensor

FP32 additive scale [channels].

eps

float

Additive variance epsilon.

output_dtype

torch.dtype

Final normalized activation dtype.

Returns: torch.Tensor

Normalized query [token_groups, heads_per_group, channels] and FP32

nemo_automodel.components.models.step5.layers._residual_rms_norm(
    hidden: torch.Tensor,
    weight: torch.Tensor,
    eps: float,
    output_dtype: torch.dtype
) -> torch.Tensor

nemo_automodel.components.models.step5.layers._rms_norm(
    hidden: torch.Tensor,
    weight: torch.Tensor,
    eps: float,
    output_dtype: torch.dtype
) -> torch.Tensor

Normalize […, width] in FP32 before rounding once to the output dtype.

nemo_automodel.components.models.step5.layers._rms_norm_with_rstd(
    hidden: torch.Tensor,
    weight: torch.Tensor,
    eps: float,
    output_dtype: torch.dtype
) -> tuple[torch.Tensor, torch.Tensor]

Normalize in FP32 and retain the row statistic needed by backward.

nemo_automodel.components.models.step5.layers.clamped_swiglu(
    gate: torch.Tensor,
    up: torch.Tensor,
    limit: float
) -> torch.Tensor

Fuse FP32 SwiGLU and its derivative without FP32 activation temporaries.

Parameters:

gate

torch.Tensor

Projection tensor [batch, sequence, intermediate].

torch.Tensor

Projection tensor with the same shape and dtype.

limit

float

Positive clamp bound; non-positive values disable clamping.

Returns: torch.Tensor

Tensor with the same shape and dtype as gate.

nemo_automodel.components.models.step5.layers.fused_attention(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    positions: torch.Tensor,
    key_positions: torch.Tensor,
    key_padding: torch.Tensor,
    documents: torch.Tensor,
    query_start: int,
    window: int | None,
    selection: nemo_automodel.components.models.step5.sparse.RegionSelection | None
) -> torch.Tensor

Run the shared compiled FlexAttention kernel without a dense token mask.

Parameters:

query

torch.Tensor

CUDA tensor [batch, local_queries, query_heads, head_dim].

key

torch.Tensor

CUDA tensor [batch, global_keys, kv_heads, head_dim].

value

torch.Tensor

CUDA tensor with the same shape and dtype as key.

positions

torch.Tensor

Document-local positions [batch, local_queries].

key_positions

torch.Tensor

Document-local positions [batch, global_keys].

key_padding

torch.Tensor

Boolean tensor [batch, global_keys], true for padding.

documents

torch.Tensor

Document ids [batch, global_keys].

query_start

int

Global offset of the contiguous query shard.

window

int | None

Sliding window size, or None for CSA layers.

selection

RegionSelection | None

Compact sorted CSA region ids, or None for sliding attention.

Returns: torch.Tensor

Attention output [batch, local_queries, query_heads, head_dim].

nemo_automodel.components.models.step5.layers.region_replay: RecomputeReplay[RegionSelection] = RecomputeReplay('Step5 CSA route')