`core.models.audio.nemo_transformer_encoder`#

Module Contents#

Classes#

`GPTConfig`	Configuration for a GPT-style transformer (vendored NeMo helper).
`TransformerEncoderConfig`	Configuration for the audio TransformerEncoder and its blocks.
`FeedForward`	Two-layer position-wise feed-forward network with a 4x hidden expansion and GELU.
`GELU`	Tanh-approximation GELU activation (vendored NeMo helper).
`LayerNorm`	Layer normalization with learnable scale and shift (upcasts to fp32 internally).
`MultiHeadAttentionWithFA`	Multi-head attention using the flash-attention `flash_attn_func` backend.
`MultiHeadAttentionWithTE`	Multi-head attention using Transformer Engine’s DotProductAttention in THD layout.
`MultiHeadAttentionWithSDPA`	Multi-head attention using PyTorch `scaled_dot_product_attention`.
`MultiHeadAttention`	Multi-head attention with an explicit softmax and optional KV cache.
`TransformerBlock`	Single transformer encoder block: attention, feed-forward, and residual norms.
`ConvSubsampling`	Convolutional subsampling that reduces the temporal dimension by 4x.
`DepthwiseConvSubsampling`	Depthwise separable conv subsampling: reduces params by replacing standard Conv1d with depthwise (groups=channels) + pointwise (1x1) convolutions for the strided layers.
`StackingSubsampling`	Stack consecutive audio frames and project them to the encoder dimension.
`TransformerEncoder`	Audio transformer encoder: subsampling front-end followed by a transformer stack.

Functions#

`_flash_attn_func`	Lazy import shim for flash_attn so this module imports without the package.
`_get_te_dot_product_attention`	Lazy import of transformer_engine.pytorch.DotProductAttention.
`_resolve_attn_impl`	Resolve ‘auto’ to ‘te’ if transformer_engine is importable, else ‘sdpa’.
`_normalize_left_context`
`_causal_window_size`
`_causal_disallow_mask`	Return a bool mask where True means the key is not visible to the query.
`compute_rope_params`	Precompute the rotary position embedding cosine and sine tables.
`apply_rope`	Apply rotary position embeddings to `x` using precomputed cos and sin tables.

Data#

`logger`
`NGPTStackingSubsampling`

API#

core.models.audio.nemo_transformer_encoder.logger#: ‘getLogger(…)’

core.models.audio.nemo_transformer_encoder._flash_attn_func(*args, **kwargs)#: Lazy import shim for flash_attn so this module imports without the package.

core.models.audio.nemo_transformer_encoder._get_te_dot_product_attention()#: Lazy import of transformer_engine.pytorch.DotProductAttention.

core.models.audio.nemo_transformer_encoder._resolve_attn_impl(impl: str) → str#: Resolve ‘auto’ to ‘te’ if transformer_engine is importable, else ‘sdpa’.

class core.models.audio.nemo_transformer_encoder.GPTConfig#

Configuration for a GPT-style transformer (vendored NeMo helper).

vocab_size: int#: 50257

context_length: int#: 1024

emb_dim: int#: 768

n_heads: int#: 12

n_layers: int#: 12

drop_rate: int#: 0.1

qkv_bias: bool#: False

theta_base: int#: 10000

class core.models.audio.nemo_transformer_encoder.TransformerEncoderConfig#

Configuration for the audio TransformerEncoder and its blocks.

n_mels: int#: 80

d_model: int#: 512

n_heads: int#: 12

n_layers: int#: 12

drop_rate: float#: 0.1

qkv_bias: bool#: False

causal_mask: bool#: False

theta_base: int#: 10000

context_length: int#: 4096

qk_norm: bool#: False

attn_impl: str#: ‘auto’

recompute_layers: bool#: False

left_context: Optional[int]#: None

core.models.audio.nemo_transformer_encoder._normalize_left_context( left_context: Optional[int], ) → Optional[int]#

core.models.audio.nemo_transformer_encoder._causal_window_size( causal_mask: bool, left_context: Optional[int], ) → Optional[Tuple[int, int]]#

core.models.audio.nemo_transformer_encoder._causal_disallow_mask( query_len: int, key_len: int, left_context: Optional[int], device, ) → torch.Tensor#: Return a bool mask where True means the key is not visible to the query.

