NVIDIA-NeMo/Speech · #16419
feat(asr): configurable blocks in the Transformer decoder
nemo/collections/asr/modules/transformer/transformer.py57 + / 2 −
@@ -14,7 +14,7 @@ # limitations under the License. from dataclasses import dataclass-from typing import Dict, List, Optional+from typing import Any, Dict, List, Optional import torch from omegaconf.omegaconf import MISSING, DictConfig@@ -24,7 +24,9 @@ from nemo.collections.asr.modules.transformer.transformer_decoders import TransformerDecoder, TransformerDecoderAdapter from nemo.collections.asr.modules.transformer.transformer_encoders import TransformerEncoder from nemo.collections.asr.modules.transformer.transformer_modules import TransformerEmbedding+from nemo.collections.asr.parts.submodules.multi_head_attention import RotaryPositionalEncoding from nemo.collections.asr.parts.utils import adapter_utils+from nemo.collections.common.parts.utils import ModuleBlock, build_module from nemo.core.classes.common import typecheck from nemo.core.classes.exportable import Exportable from nemo.core.classes.mixins import adapter_mixins@@ -68,6 +70,13 @@ class NeMoTransformerEncoderConfig(NeMoTransformerConfig): class NeMoTransformerDecoderConfig(NeMoTransformerConfig): r2l: bool = False + # configurable blocks (Hydra sections with _target_ and _partial_: true), see TransformerDecoderNM+ embedding: Optional[Any] = None+ norm: Optional[Any] = None+ self_attn: Optional[Any] = None+ cross_attn: Optional[Any] = None+ feed_forward: Optional[Any] = None+ class TransformerEncoderNM(EncoderModule, Exportable): def __init__(@@ -158,6 +167,36 @@ def input_example(self, max_batch=1, max_dim=256): class TransformerDecoderNM(DecoderModule, Exportable):+ """+ Transformer decoder with its token and position embedding, as a NeMo decoder module.++ Args:+ vocab_size: size of the vocabulary+ hidden_size: size of the embeddings in the model, also known as d_model+ num_layers: number of decoder blocks+ inner_size: number of neurons in the intermediate part of the feed-forward net+ num_attention_heads: number of heads in multi-head attention+ max_sequence_length: maximum allowed length of the input sequence+ num_token_types: number of different token types, 0 for none+ embedding_dropout: probability of dropout applied to embeddings+ learn_positional_encodings: whether to learn positional encodings or use fixed (sine-cosine) ones+ ffn_dropout: probability of dropout applied to the feed-forward output+ attn_score_dropout: probability of dropout applied to attention scores+ attn_layer_dropout: probability of dropout applied to the output of the attention layers,+ but before layer normalization+ hidden_act: activation function used between the two linear layers of the feed-forward net+ pre_ln: whether to apply the layer norms before the sub-layers instead of after them+ pre_ln_final_layer_norm: with `pre_ln`, whether a final layer norm follows the last block+ embedding: block of the embedding, default `TransformerEmbedding`,+ built with the vocabulary and embedding arguments above;+ a block is a class, a callable or a Hydra partial (``_target_`` with ``_partial_: true``),+ built from the arguments the parent supplies, so a block section only names what deviates+ norm: block of every layer norm of the decoder, see `TransformerDecoderBlock`+ self_attn: block of the self-attention of every block, see `TransformerDecoderBlock`+ cross_attn: block of the attention over the encoder states of every block, see `TransformerDecoderBlock`+ feed_forward: block of the feed-forward sub-layer of every block, see `TransformerDecoderBlock`+ """+ DECODER_TYPE: type = TransformerDecoder def __init__(@@ -177,6 +216,11 @@ def __init__( hidden_act: str = 'relu', pre_ln: bool = False, pre_ln_final_layer_norm: bool = True,+ embedding: ModuleBlock | None = None,+ norm: ModuleBlock | None = None,+ self_attn: ModuleBlock | None = None,+ cross_attn: ModuleBlock | None = None,+ feed_forward: ModuleBlock | None = None, ): super().__init__() @@ -188,7 +232,8 @@ def __init__( if pre_ln_final_layer_norm: self.num_states += 1 - self._embedding = TransformerEmbedding(+ self._embedding = build_module(+ embedding or TransformerEmbedding, vocab_size=self.vocab_size, hidden_size=self.hidden_size, max_sequence_length=max_sequence_length,@@ -208,6 +253,10 @@ def __init__( hidden_act=hidden_act, pre_ln=pre_ln, pre_ln_final_layer_norm=pre_ln_final_layer_norm,+ norm=norm,+ self_attn=self_attn,+ cross_attn=cross_attn,+ feed_forward=feed_forward, ) @typecheck()@@ -278,6 +327,12 @@ def input_example(self, max_batch=1, max_dim=256): def _prepare_for_export(self, **kwargs): self._decoder.diagonal = None self.return_mems = True+ # rotary tables grow with the key length in eager mode but are constants in the exported graph,+ # so they cover every position the decoder supports before the example is traced+ sample = next(self.parameters())+ for module in self._decoder.modules():+ if isinstance(module, RotaryPositionalEncoding):+ module.extend_pe(self.max_sequence_length, sample.device, sample.dtype) super()._prepare_for_export(**kwargs) @propertynemo/collections/asr/modules/transformer/transformer_decoders.py61 + / 13 −
@@ -24,6 +24,7 @@ from nemo.collections.asr.parts.submodules.adapters.attention_adapter_mixin import AttentionAdapterModuleMixin from nemo.collections.asr.parts.utils import adapter_utils from nemo.collections.common.parts import form_attention_mask+from nemo.collections.common.parts.utils import ModuleBlock, build_module from nemo.core.classes.mixins import adapter_mixins __all__ = ["TransformerDecoder"]@@ -43,6 +44,13 @@ class TransformerDecoderBlock(nn.Module, AttentionAdapterModuleMixin): attention layers, but before layer normalization ffn_dropout: probability of dropout applied to FFN output hidden_act: activation function used between two linear layers in FFN+ pre_ln: whether to apply the layer norms before the sub-layers instead of after them+ norm: block of the three layer norms, default `torch.nn.LayerNorm`;+ a block is a class, a callable or a Hydra partial (``_target_`` with ``_partial_: true``),+ built from the arguments the parent supplies, so a block section only names what deviates+ self_attn: block of the self-attention, default `MultiHeadAttention`+ cross_attn: block of the attention over the encoder states, default `MultiHeadAttention`+ feed_forward: block of the feed-forward sub-layer, default `PositionWiseFF` """ def __init__(@@ -55,23 +63,34 @@ def __init__( ffn_dropout: float = 0.0, hidden_act: str = "relu", pre_ln: bool = False,+ norm: ModuleBlock | None = None,+ self_attn: ModuleBlock | None = None,+ cross_attn: ModuleBlock | None = None,+ feed_forward: ModuleBlock | None = None, ): super().__init__() self.pre_ln = pre_ln- self.layer_norm_1 = nn.LayerNorm(hidden_size, eps=1e-5)- self.first_sub_layer = MultiHeadAttention(- hidden_size, num_attention_heads, attn_score_dropout, attn_layer_dropout+ norm = norm or nn.LayerNorm+ attention_defaults = dict(+ hidden_size=hidden_size,+ num_attention_heads=num_attention_heads,+ attn_score_dropout=attn_score_dropout,+ attn_layer_dropout=attn_layer_dropout, )- self.layer_norm_2 = nn.LayerNorm(hidden_size, eps=1e-5)- self.second_sub_layer = MultiHeadAttention(- hidden_size,- num_attention_heads,- attn_score_dropout,- attn_layer_dropout,- return_xatt_scores=True,+ self.layer_norm_1 = build_module(norm, normalized_shape=hidden_size)+ self.first_sub_layer = build_module(self_attn or MultiHeadAttention, **attention_defaults)+ self.layer_norm_2 = build_module(norm, normalized_shape=hidden_size)+ self.second_sub_layer = build_module(+ cross_attn or MultiHeadAttention, return_xatt_scores=True, **attention_defaults+ )+ self.layer_norm_3 = build_module(norm, normalized_shape=hidden_size)+ self.third_sub_layer = build_module(+ feed_forward or PositionWiseFF,+ hidden_size=hidden_size,+ inner_size=inner_size,+ ffn_dropout=ffn_dropout,+ hidden_act=hidden_act, )- self.layer_norm_3 = nn.LayerNorm(hidden_size, eps=1e-5)- self.third_sub_layer = PositionWiseFF(hidden_size, inner_size, ffn_dropout, hidden_act) # Information for the adapter module mixin self.self_attention_model = "transf_abs"@@ -183,6 +202,27 @@ def get_accepted_adapter_types(self) -> Set[type]: class TransformerDecoder(nn.Module):+ """+ Stack of Transformer decoder blocks with an optional final layer norm.++ Args:+ num_layers: number of decoder blocks+ hidden_size: size of the embeddings in the model, also known as d_model+ inner_size: number of neurons in the intermediate part of the feed-forward net+ num_attention_heads: number of heads in multi-head attention+ attn_score_dropout: probability of dropout applied to attention scores+ attn_layer_dropout: probability of dropout applied to the output of the attention layers,+ but before layer normalization+ ffn_dropout: probability of dropout applied to the feed-forward output+ hidden_act: activation function used between the two linear layers of the feed-forward net+ pre_ln: whether to apply the layer norms before the sub-layers instead of after them+ pre_ln_final_layer_norm: with `pre_ln`, whether a final layer norm follows the last block+ norm: block of every layer norm, including the final one, see `TransformerDecoderBlock`+ self_attn: block of the self-attention of every block, see `TransformerDecoderBlock`+ cross_attn: block of the attention over the encoder states of every block, see `TransformerDecoderBlock`+ feed_forward: block of the feed-forward sub-layer of every block, see `TransformerDecoderBlock`+ """+ def __init__( self, num_layers: int,@@ -195,11 +235,15 @@ def __init__( hidden_act: str = "relu", pre_ln: bool = False, pre_ln_final_layer_norm: bool = True,+ norm: ModuleBlock | None = None,+ self_attn: ModuleBlock | None = None,+ cross_attn: ModuleBlock | None = None,+ feed_forward: ModuleBlock | None = None, ): super().__init__() if pre_ln and pre_ln_final_layer_norm:- self.final_layer_norm = nn.LayerNorm(hidden_size, eps=1e-5)+ self.final_layer_norm = build_module(norm or nn.LayerNorm, normalized_shape=hidden_size) else: self.final_layer_norm = None @@ -214,6 +258,10 @@ def __init__( ffn_dropout, hidden_act, pre_ln,+ norm=norm,+ self_attn=self_attn,+ cross_attn=cross_attn,+ feed_forward=feed_forward, ) self.layers = nn.ModuleList([copy.deepcopy(layer) for _ in range(num_layers)]) self.diagonal = 0nemo/collections/asr/modules/transformer/transformer_modules.py143 + / 21 −
@@ -20,11 +20,38 @@ import numpy as np import torch from torch import nn-from torch.nn.functional import gelu +from nemo.collections.asr.parts.submodules.multi_head_attention import RotaryPositionalEncoding from nemo.collections.common.parts import form_attention_mask+from nemo.collections.common.parts.utils import build_activation -__all__ = ["TransformerEmbedding", "AttentionBridge"]+__all__ = ["TransformerEmbedding", "AttentionBridge", "RMSNorm", "PositionWiseFF", "GatedPositionWiseFF"]+++class RMSNorm(nn.Module):+ """+ Root mean square normalization (Zhang and Sennrich, 2019) with a learned scale and no bias.++ Written with elementary ops so that ONNX export works:+ ``torch.nn.RMSNorm`` lowers to ``aten::rms_norm``, which the exporter does not support.+ The mean square is accumulated in float32, as ``torch.nn.functional.rms_norm`` does for half inputs,+ and only the scale is cast back to the input dtype, which the output keeps.+ The state is ``weight``, as in ``torch.nn.RMSNorm``.++ Args:+ normalized_shape: size of the normalized last axis+ eps: added to the mean square before the square root+ """++ def __init__(self, normalized_shape: int, eps: float = 1e-5):+ super().__init__()+ self.eps = eps+ self.weight = nn.Parameter(torch.ones(normalized_shape))++ def forward(self, x):+ # in float16 the squares overflow from magnitude 256 on, so the mean square is taken in float32+ variance = x.float().pow(2).mean(dim=-1, keepdim=True)+ return x * torch.rsqrt(variance + self.eps).to(x.dtype) * self.weight class FixedPositionalEncoding(nn.Module):@@ -35,13 +62,24 @@ class FixedPositionalEncoding(nn.Module): Args: hidden_size: size of the embeddings in the model, also known as d_model max_sequence_length: maximum allowed length of the input sequence+ layout: ``interleaved`` (default) alternates sine and cosine along the feature axis,+ with the frequencies ``10000**(-2i/hidden_size)``;+ ``tensor2tensor`` stores all sines before all cosines,+ with the frequencies ``10000**(-i/(hidden_size/2 - 1))`` of the original tensor2tensor implementation+ scale: multiplier of the table;+ the default ``1/sqrt(hidden_size)`` matches token embeddings of unit scale,+ ``1.0`` matches token embeddings scaled by ``sqrt(hidden_size)`` """ - def __init__(self, hidden_size, max_sequence_length=512):+ def __init__(self, hidden_size, max_sequence_length=512, layout: str = "interleaved", scale: float | None = None): super().__init__() + if layout not in ("interleaved", "tensor2tensor"):+ raise ValueError(f"layout must be 'interleaved' or 'tensor2tensor', got {layout!r}") self._hidden_size = hidden_size self._max_sequence_length = max_sequence_length+ self._layout = layout+ self._scale = 1.0 / math.sqrt(hidden_size) if scale is None else scale self._build_pos_enc(hidden_size=self._hidden_size, max_sequence_length=self._max_sequence_length) def _build_pos_enc(self, hidden_size, max_sequence_length, device=None):@@ -50,11 +88,17 @@ def _build_pos_enc(self, hidden_size, max_sequence_length, device=None): """ pos_enc = torch.zeros(max_sequence_length, hidden_size, device=device) position = torch.arange(0.0, max_sequence_length).unsqueeze(1)- coef = -math.log(10000.0) / hidden_size- div_term = torch.exp(coef * torch.arange(0.0, hidden_size, 2))- pos_enc[:, 0::2] = torch.sin(position * div_term)- pos_enc[:, 1::2] = torch.cos(position * div_term)- pos_enc.div_(math.sqrt(hidden_size))+ if self._layout == "tensor2tensor":+ half = hidden_size // 2+ div_term = torch.exp(-math.log(10000.0) / max(half - 1, 1) * torch.arange(0.0, half))+ pos_enc[:, :half] = torch.sin(position * div_term)+ pos_enc[:, half:] = torch.cos(position * div_term)+ else:+ coef = -math.log(10000.0) / hidden_size+ div_term = torch.exp(coef * torch.arange(0.0, hidden_size, 2))+ pos_enc[:, 0::2] = torch.sin(position * div_term)+ pos_enc[:, 1::2] = torch.cos(position * div_term)+ pos_enc.mul_(self._scale) self.register_buffer('pos_enc', pos_enc) def forward(self, position_ids):@@ -76,6 +120,10 @@ class TransformerEmbedding(nn.Module): embedding_dropout: probability of dropout applied to embeddings learn_positional_encodings: whether to learn positional encodings or use fixed (sine-cosine) ones+ scale_embeddings: multiply the token embeddings by ``sqrt(hidden_size)`` as in the original Transformer;+ the fixed positional table is then used unscaled+ positional_encoding_layout: layout of the fixed positional table, see ``FixedPositionalEncoding``+ use_layer_norm: apply the LayerNorm after adding the positions (default True) """ def __init__(@@ -86,19 +134,28 @@ def __init__( num_token_types=2, embedding_dropout=0.0, learn_positional_encodings=False,+ scale_embeddings: bool = False,+ positional_encoding_layout: str = "interleaved",+ use_layer_norm: bool = True, ): super().__init__() self.max_sequence_length = max_sequence_length self.learn_positional_encodings = learn_positional_encodings+ self.embedding_scale = math.sqrt(hidden_size) if scale_embeddings else None self.token_embedding = nn.Embedding(vocab_size, hidden_size, padding_idx=0) if learn_positional_encodings: self.position_embedding = nn.Embedding(max_sequence_length, hidden_size) else:- self.position_embedding = FixedPositionalEncoding(hidden_size, max_sequence_length)+ self.position_embedding = FixedPositionalEncoding(+ hidden_size,+ max_sequence_length,+ layout=positional_encoding_layout,+ scale=1.0 if scale_embeddings else None,+ ) if num_token_types > 0: self.token_type_embedding = nn.Embedding(num_token_types, hidden_size)- self.layer_norm = nn.LayerNorm(hidden_size, eps=1e-5)+ self.layer_norm = nn.LayerNorm(hidden_size, eps=1e-5) if use_layer_norm else None self.dropout = nn.Dropout(embedding_dropout) def forward(self, input_ids, token_type_ids=None, start_pos=0):@@ -129,14 +186,17 @@ def forward(self, input_ids, token_type_ids=None, start_pos=0): position_ids = position_ids + shift_pos token_embeddings = self.token_embedding(input_ids)+ if self.embedding_scale is not None:+ token_embeddings = token_embeddings * self.embedding_scale position_embeddings = self.position_embedding(position_ids) embeddings = token_embeddings + position_embeddings if token_type_ids is not None: token_type_embeddings = self.token_type_embedding(token_type_ids) embeddings = embeddings + token_type_embeddings - embeddings = self.layer_norm(embeddings)+ if self.layer_norm is not None:+ embeddings = self.layer_norm(embeddings) embeddings = self.dropout(embeddings) return embeddings@@ -152,6 +212,10 @@ class MultiHeadAttention(nn.Module): attn_score_dropout: probability of dropout applied to attention scores attn_layer_dropout: probability of dropout applied to the output of the whole layer, but before layer normalization+ use_bias: bias in the query, key, value and output projections (default True)+ rotary_base: theta base of a rotary position embedding applied to queries and keys, self-attention only;+ keys longer than the queries are treated as a cache prefix;+ None (default) for no rotation """ def __init__(@@ -161,6 +225,8 @@ def __init__( attn_score_dropout=0.0, attn_layer_dropout=0.0, return_xatt_scores=False,+ use_bias: bool = True,+ rotary_base: float | None = None, ): super().__init__() if hidden_size % num_attention_heads != 0:@@ -173,10 +239,15 @@ def __init__( self.attn_head_size = int(hidden_size / num_attention_heads) self.attn_scale = math.sqrt(math.sqrt(self.attn_head_size)) - self.query_net = nn.Linear(hidden_size, hidden_size)- self.key_net = nn.Linear(hidden_size, hidden_size)- self.value_net = nn.Linear(hidden_size, hidden_size)- self.out_projection = nn.Linear(hidden_size, hidden_size)+ self.query_net = nn.Linear(hidden_size, hidden_size, bias=use_bias)+ self.key_net = nn.Linear(hidden_size, hidden_size, bias=use_bias)+ self.value_net = nn.Linear(hidden_size, hidden_size, bias=use_bias)+ self.out_projection = nn.Linear(hidden_size, hidden_size, bias=use_bias)+ self.rotary_pos_enc = (+ RotaryPositionalEncoding(d_k=self.attn_head_size, rope_base=rotary_base)+ if rotary_base is not None+ else None+ ) self.attn_dropout = nn.Dropout(attn_score_dropout) self.layer_dropout = nn.Dropout(attn_layer_dropout)@@ -198,6 +269,10 @@ def forward(self, queries, keys, values, attention_mask): query = self.transpose_for_scores(query) / self.attn_scale key = self.transpose_for_scores(key) / self.attn_scale value = self.transpose_for_scores(value)+ if self.rotary_pos_enc is not None:+ # the rotation commutes with the scalar scaling above; the table grows with the key length+ self.rotary_pos_enc.extend_pe(key.size(2), key.device, key.dtype)+ query, key = self.rotary_pos_enc(query, key) # for numerical stability we pre-divide query and key by sqrt(sqrt(d)) attention_scores = torch.matmul(query, key.transpose(-1, -2))@@ -232,16 +307,19 @@ class PositionWiseFF(nn.Module): inner_size: number of neurons in the intermediate part of feed-forward net, usually is (4-8 x hidden_size) in the papers ffn_dropout: probability of dropout applied to net output- hidden_act: activation function used between two linear layers+ hidden_act: activation between the two linear layers, a module or its name in the activation registry+ (``relu``, ``gelu``, ``swish``, ``squared_relu``, ...)+ use_bias: bias in both projections (default True) """ - def __init__(self, hidden_size, inner_size, ffn_dropout=0.0, hidden_act="relu"):+ def __init__(+ self, hidden_size, inner_size, ffn_dropout=0.0, hidden_act: nn.Module | str = "relu", use_bias: bool = True+ ): super().__init__()- self.dense_in = nn.Linear(hidden_size, inner_size)- self.dense_out = nn.Linear(inner_size, hidden_size)+ self.dense_in = nn.Linear(hidden_size, inner_size, bias=use_bias)+ self.dense_out = nn.Linear(inner_size, hidden_size, bias=use_bias) self.layer_dropout = nn.Dropout(ffn_dropout)- ACT2FN = {"gelu": gelu, "relu": torch.relu}- self.act_fn = ACT2FN[hidden_act]+ self.act_fn = build_activation(hidden_act) def forward(self, hidden_states): output_states = self.dense_in(hidden_states)@@ -251,6 +329,50 @@ def forward(self, hidden_states): return output_states +class GatedPositionWiseFF(nn.Module):+ """+ Gated position-wise feed-forward network (the GLU variants of Shazeer, 2020), SwiGLU with the defaults.++ The first projection has twice the inner size; its first half goes through ``hidden_act``+ and gates its second half, which goes through ``gate_act``:+ ``dense_out(hidden_act(a) * gate_act(b))`` with ``a, b = dense_in(x).chunk(2)``.+ A decoder block passes its own ``hidden_act`` on to this block, so a block section names ``hidden_act: swish``+ for SwiGLU unless the decoder's activation is already the one wanted.++ Args:+ hidden_size: size of the embeddings in the model, also known as d_model+ inner_size: size of the gated intermediate representation;+ the paper keeps the parameter count of the plain network with two thirds of its inner size+ ffn_dropout: probability of dropout applied to net output+ hidden_act: activation of the first half, a module or its name in the activation registry; default ``swish``+ gate_act: activation of the second half; default ``identity``+ use_bias: bias in both projections; default False, as the gated variants are usually trained+ """++ def __init__(+ self,+ hidden_size,+ inner_size,+ ffn_dropout=0.0,+ hidden_act: nn.Module | str = "swish",+ gate_act: nn.Module | str = "identity",+ use_bias: bool = False,+ ):+ super().__init__()+ self.dense_in = nn.Linear(hidden_size, 2 * inner_size, bias=use_bias)+ self.dense_out = nn.Linear(inner_size, hidden_size, bias=use_bias)+ self.layer_dropout = nn.Dropout(ffn_dropout)+ self.act_fn = build_activation(hidden_act)+ self.gate_act_fn = build_activation(gate_act)++ def forward(self, hidden_states):+ activated, gate = self.dense_in(hidden_states).chunk(2, dim=-1)+ output_states = self.act_fn(activated) * self.gate_act_fn(gate)+ output_states = self.dense_out(output_states)+ output_states = self.layer_dropout(output_states)+ return output_states++ class AttentionBridge(torch.nn.Module): """ A multi-head attention bridge to project a variable-size hidden statesnemo/collections/asr/modules/transformer/transformer_utils.py5 + / 0 −
@@ -131,6 +131,11 @@ def get_nemo_transformer( pre_ln=cfg.get('pre_ln', False), pre_ln_final_layer_norm=pre_ln_final_layer_norm, num_token_types=cfg.get('num_token_types', 2),+ embedding=cfg.get('embedding', None),+ norm=cfg.get('norm', None),+ self_attn=cfg.get('self_attn', None),+ cross_attn=cfg.get('cross_attn', None),+ feed_forward=cfg.get('feed_forward', None), ) return modelnemo/collections/common/parts/utils.py27 + / 1 −
@@ -29,7 +29,16 @@ import torch.nn as nn from omegaconf import DictConfig -__all__ = ['if_exist', '_compute_softmax', 'flatten', 'build_module', 'block_binds', 'ModuleBlock', 'SquaredReLU']+__all__ = [+ 'if_exist',+ '_compute_softmax',+ 'flatten',+ 'build_module',+ 'block_binds',+ 'build_activation',+ 'ModuleBlock',+ 'SquaredReLU',+] class SquaredReLU(nn.Module):@@ -50,6 +59,23 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: "squared_relu": SquaredReLU, } ++def build_activation(activation: nn.Module | str) -> nn.Module:+ """The activation module given, or a new instance of the registry's class under this name.++ Args:+ activation: an activation module, or its name in ``activation_registry`` (``relu``, ``gelu``, ``swish``, ...).++ Returns:+ The activation module.+ """+ if isinstance(activation, str):+ if activation not in activation_registry:+ raise ValueError(f"unknown activation {activation!r}; the registry has {sorted(activation_registry)}")+ return activation_registry[activation]()+ return activation++ ModuleBlock = Callable[..., nn.Module] | functools.partial | Mapping tests/collections/asr/test_transformer_decoder.pyadded248 + / 0 −
@@ -0,0 +1,248 @@+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.+# SPDX-License-Identifier: Apache-2.0+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+# http://www.apache.org/licenses/LICENSE-2.0+#+# Unless required by applicable law or agreed to in writing, software+# distributed under the License is distributed on an "AS IS" BASIS,+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.+# See the License for the specific language governing permissions and+# limitations under the License.++"""Transformer decoder modules: configurable blocks (embedding, norm, self- and cross-attention, feed-forward)."""++import functools+import math++import pytest+import torch+from omegaconf import DictConfig++from nemo.collections.asr.modules.transformer.transformer import TransformerDecoderNM+from nemo.collections.asr.modules.transformer.transformer_modules import (+ FixedPositionalEncoding,+ GatedPositionWiseFF,+ MultiHeadAttention,+ PositionWiseFF,+ RMSNorm,+ TransformerEmbedding,+)+from nemo.collections.asr.modules.transformer.transformer_utils import get_nemo_transformer+++def _composed_decoder(**overrides):+ options = dict(+ vocab_size=11,+ hidden_size=16,+ num_layers=2,+ inner_size=24,+ num_attention_heads=2,+ max_sequence_length=32,+ num_token_types=0,+ pre_ln=True,+ embedding=functools.partial(+ TransformerEmbedding,+ scale_embeddings=True,+ positional_encoding_layout="tensor2tensor",+ use_layer_norm=False,+ ),+ norm=functools.partial(RMSNorm, eps=1e-6),+ self_attn=functools.partial(MultiHeadAttention, use_bias=False, rotary_base=10000.0),+ feed_forward=functools.partial(GatedPositionWiseFF, hidden_act="swish"),+ )+ options.update(overrides)+ return TransformerDecoderNM(**options).eval()+++class TestTransformerDecoderBlockComposition:+ @pytest.mark.unit+ def test_blocks_from_partials(self):+ decoder = _composed_decoder()+ block = decoder.decoder.layers[0]+ assert isinstance(block.layer_norm_1, RMSNorm) and block.layer_norm_1.eps == 1e-6+ assert isinstance(decoder.decoder.final_layer_norm, RMSNorm)+ assert block.first_sub_layer.query_net.bias is None and block.first_sub_layer.rotary_pos_enc is not None+ # the cross-attention keeps the default block+ assert block.second_sub_layer.key_net.bias is not None and block.second_sub_layer.rotary_pos_enc is None+ assert block.third_sub_layer.dense_in.weight.shape == (48, 16) and block.third_sub_layer.dense_in.bias is None+ assert decoder.embedding.layer_norm is None and decoder.embedding.embedding_scale == math.sqrt(16)+ assert decoder.embedding.token_embedding.num_embeddings == 11++ @pytest.mark.unit+ @pytest.mark.parametrize("rotary", [True, False])+ def test_incremental_decoding_matches_full_pass(self, rotary):+ torch.manual_seed(0)+ overrides = {} if rotary else dict(self_attn=functools.partial(MultiHeadAttention, use_bias=False))+ decoder = _composed_decoder(**overrides)+ for parameter in decoder.parameters():+ parameter.data.normal_(0, 0.3)+ input_ids = torch.tensor([[10, 1, 5, 7, 2], [10, 3, 4, 0, 0]])+ encoder_states = torch.randn(2, 9, 16)+ encoder_mask = torch.ones(2, 9)+ with torch.no_grad():+ full, _ = decoder.decoder(+ decoder.embedding(input_ids), torch.ones_like(input_ids), encoder_states, encoder_mask+ )+ mems = None+ for step in range(input_ids.shape[1]):+ embedded = decoder.embedding(input_ids[:, step : step + 1], start_pos=step)+ mems, _ = decoder.decoder(+ embedded, torch.ones(2, 1), encoder_states, encoder_mask, mems, return_mems=True+ )+ torch.testing.assert_close(mems[-1], full, atol=1e-5, rtol=1e-5)++ @pytest.mark.unit+ def test_blocks_from_hydra_sections(self):+ config = DictConfig(+ {+ "vocab_size": 11,+ "hidden_size": 16,+ "num_layers": 1,+ "inner_size": 24,+ "num_attention_heads": 2,+ "max_sequence_length": 32,+ "num_token_types": 0,+ "pre_ln": True,+ "norm": {+ "_target_": "nemo.collections.asr.modules.transformer.transformer_modules.RMSNorm",+ "_partial_": True,+ "eps": 1e-6,+ },+ "self_attn": {+ "_target_": "nemo.collections.asr.modules.transformer.transformer_modules.MultiHeadAttention",+ "_partial_": True,+ "use_bias": False,+ "rotary_base": 10000.0,+ },+ "feed_forward": {+ "_target_": "nemo.collections.asr.modules.transformer.transformer_modules.GatedPositionWiseFF",+ "_partial_": True,+ "hidden_act": "swish",+ },+ }+ )+ decoder = get_nemo_transformer(config_dict=config, encoder=False)+ block = decoder.decoder.layers[0]+ assert isinstance(block.layer_norm_1, RMSNorm) and block.layer_norm_1.eps == 1e-6+ assert block.first_sub_layer.query_net.bias is None and block.first_sub_layer.rotary_pos_enc is not None+ assert type(block.third_sub_layer) is GatedPositionWiseFF+ assert block.third_sub_layer.dense_in.weight.shape == (48, 16) and block.third_sub_layer.dense_in.bias is None+ assert isinstance(block.third_sub_layer.act_fn, torch.nn.SiLU)+ assert isinstance(decoder.embedding.layer_norm, torch.nn.LayerNorm)++ @pytest.mark.unit+ def test_rms_norm(self):+ torch.manual_seed(0)+ norm = RMSNorm(8, eps=1e-6)+ norm.weight.data.normal_()+ x = torch.randn(3, 8)+ torch.testing.assert_close(norm(x), torch.nn.functional.rms_norm(x, (8,), norm.weight, eps=1e-6))+ assert list(norm.state_dict()) == ["weight"]++ @pytest.mark.unit+ @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])+ def test_rms_norm_half_precision(self, dtype):+ # the mean square must not overflow: 1000 squared is outside the float16 range+ norm = RMSNorm(8, eps=1e-6).to(dtype)+ large = torch.full((2, 8), 1000.0, dtype=dtype)+ torch.testing.assert_close(norm(large), torch.ones_like(large))+ # forward and backward against the float32 computation; mixed signs, a zero row, a non-unit weight+ torch.manual_seed(0)+ x32 = torch.randn(5, 8) * 300.0+ x32[2] = 0.0+ w32 = torch.randn(8)+ norm.weight.data = w32.to(dtype)+ x = x32.to(dtype).requires_grad_()+ out = norm(x)+ assert out.dtype == dtype+ out.sum().backward()+ reference = RMSNorm(8, eps=1e-6)+ reference.weight.data = w32.clone()+ x_ref = x32.clone().requires_grad_()+ expected = reference(x_ref)+ expected.sum().backward()+ torch.testing.assert_close(out.float(), expected, rtol=2e-2, atol=5e-2)+ torch.testing.assert_close(x.grad.float(), x_ref.grad, rtol=2e-2, atol=5e-2)+ torch.testing.assert_close(norm.weight.grad.float(), reference.weight.grad, rtol=2e-2, atol=5e-2)++ @pytest.mark.unit+ def test_export_covers_rotary_positions_beyond_the_example(self, tmp_path):+ # the exported graph holds the rotary table as a constant,+ # so it must cover every position the decoder supports, not only the example's key length+ decoder = _composed_decoder(max_sequence_length=8)+ for parameter in decoder.parameters():+ parameter.data.normal_(0, 0.3)++ def example(length):+ torch.manual_seed(length)+ input_ids = torch.randint(0, 11, (1, length))+ decoder_mask = torch.ones(1, length, dtype=torch.int64)+ encoder_states = torch.randn(1, 9, 16)+ encoder_mask = torch.ones(1, 9, dtype=torch.int64)+ decoder_mems = torch.randn(1, decoder.num_states, length - 1, 16)+ return input_ids, decoder_mask, encoder_states, encoder_mask, decoder_mems++ path = str(tmp_path / "decoder.ts")+ decoder.export(path, input_example=example(4), check_trace=False)+ exported = torch.jit.load(path)+ for length in (4, 5, 8):+ input_ids, decoder_mask, encoder_states, encoder_mask, decoder_mems = example(length)+ with torch.no_grad():+ # export switched the module to returning the states of every layer, as the graph does+ expected = decoder(+ input_ids=input_ids,+ decoder_mask=decoder_mask,+ encoder_embeddings=encoder_states,+ encoder_mask=encoder_mask,+ decoder_mems=decoder_mems,+ )+ torch.testing.assert_close(+ exported(input_ids, decoder_mask, encoder_states, encoder_mask, decoder_mems), expected+ )++ @pytest.mark.unit+ def test_gated_feed_forward(self):+ torch.manual_seed(0)+ ff = GatedPositionWiseFF(16, 24)+ assert ff.dense_in.bias is None and ff.dense_in.weight.shape == (48, 16)+ x = torch.randn(3, 16)+ activated, gate = ff.dense_in(x).chunk(2, dim=-1)+ torch.testing.assert_close(ff(x), ff.dense_out(torch.nn.functional.silu(activated) * gate))+ # both activations are free: GEGLU-style first half with a gated second half, from registry names+ geglu = GatedPositionWiseFF(16, 24, hidden_act="gelu", gate_act="relu", use_bias=True)+ assert isinstance(geglu.act_fn, torch.nn.GELU) and isinstance(geglu.gate_act_fn, torch.nn.ReLU)+ assert geglu.dense_in.bias is not None+ # the plain network takes its activation from the registry as well+ assert isinstance(PositionWiseFF(16, 24, hidden_act="gelu").act_fn, torch.nn.GELU)+ with pytest.raises(ValueError, match="activation"):+ PositionWiseFF(16, 24, hidden_act="tanh")++ @pytest.mark.unit+ def test_tensor2tensor_positional_table(self):+ table = FixedPositionalEncoding(8, max_sequence_length=4, layout="tensor2tensor", scale=1.0).pos_enc+ positions = torch.arange(4.0)[:, None]+ frequencies = torch.exp(-math.log(10000.0) / 3 * torch.arange(4.0))+ torch.testing.assert_close(table[:, :4], torch.sin(positions * frequencies))+ torch.testing.assert_close(table[:, 4:], torch.cos(positions * frequencies))+ default = FixedPositionalEncoding(8, max_sequence_length=4).pos_enc+ torch.testing.assert_close(default[1, 0], torch.sin(torch.tensor(1.0)) / math.sqrt(8))++ @pytest.mark.unit+ def test_defaults_unchanged(self):+ decoder = TransformerDecoderNM(+ vocab_size=8, hidden_size=16, num_layers=1, inner_size=24, num_attention_heads=2+ )+ block = decoder.decoder.layers[0]+ assert type(block.layer_norm_1) is torch.nn.LayerNorm and block.layer_norm_1.eps == 1e-5+ assert block.first_sub_layer.query_net.bias is not None and block.first_sub_layer.rotary_pos_enc is None+ assert type(block.third_sub_layer) is PositionWiseFF+ assert isinstance(block.third_sub_layer.act_fn, torch.nn.ReLU)+ assert block.third_sub_layer.dense_in.weight.shape == (24, 16)+ assert (+ isinstance(decoder.embedding.layer_norm, torch.nn.LayerNorm) and decoder.embedding.embedding_scale is None+ )+ assert decoder.decoder.final_layer_norm is Nonetests/collections/common/test_utils.py9 + / 0 −
@@ -29,6 +29,7 @@ SquaredReLU, activation_registry, block_binds,+ build_activation, build_module, flatten, mask_sequence_tensor,@@ -449,3 +450,11 @@ def test_squared_relu(self): x = torch.tensor([-2.0, 0.0, 3.0]) torch.testing.assert_close(SquaredReLU()(x), torch.tensor([0.0, 0.0, 9.0])) assert activation_registry["squared_relu"] is SquaredReLU++ @pytest.mark.unit+ def test_build_activation(self):+ assert isinstance(build_activation("gelu"), torch.nn.GELU)+ module = torch.nn.Tanh()+ assert build_activation(module) is module+ with pytest.raises(ValueError, match="unknown activation 'tanh'"):+ build_activation("tanh")