class core.models.audio.nemo_transformer_encoder.FeedForward(dim)#

Bases: torch.nn.Module

Two-layer position-wise feed-forward network with a 4x hidden expansion and GELU.

Initialization

forward(x)#: Apply the feed-forward network to x.

class core.models.audio.nemo_transformer_encoder.GELU#

Bases: torch.nn.Module

Tanh-approximation GELU activation (vendored NeMo helper).

Initialization

forward(x)#: Apply the tanh-approximation GELU to x.

class core.models.audio.nemo_transformer_encoder.LayerNorm(dim, eps=1e-05)#

Bases: torch.nn.Module

Layer normalization with learnable scale and shift (upcasts to fp32 internally).

Initialization

forward(x)#: Normalize x over its last dimension and apply scale and shift.

core.models.audio.nemo_transformer_encoder.compute_rope_params( head_dim, theta_base=10000, context_length=4096, dtype=torch.float32, )#: Precompute the rotary position embedding cosine and sine tables.

core.models.audio.nemo_transformer_encoder.apply_rope(x, cos, sin)#: Apply rotary position embeddings to x using precomputed cos and sin tables.

class core.models.audio.nemo_transformer_encoder.MultiHeadAttentionWithFA( dim_in, dim_out, dropout=0.0, qkv_bias=False, context_length=1024, num_heads=8, causal_mask=False, left_context=None, )#

Bases: torch.nn.Module

Multi-head attention using the flash-attention flash_attn_func backend.

Initialization

forward(x)#: Run flash-attention over x and project the result.

class core.models.audio.nemo_transformer_encoder.MultiHeadAttentionWithTE(

dim_in,

dim_out,

dropout=0.0,

qkv_bias=False,

num_heads=8,

causal_mask=False,

left_context=None,

qk_norm=False,

**_,

)#

Bases: torch.nn.Module

Multi-head attention using Transformer Engine’s DotProductAttention in THD layout.

Always runs attention in TE’s thd (packed) format so flash-attention handles variable-length sequences efficiently. Two callsite modes:

Unpacked: forward(x, lengths=…) where x is (B, T, C). The module packs valid tokens into THD using lengths, runs TE attention, then scatters results back into a (B, T, C) tensor (padding positions remain zero).
Packed: forward(x, packed_seq_params=…) where x is a flat (Ttot, C) (or (Ttot, 1, C)) tensor and the caller supplies cu_seqlens / max_seqlens. The module returns (Ttot, C) in the same packed layout.

Linear projection parameter names (w_query, w_key, w_value, out_proj) match the SDPA/FA variants so .nemo checkpoints load into either backend.

Initialization

forward(x, lengths=None, packed_seq_params=None, **_)#: Run TE attention in THD layout, packing/unpacking as needed, and project.

class core.models.audio.nemo_transformer_encoder.MultiHeadAttentionWithSDPA( dim_in, dim_out, dropout=0.0, qkv_bias=False, context_length=1024, num_heads=8, causal_mask=False, left_context=None, qk_norm=False, )#

Bases: torch.nn.Module

Multi-head attention using PyTorch scaled_dot_product_attention.

Initialization

forward(x, attn_mask=None, use_cache=False)#: Run scaled-dot-product attention over x and project the result.

class core.models.audio.nemo_transformer_encoder.MultiHeadAttention( dim_in, dim_out, dropout=0.0, qkv_bias=False, context_length=1024, num_heads=8, causal_mask=False, left_context=None, )#

Bases: torch.nn.Module

Multi-head attention with an explicit softmax and optional KV cache.

Initialization

forward(x, use_cache=False)#: Run masked multi-head attention over x, optionally using the KV cache.

reset_cache()#: Clear the cached keys and values.

class core.models.audio.nemo_transformer_encoder.TransformerBlock( cfg: core.models.audio.nemo_transformer_encoder.TransformerEncoderConfig, )#

Bases: torch.nn.Module

Single transformer encoder block: attention, feed-forward, and residual norms.

Initialization

forward( x, attn_mask=None, lengths=None, packed_seq_params=None, use_cache=False, )#: Apply attention and feed-forward sublayers with residual connections.

reset_cache()#: Reset KV cache on the attention module if it supports it (e.g. MultiHeadAttention).

class core.models.audio.nemo_transformer_encoder.ConvSubsampling(n_mels: int = 80, d_model: int = 512)#

Bases: torch.nn.Module

Convolutional subsampling that reduces the temporal dimension by 4x.

Initialization

forward(x, length)#: Subsample x by 4x and return the features and updated lengths.

class core.models.audio.nemo_transformer_encoder.DepthwiseConvSubsampling(n_mels: int = 80, d_model: int = 512)#

Bases: torch.nn.Module

Depthwise separable conv subsampling: reduces params by replacing standard Conv1d with depthwise (groups=channels) + pointwise (1x1) convolutions for the strided layers.

Initialization

forward(x, length)#: Subsample x by 4x via depthwise separable convs and update lengths.

class core.models.audio.nemo_transformer_encoder.StackingSubsampling( subsampling_factor: int, feat_in: int, feat_out: int, use_bias: bool = False, *, pad_mode: Literal[learned, zeros] = 'learned', )#

Bases: torch.nn.Module

Stack consecutive audio frames and project them to the encoder dimension.

Parameters:

subsampling_factor – Number of adjacent frames in each stack.
feat_in – Input feature dimension.
feat_out – Output encoder dimension.
use_bias – Whether the output projection has a bias.
pad_mode – How missing frames in a partial final stack are represented. "learned" preserves the legacy NGPT behavior and checkpoint layout; "zeros" matches NeMo’s FeatureStacking.

Initialization

_projection() → torch.nn.Linear#

_replace_zero_frames(x: torch.Tensor) → torch.Tensor#

_pad_partial_stack(x: torch.Tensor, pad_size: int) → torch.Tensor#

forward( x: torch.Tensor, length: torch.Tensor, ) → tuple[torch.Tensor, torch.Tensor]#

Stack and project a dense batch of padded audio features.

Parameters:

x – Audio features with shape [batch, features, time].
length – Valid input frame count for each sample.

Returns:

Projected features with shape [batch, time', feat_out] and the corresponding per-sample lengths.

forward_packed( x: torch.Tensor, lengths: torch.Tensor, ) → tuple[torch.Tensor, torch.Tensor]#

Stack only valid frames and return contiguous per-audio token runs.

x may be a padded [B, F, T] batch or already-packed [Ttotal, F] features. Each sample is padded only through its final partial stack, avoiding projection work on suffix padding.

core.models.audio.nemo_transformer_encoder.NGPTStackingSubsampling#: None

class core.models.audio.nemo_transformer_encoder.TransformerEncoder( n_mels: int = 80, d_model: int = 512, n_heads: int = 8, n_layers: int = 17, drop_rate: float = 0.1, qkv_bias: bool = False, causal_mask: bool = False, pre_encode: str = 'conv', nan_debug: bool = True, qk_norm: bool = False, subsampling_factor: int = 4, stacking_pad_mode: Literal[learned, zeros] = 'learned', attn_impl: str = 'auto', recompute_layers: bool = False, left_context: Optional[int] = None, )#

Bases: torch.nn.Module

Audio transformer encoder: subsampling front-end followed by a transformer stack.

Initialization

forward( audio_signal, length, packed_seq_params=None, return_packed: bool = False, )#

Parameters:

audio_signal (torch.Tensor) – (B, C, T) audio features.
length (torch.Tensor) – (B,) input frame counts.
packed_seq_params – optional caller-supplied PackedSeqParams. Unused by production callers today; reserved for future dataloader-side packing. When None and attn_impl==”te”, the encoder builds its own PackedSeqParams from length and runs the block stack on packed features for efficiency.

Returns:

(B, D_model, T’) with zero-padded positions when return_packed is false, otherwise (Ttot, D_model) with only valid post-subsampling positions. length (torch.Tensor): (B,) post-subsampling lengths.

Return type:

x (torch.Tensor)

_check_nan(x, name)#

reset_cache()#: Reset KV cache on every block (no-op for attention impls without cache).

freeze()#: Disable gradients for all encoder parameters.

unfreeze(partial=False)#: Enable gradients for all encoder parameters.

core.models.audio.nemo_transformer_encoder#

Module Contents#

Classes#

Functions#

Data#

API#

`core.models.audio.nemo_transformer_encoder`#