wehub-resource-sync
ba4be087d5
Create PR to main with cherry-pick from release / cherry-pick (push) Failing after 0s
CICD NeMo / pre-flight (push) Failing after 0s
CICD NeMo / configure (push) Has been skipped
Build, validate, and release Neural Modules / pre-flight (push) Failing after 1s
CICD NeMo / code-linting (push) Has been skipped
Build, validate, and release Neural Modules / release (push) Has been skipped
Build, validate, and release Neural Modules / release-summary (push) Has been cancelled
CICD NeMo / cicd-test-container-build (push) Has been cancelled
CICD NeMo / cicd-import-tests (push) Has been cancelled
CICD NeMo / L0_Setup_Test_Data_And_Models (push) Has been cancelled
CICD NeMo / cicd-main-unit-tests (push) Has been cancelled
CICD NeMo / cicd-main-speech (push) Has been cancelled
CICD NeMo / Nemo_CICD_Test (push) Has been cancelled
CICD NeMo / Coverage (e2e) (push) Has been cancelled
CICD NeMo / Coverage (unit-test) (push) Has been cancelled
CodeQL / Analyze (python) (push) Has been cancelled
CICD NeMo / cicd-wait-in-queue (push) Has been cancelled
622 lines
30 KiB
Python
622 lines
30 KiB
Python
# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
|
|
#
|
|
# 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.
|
|
|
|
import math
|
|
from dataclasses import dataclass
|
|
from typing import Optional
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
from torch.nn.attention.flex_attention import and_masks, create_block_mask, flex_attention
|
|
|
|
from nemo.collections.asr.parts.submodules.multi_head_attention import (
|
|
PositionalEncoding,
|
|
RelPositionalEncoding,
|
|
RelPositionMultiHeadAttention,
|
|
RotaryPositionalEncoding,
|
|
)
|
|
from nemo.collections.asr.parts.submodules.subsampling import FeatureStacking, StackingSubsampling
|
|
from nemo.core.classes.module import freeze, unfreeze
|
|
from nemo.utils.decorators import experimental
|
|
|
|
flex_attention_compiled = torch.compile(flex_attention, dynamic=True)
|
|
|
|
|
|
@dataclass
|
|
class TransformerEncoderConfig:
|
|
"""Configuration for ``TransformerEncoder`` and its sub-blocks.
|
|
|
|
Args:
|
|
feat_in: Input feature dimension (e.g. number of mel bins).
|
|
d_model: Transformer encoder state dimension, i.e. the size of the residual stream that flows
|
|
through every block (token/frame embedding size, attention input/output size, and
|
|
feed-forward input/output size). Also known as ``hidden_size`` in HuggingFace
|
|
``transformers`` configs and ``embed_dim``/``d_model`` in PyTorch's
|
|
``nn.TransformerEncoderLayer``.
|
|
n_heads: Number of attention heads.
|
|
n_layers: Number of Transformer blocks.
|
|
drop_rate: Dropout probability applied inside attention and feed-forward sublayers.
|
|
qkv_bias: If True, add a learnable bias to the fused Q/K/V projection. Many modern ASR/LM
|
|
Transformers (e.g. HuggingFace Whisper) drop the bias on the K projection because a
|
|
constant K bias adds the same scalar to every key and is wiped out by softmax's
|
|
shift-invariance, making it a redundant parameter. Default ``False`` matches that style.
|
|
qk_norm: If True, apply per-head ``LayerNorm`` to Q and K before the dot product. Stabilizes
|
|
training by preventing exponential Q/K-norm growth and "attention entropy collapse"
|
|
(Henry et al. 2020; used in OLMo 2, Gemma 3, Qwen 3). Cheap, ~no-op for inference.
|
|
ff_expansion: Multiplier for the per-block FFN inner hidden size:
|
|
``ffn_hidden_size = int(ff_expansion * d_model)``. Only widens the intermediate FFN
|
|
projection; FFN input/output stays at ``d_model``. Typical value ``4.0``; ``float``
|
|
allows sub-1x experts for MoE. Equivalent to ``intermediate_size / hidden_size`` in
|
|
HuggingFace and ``dim_feedforward / d_model`` in PyTorch's ``nn.TransformerEncoderLayer``.
|
|
pre_block_norm: If True, apply ``LayerNorm`` to embeddings before the first Transformer block
|
|
(BERT/ViT-style). Set False to match pre-norm Transformers such as Whisper or GPT-2.
|
|
subsampling_factor: Frame-level subsampling factor performed by the pre-encoder.
|
|
attn_mode: Attention pattern. Currently only ``"full"`` (bidirectional) is supported.
|
|
Future modes: ``"causal"``, ``"lookahead"``, ``"local"``, ``"sliding_window"``.
|
|
self_attention_model: Positional encoding / attention scoring scheme.
|
|
|
|
- ``"rel_pos"`` (default): Transformer-XL relative positional encoding
|
|
(https://arxiv.org/abs/1901.02860). The (b)+(d) cross/positional bias is computed
|
|
from the relative-position embedding and injected into FlexAttention via a
|
|
``score_mod`` closure; the (c) global-content bias is folded into the query as
|
|
``Q + pos_bias_u``.
|
|
- ``"abs_pos"``: sinusoidal absolute positional encoding added to embeddings
|
|
before the first block; standard scaled dot-product attention.
|
|
- ``"no_pos"`` (or ``None``): no positional encoding at all. The pre-encoder output
|
|
is consumed directly by the Transformer blocks. ``xscaling``, ``pos_emb_max_len``,
|
|
``dropout_pre_encoder`` and ``dropout_emb`` are unused in this mode.
|
|
- ``"rope"``: rotary position embedding applied to Q and K inside attention. No
|
|
additive positional embedding is added to the embeddings; the standard scaled
|
|
dot-product attention runs through FlexAttention with ``score_mod=None``.
|
|
rope_base: Theta base for the rotary position embedding. Only used when
|
|
``self_attention_model='rope'``.
|
|
rotary_fraction: Fraction of the per-head dim to rotate. Only used when
|
|
``self_attention_model='rope'``.
|
|
"""
|
|
|
|
feat_in: int = 128
|
|
d_model: int = 512
|
|
n_heads: int = 8
|
|
n_layers: int = 17
|
|
drop_rate: float = 0.1
|
|
qkv_bias: bool = False
|
|
qk_norm: bool = False
|
|
ff_expansion: float = 4.0
|
|
pre_block_norm: bool = True
|
|
subsampling_factor: int = 4
|
|
# Attention mode: "full" (bidirectional) or "causal" (each token only attends to itself and earlier tokens).
|
|
# Future: "lookahead", "local", "sliding_window".
|
|
attn_mode: str = "full"
|
|
self_attention_model: str = "rel_pos"
|
|
rope_base: float = 10000.0
|
|
rotary_fraction: float = 1.0
|
|
|
|
|
|
def _make_padding_mod(lengths):
|
|
"""Mask out padding positions based on per-sample lengths."""
|
|
|
|
def pad_mask(b, h, q_idx, kv_idx):
|
|
return kv_idx < lengths[b]
|
|
|
|
return pad_mask
|
|
|
|
|
|
def _make_causal_mod():
|
|
"""Strictly causal — each query only attends to its own and earlier kv positions."""
|
|
|
|
def causal(b, h, q_idx, kv_idx):
|
|
return q_idx >= kv_idx
|
|
|
|
return causal
|
|
|
|
|
|
_SUPPORTED_ATTENTION_MODES = ("full", "causal")
|
|
_SUPPORTED_SELF_ATTENTION_MODELS = ("abs_pos", "rel_pos", "no_pos", "rope")
|
|
|
|
|
|
class FeedForward(nn.Module):
|
|
def __init__(self, cfg: TransformerEncoderConfig):
|
|
super().__init__()
|
|
ff_hidden = int(cfg.ff_expansion * cfg.d_model)
|
|
self.net = nn.Sequential(
|
|
nn.Linear(cfg.d_model, ff_hidden),
|
|
nn.GELU(),
|
|
nn.Dropout(cfg.drop_rate),
|
|
nn.Linear(ff_hidden, cfg.d_model),
|
|
nn.Dropout(cfg.drop_rate),
|
|
)
|
|
|
|
def forward(self, x):
|
|
return self.net(x)
|
|
|
|
|
|
class MultiHeadAttention(nn.Module):
|
|
def __init__(self, cfg: TransformerEncoderConfig, pos_enc=None):
|
|
super().__init__()
|
|
self.n_heads = cfg.n_heads
|
|
self.head_dim = cfg.d_model // cfg.n_heads
|
|
self.d_model = cfg.d_model
|
|
self.self_attention_model = cfg.self_attention_model
|
|
self._uses_rel_pos = self.self_attention_model == "rel_pos"
|
|
self._uses_rope = self.self_attention_model == "rope"
|
|
if self.self_attention_model not in _SUPPORTED_SELF_ATTENTION_MODELS:
|
|
raise ValueError(
|
|
f"self_attention_model='{self.self_attention_model}' is not supported. "
|
|
f"Supported modes: {_SUPPORTED_SELF_ATTENTION_MODELS}."
|
|
)
|
|
if self.head_dim < 16:
|
|
raise ValueError(
|
|
"PyTorch FlexAttention CUDA backend requires per-head embedding dimension >= 16, "
|
|
f"but got head_dim={self.head_dim} from d_model={self.d_model}, n_heads={self.n_heads}."
|
|
)
|
|
|
|
# Rotary position embedding shared across layers; rotates Q/K before the attention
|
|
# kernel. The shared module owns the cos/sin buffers (see ``TransformerEncoder``).
|
|
if self._uses_rope:
|
|
if pos_enc is None:
|
|
raise ValueError("'rope' attention requires a RotaryPositionalEncoding via pos_enc.")
|
|
self.rope = pos_enc
|
|
else:
|
|
self.rope = None
|
|
|
|
self.w_qkv = nn.Linear(cfg.d_model, 3 * cfg.d_model, bias=cfg.qkv_bias)
|
|
self.out_proj = nn.Linear(cfg.d_model, cfg.d_model)
|
|
|
|
self.qk_norm = cfg.qk_norm
|
|
if cfg.qk_norm:
|
|
self.q_norm = nn.LayerNorm(self.head_dim)
|
|
self.k_norm = nn.LayerNorm(self.head_dim)
|
|
|
|
# Transformer-XL relative-position parameters (matrix b and matrix d from
|
|
# https://arxiv.org/abs/1901.02860 Section 3.3). The "matrix c" term `u @ K^T` is
|
|
# absorbed by passing `Q + pos_bias_u` as the query to FlexAttention.
|
|
if self._uses_rel_pos:
|
|
self.linear_pos = nn.Linear(cfg.d_model, cfg.d_model, bias=False)
|
|
self.pos_bias_u = nn.Parameter(torch.zeros(self.n_heads, self.head_dim))
|
|
self.pos_bias_v = nn.Parameter(torch.zeros(self.n_heads, self.head_dim))
|
|
else:
|
|
self.linear_pos = None
|
|
self.pos_bias_u = None
|
|
self.pos_bias_v = None
|
|
|
|
def _rel_shift(self, x):
|
|
"""Transformer-XL relative-position shift.
|
|
|
|
Delegates to ``RelPositionMultiHeadAttention.rel_shift`` (which does not reference
|
|
``self``) so the logic lives in a single place — NeMo's existing reference
|
|
implementation in ``parts/submodules/multi_head_attention.py``.
|
|
"""
|
|
return RelPositionMultiHeadAttention.rel_shift(None, x)
|
|
|
|
def _build_rel_pos_score_mod(self, q, pos_emb):
|
|
"""Build the FlexAttention inputs that realize Transformer-XL relative attention.
|
|
|
|
Implements the (b), (c), (d) terms of Transformer-XL Section 3.3
|
|
(https://arxiv.org/abs/1901.02860) on top of FlexAttention:
|
|
|
|
- Matrices (b) + (d) — the position-dependent score bias ``(Q + v) @ R^T`` rel-
|
|
shifted into ``(q_idx, kv_idx)`` coordinates — are precomputed into a
|
|
``(B, H, T, T)`` tensor, scaled by ``1/sqrt(D)`` (to match FlexAttention's
|
|
already-scaled ``QK^T`` scores), and captured by a local ``score_mod`` closure.
|
|
FlexAttention's ``score_mod`` API fixes the callable signature, so the per-forward
|
|
tensor is threaded in via closure capture rather than as an explicit argument.
|
|
Keeping it local (instead of on ``self``) lets the ``(B, H, T, T)`` bias be freed
|
|
as soon as the layer's attention call returns, so peak memory holds at most one
|
|
layer's bias rather than all layers' biases at once.
|
|
- Matrix (c) — the global-content bias ``u @ K^T`` — is folded into FlexAttention
|
|
by rewriting the query as ``Q + pos_bias_u``, which is returned.
|
|
|
|
Args:
|
|
q: Query tensor with shape ``(B, H, T, D)``.
|
|
pos_emb: Relative positional embedding ``(1, 2T - 1, d_model)`` produced by
|
|
``RelPositionalEncoding``.
|
|
|
|
Returns:
|
|
score_mod: Callable to pass as ``flex_attention(..., score_mod=...)``.
|
|
q_with_bias_u: ``Q + pos_bias_u`` — the (c) "matrix c" query rewrite.
|
|
"""
|
|
H, D = self.n_heads, self.head_dim
|
|
T = q.size(-2)
|
|
# pos_emb: (1, 2T - 1, d_model) -> p: (1, H, 2T - 1, D)
|
|
p = self.linear_pos(pos_emb).view(pos_emb.size(0), -1, H, D).transpose(1, 2)
|
|
# pos_bias_{u,v}: (H, D) -> (1, H, 1, D) so they broadcast over the (B, H, T, D)
|
|
# Q tensor against the head/depth axes rather than (incorrectly) against time.
|
|
bias_u = self.pos_bias_u.view(1, H, 1, D).to(q.dtype)
|
|
bias_v = self.pos_bias_v.view(1, H, 1, D).to(q.dtype)
|
|
# Matrix b + d: ((Q + v) @ R^T) shifted into (q_idx, kv_idx) space, then scaled
|
|
# by 1/sqrt(D) so it can be added directly to FlexAttention's already-scaled scores.
|
|
q_with_bias_v = q + bias_v
|
|
matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1)) # (B, H, T, 2T - 1)
|
|
# rel_shift converts absolute-relative-position columns into (query, key) columns;
|
|
# keep the first T to land in (B, H, T, T) bias space.
|
|
rel_pos_bias = self._rel_shift(matrix_bd)[..., :T] * (D**-0.5)
|
|
|
|
def score_mod(score, b, h, q_idx, kv_idx):
|
|
return score + rel_pos_bias[b, h, q_idx, kv_idx]
|
|
|
|
# Matrix c: fold u @ K^T into FlexAttention by rewriting Q as (Q + u).
|
|
return score_mod, q + bias_u
|
|
|
|
def forward(self, x, block_mask=None, pos_emb=None):
|
|
B, T, _ = x.shape
|
|
H, D = self.n_heads, self.head_dim
|
|
|
|
qkv = self.w_qkv(x).view(B, T, 3, H, D).permute(2, 0, 3, 1, 4)
|
|
q, k, v = qkv.unbind(0)
|
|
|
|
if self.qk_norm:
|
|
q = self.q_norm(q).to(v.dtype)
|
|
k = self.k_norm(k).to(v.dtype)
|
|
|
|
if self._uses_rope:
|
|
# RoPE rotates Q/K in place; it is orthogonal to FlexAttention's score_mod.
|
|
q, k = self.rope(q, k)
|
|
|
|
score_mod = None
|
|
if self._uses_rel_pos:
|
|
score_mod, q = self._build_rel_pos_score_mod(q, pos_emb)
|
|
|
|
attn_fn = flex_attention_compiled if q.is_cuda else flex_attention
|
|
out = attn_fn(q, k, v, block_mask=block_mask, score_mod=score_mod)
|
|
out = out.transpose(1, 2).contiguous().view(B, T, self.d_model)
|
|
return self.out_proj(out)
|
|
|
|
|
|
class TransformerBlock(nn.Module):
|
|
def __init__(self, cfg: TransformerEncoderConfig, pos_enc=None):
|
|
super().__init__()
|
|
self.norm1 = nn.LayerNorm(cfg.d_model)
|
|
self.attn = MultiHeadAttention(cfg, pos_enc=pos_enc)
|
|
self.drop = nn.Dropout(cfg.drop_rate)
|
|
self.norm2 = nn.LayerNorm(cfg.d_model)
|
|
self.ffn = FeedForward(cfg)
|
|
|
|
def forward(self, x, block_mask=None, pos_emb=None):
|
|
x = x + self.drop(self.attn(self.norm1(x), block_mask=block_mask, pos_emb=pos_emb))
|
|
x = x + self.drop(self.ffn(self.norm2(x)))
|
|
return x
|
|
|
|
|
|
@experimental
|
|
class TransformerEncoder(nn.Module):
|
|
"""Pre-norm Transformer encoder for ASR.
|
|
|
|
Architecture: PreEncode -> PositionalEncoding -> LayerNorm -> N x TransformerBlock -> FinalNorm
|
|
|
|
Uses PyTorch FlexAttention for attention computation. On CUDA, mask functions
|
|
are compiled into fused Triton kernels with block-sparse optimization. On CPU,
|
|
FlexAttention falls back to an unfused implementation automatically.
|
|
|
|
Args:
|
|
feat_in: Input feature dimension (number of mel bins).
|
|
d_model: Transformer encoder state dimension, i.e. the size of the residual stream that flows
|
|
through every block (token/frame embedding size, attention input/output size, and
|
|
feed-forward input/output size). Also known as ``hidden_size`` in HuggingFace
|
|
``transformers`` configs and ``embed_dim``/``d_model`` in PyTorch's
|
|
``nn.TransformerEncoderLayer``.
|
|
|
|
n_heads: Number of attention heads.
|
|
n_layers: Number of Transformer blocks.
|
|
feat_out: Output feature dimension. Defaults to ``d_model``.
|
|
subsampling: Subsampling method. Supports ``feature_stacking`` for the
|
|
Transformer-native ``FeatureStacking`` module, plus ``stacking`` and
|
|
``stacking_norm`` for linear frame stacking.
|
|
subsampling_factor: Subsampling factor for the pre-encoder.
|
|
drop_rate: Dropout probability.
|
|
dropout_pre_encoder: Dropout probability after positional encoding. Defaults to ``drop_rate``.
|
|
dropout_emb: Dropout probability for positional embeddings.
|
|
qkv_bias: If True, add a learnable bias to the fused Q/K/V projection. Many modern ASR/LM
|
|
Transformers (e.g. HuggingFace Whisper) drop the bias on the K projection because a
|
|
constant K bias adds the same scalar to every key and is wiped out by softmax's
|
|
shift-invariance, making it a redundant parameter. Default ``False`` matches that style.
|
|
qk_norm: If True, apply per-head ``LayerNorm`` to Q and K before the dot product. Stabilizes
|
|
training by preventing exponential Q/K-norm growth and "attention entropy collapse"
|
|
(Henry et al. 2020; used in OLMo 2, Gemma 3, Qwen 3). Cheap, ~no-op for inference.
|
|
ff_expansion: Multiplier for the per-block FFN inner hidden size:
|
|
``ffn_hidden_size = int(ff_expansion * d_model)``. Only widens the intermediate FFN
|
|
projection; FFN input/output stays at ``d_model``. Typical value ``4.0``; ``float``
|
|
allows sub-1x experts for MoE. Equivalent to ``intermediate_size / hidden_size`` in
|
|
HuggingFace and ``dim_feedforward / d_model`` in PyTorch's ``nn.TransformerEncoderLayer``.
|
|
pre_block_norm: If True (default), apply LayerNorm to embeddings before the first
|
|
Transformer block (BERT/ViT-style). Set False to match pre-norm Transformers
|
|
such as Whisper or GPT-2 — required when loading pretrained weights from those
|
|
checkpoints.
|
|
self_attention_model: Type of positional encoding and attention scoring scheme. Mirrors
|
|
the Conformer encoder's ``self_attention_model`` choices, plus a ``"no_pos"`` option:
|
|
|
|
- ``"rel_pos"`` (default): Transformer-XL relative positional encoding
|
|
(https://arxiv.org/abs/1901.02860). The relative-position bias is computed in each
|
|
layer and injected into FlexAttention via a ``score_mod`` closure (the (b)+(d)
|
|
terms) plus a ``Q + pos_bias_u`` query rewrite (the (c) term), so the kernel stays
|
|
FlexAttention.
|
|
- ``"abs_pos"``: sinusoidal absolute positional encoding added to the embeddings
|
|
before the first block; standard ``Q @ K^T`` attention via FlexAttention.
|
|
- ``"no_pos"`` (or ``None``): no positional encoding at all — pre-encoder output
|
|
flows straight into ``embed_norm`` and the Transformer blocks. ``xscaling``,
|
|
``pos_emb_max_len``, ``dropout_pre_encoder`` and ``dropout_emb`` have no effect
|
|
in this mode. ``None`` is accepted as a YAML-friendly alias for ``"no_pos"``
|
|
(an unset field in a config maps to ``None``).
|
|
- ``"rope"``: rotary position embedding applied to Q/K inside attention. No additive
|
|
positional embedding is added to the embeddings; ``dropout_pre_encoder`` (and
|
|
``xscaling`` if set) are still applied to the pre-encoder output, and ``dropout_emb``
|
|
has no effect in this mode.
|
|
|
|
``"rel_pos_local_attn"`` is not implemented yet.
|
|
rope_base: Theta base for the rotary position embedding. Only used when
|
|
``self_attention_model='rope'``. Defaults to 10000.0.
|
|
rotary_fraction: Fraction of the per-head dim to rotate. Only used when
|
|
``self_attention_model='rope'``. Defaults to 1.0.
|
|
pos_emb_max_len: Initial maximum length for sinusoidal positional embeddings.
|
|
xscaling: If True, scale embeddings by ``sqrt(d_model)`` before adding positional encodings,
|
|
following "Attention Is All You Need" article. Originally intended to balance the magnitude
|
|
of small-variance token embeddings against unit-bounded sinusoidal positions and to keep
|
|
tied input/pre-softmax logits well-scaled. With modern unit-variance ``nn.Linear``
|
|
pre-encoders and the LayerNorm directly after the positional sum, this scaling is
|
|
largely a no-op for activation magnitudes. Only meaningful when ``pre_block_norm=False``
|
|
or when matching pretrained checkpoints that expect this scaling.
|
|
attn_mode: Attention pattern — currently only "full" (bidirectional) is supported.
|
|
sync_max_audio_length: When true, sync positional encoding allocation length across distributed ranks.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
feat_in: int = 128,
|
|
d_model: int = 512,
|
|
n_heads: int = 8,
|
|
n_layers: int = 17,
|
|
feat_out: int = -1,
|
|
subsampling: str = 'feature_stacking',
|
|
subsampling_factor: int = 4,
|
|
drop_rate: float = 0.1,
|
|
dropout_pre_encoder: float = None,
|
|
dropout_emb: float = 0.0,
|
|
qkv_bias: bool = False,
|
|
qk_norm: bool = False,
|
|
ff_expansion: float = 4.0,
|
|
pre_block_norm: bool = True,
|
|
self_attention_model: Optional[str] = "rel_pos",
|
|
rope_base: float = 10000.0,
|
|
rotary_fraction: float = 1.0,
|
|
pos_emb_max_len: int = 5000,
|
|
xscaling: bool = False,
|
|
attn_mode: str = "full",
|
|
sync_max_audio_length: bool = True,
|
|
):
|
|
super().__init__()
|
|
if d_model % n_heads != 0:
|
|
raise ValueError(f"d_model ({d_model}) must be divisible by n_heads ({n_heads}).")
|
|
if attn_mode not in _SUPPORTED_ATTENTION_MODES:
|
|
raise ValueError(
|
|
f"attn_mode='{attn_mode}' is not yet supported. Supported modes: {_SUPPORTED_ATTENTION_MODES}."
|
|
)
|
|
# ``None`` is accepted as a YAML-friendly alias for ``"no_pos"`` (an unset field in a
|
|
# config simply maps to None) — normalize here so the rest of the module only deals with
|
|
# the string form.
|
|
if self_attention_model is None:
|
|
self_attention_model = "no_pos"
|
|
if self_attention_model not in _SUPPORTED_SELF_ATTENTION_MODELS:
|
|
raise ValueError(
|
|
f"self_attention_model='{self_attention_model}' is not supported. "
|
|
"Currently only 'abs_pos', 'rel_pos', 'rope', and 'no_pos' (or None) are available."
|
|
)
|
|
if dropout_pre_encoder is None:
|
|
dropout_pre_encoder = drop_rate
|
|
|
|
cfg = TransformerEncoderConfig(
|
|
feat_in=feat_in,
|
|
d_model=d_model,
|
|
n_heads=n_heads,
|
|
n_layers=n_layers,
|
|
drop_rate=drop_rate,
|
|
qkv_bias=qkv_bias,
|
|
qk_norm=qk_norm,
|
|
ff_expansion=ff_expansion,
|
|
pre_block_norm=pre_block_norm,
|
|
subsampling_factor=subsampling_factor,
|
|
attn_mode=attn_mode,
|
|
self_attention_model=self_attention_model,
|
|
rope_base=rope_base,
|
|
rotary_fraction=rotary_fraction,
|
|
)
|
|
self.d_model = d_model
|
|
self.n_layers = n_layers
|
|
self._feat_in = feat_in
|
|
self.subsampling = subsampling
|
|
self.subsampling_factor = subsampling_factor
|
|
self.sync_max_audio_length = sync_max_audio_length
|
|
self.self_attention_model = self_attention_model
|
|
self.attn_mode = attn_mode
|
|
|
|
if subsampling == 'feature_stacking':
|
|
self.pre_encode = FeatureStacking(subsampling_factor, feat_in, d_model)
|
|
elif subsampling and subsampling_factor > 1:
|
|
if subsampling in ['stacking', 'stacking_norm']:
|
|
self.pre_encode = StackingSubsampling(
|
|
subsampling_factor=subsampling_factor,
|
|
feat_in=feat_in,
|
|
feat_out=d_model,
|
|
norm=True if subsampling == 'stacking_norm' else False,
|
|
)
|
|
else:
|
|
raise ValueError(
|
|
f"subsampling='{subsampling}' is not supported. "
|
|
"Currently only 'feature_stacking', 'stacking', and 'stacking_norm' are available."
|
|
)
|
|
else:
|
|
self.pre_encode = nn.Linear(feat_in, d_model)
|
|
|
|
self._feat_out = d_model
|
|
if xscaling:
|
|
self.xscale = math.sqrt(d_model)
|
|
else:
|
|
self.xscale = None
|
|
self.pos_emb_max_len = pos_emb_max_len
|
|
if self_attention_model == "rel_pos":
|
|
self.pos_enc = RelPositionalEncoding(
|
|
d_model=d_model,
|
|
dropout_rate=dropout_pre_encoder,
|
|
max_len=pos_emb_max_len,
|
|
xscale=self.xscale,
|
|
dropout_rate_emb=dropout_emb,
|
|
)
|
|
elif self_attention_model == "abs_pos":
|
|
self.pos_enc = PositionalEncoding(
|
|
d_model=d_model,
|
|
dropout_rate=dropout_pre_encoder,
|
|
max_len=pos_emb_max_len,
|
|
xscale=self.xscale,
|
|
dropout_rate_emb=dropout_emb,
|
|
)
|
|
elif self_attention_model == "rope":
|
|
# RoPE has no additive pos-emb step to host dropout, so pre-encoder
|
|
# dropout is applied separately in forward_internal.
|
|
self.dropout_pre_encoder = nn.Dropout(dropout_pre_encoder)
|
|
self.pos_enc = RotaryPositionalEncoding(
|
|
d_k=d_model // n_heads,
|
|
rotary_fraction=rotary_fraction,
|
|
rope_base=rope_base,
|
|
max_len=pos_emb_max_len,
|
|
)
|
|
else: # "no_pos"
|
|
self.pos_enc = None
|
|
self.embed_norm = nn.LayerNorm(d_model) if pre_block_norm else nn.Identity()
|
|
# For 'rope', the shared RotaryPositionalEncoding is passed into each block so the
|
|
# cos/sin buffers are computed once and reused across all attention modules.
|
|
layer_pos_enc = self.pos_enc if self_attention_model == "rope" else None
|
|
self.layers = nn.ModuleList([TransformerBlock(cfg, pos_enc=layer_pos_enc) for _ in range(n_layers)])
|
|
self.final_norm = nn.LayerNorm(d_model)
|
|
if feat_out > 0 and feat_out != self._feat_out:
|
|
self.out_proj = nn.Linear(self._feat_out, feat_out)
|
|
self._feat_out = feat_out
|
|
else:
|
|
self.out_proj = None
|
|
|
|
self.set_max_audio_length(self.pos_emb_max_len)
|
|
|
|
def forward(self, audio_signal, length, bypass_pre_encode=False):
|
|
"""
|
|
Args:
|
|
audio_signal: ``(B, C, T)`` mel spectrogram when ``bypass_pre_encode=False``,
|
|
or ``(B, T, D)`` pre-encoded embeddings when ``bypass_pre_encode=True``.
|
|
length: (B,) — valid frame counts per sample.
|
|
bypass_pre_encode: If true, skip the pre-encoder and consume frame-level embeddings.
|
|
This option is used when pre-encoded embeddings are used as speaker cache
|
|
such as speaker diarization models.
|
|
|
|
Returns:
|
|
x: (B, D, T') — encoded representation (channels-first).
|
|
length: (B,) — output lengths after subsampling.
|
|
"""
|
|
if not bypass_pre_encode and audio_signal.shape[-2] != self._feat_in:
|
|
raise ValueError(
|
|
f"If bypass_pre_encode is False, audio_signal should have shape "
|
|
f"(batch, {self._feat_in}, n_frame) but got last dimension {audio_signal.shape[-2]}."
|
|
)
|
|
if bypass_pre_encode and audio_signal.shape[-1] != self.d_model:
|
|
raise ValueError(
|
|
f"If bypass_pre_encode is True, audio_signal should have shape "
|
|
f"(batch, n_frame, {self.d_model}) but got last dimension {audio_signal.shape[-1]}."
|
|
)
|
|
|
|
if bypass_pre_encode:
|
|
self.update_max_seq_length(seq_length=audio_signal.size(1), device=audio_signal.device)
|
|
else:
|
|
self.update_max_seq_length(seq_length=audio_signal.size(2), device=audio_signal.device)
|
|
return self.forward_internal(audio_signal, length, bypass_pre_encode=bypass_pre_encode)
|
|
|
|
def forward_internal(self, audio_signal, length, bypass_pre_encode=False):
|
|
if length is None:
|
|
length = audio_signal.new_full(
|
|
(audio_signal.size(0),),
|
|
audio_signal.size(1) if bypass_pre_encode else audio_signal.size(-1),
|
|
dtype=torch.int64,
|
|
device=audio_signal.device,
|
|
)
|
|
|
|
if not bypass_pre_encode:
|
|
if isinstance(self.pre_encode, FeatureStacking):
|
|
x, length = self.pre_encode(audio_signal, length)
|
|
else:
|
|
x = torch.transpose(audio_signal, 1, 2)
|
|
if isinstance(self.pre_encode, nn.Linear):
|
|
x = self.pre_encode(x)
|
|
elif not isinstance(self.pre_encode, FeatureStacking):
|
|
x, length = self.pre_encode(x=x, lengths=length)
|
|
length = length.to(torch.int64)
|
|
else:
|
|
x = audio_signal
|
|
length = length.to(torch.int64)
|
|
|
|
if self.self_attention_model == "rope":
|
|
# RoPE: no pos emb added; just apply xscale (if set) + pre-encoder dropout here.
|
|
if self.xscale:
|
|
x = x * self.xscale
|
|
x = self.dropout_pre_encoder(x)
|
|
pos_emb = None
|
|
elif self.pos_enc is not None:
|
|
x, pos_emb = self.pos_enc(x=x)
|
|
else: # "no_pos": pre-encoder output flows in unchanged
|
|
pos_emb = None
|
|
x = self.embed_norm(x)
|
|
|
|
B, T, _ = x.shape
|
|
if self.attn_mode == "causal":
|
|
mask_mod = and_masks(_make_causal_mod(), _make_padding_mod(length))
|
|
else:
|
|
mask_mod = _make_padding_mod(length)
|
|
block_mask = create_block_mask(mask_mod, B=B, H=1, Q_LEN=T, KV_LEN=T, device=x.device)
|
|
# For ``abs_pos`` the positional information is already baked into ``x``, so we don't
|
|
# need to thread ``pos_emb`` through each layer; only ``rel_pos`` consumes it.
|
|
layer_pos_emb = pos_emb if self.self_attention_model == "rel_pos" else None
|
|
for layer in self.layers:
|
|
x = layer(x, block_mask=block_mask, pos_emb=layer_pos_emb)
|
|
|
|
x = self.final_norm(x)
|
|
if self.out_proj is not None:
|
|
x = self.out_proj(x)
|
|
x = x.transpose(1, 2) # (B, T, D) -> (B, D, T)
|
|
length = length.to(dtype=torch.int64)
|
|
return x, length
|
|
|
|
def update_max_seq_length(self, seq_length: int, device):
|
|
"""
|
|
Updates the maximum sequence length for positional encodings.
|
|
|
|
Args:
|
|
seq_length: New maximum sequence length.
|
|
device: Device to use for computations.
|
|
"""
|
|
if self.sync_max_audio_length and torch.distributed.is_initialized():
|
|
global_max_len = torch.tensor([seq_length], dtype=torch.float32, device=device)
|
|
torch.distributed.all_reduce(global_max_len, op=torch.distributed.ReduceOp.MAX)
|
|
seq_length = global_max_len.int().item()
|
|
|
|
if seq_length > self.max_audio_length:
|
|
self.set_max_audio_length(seq_length)
|
|
|
|
def set_max_audio_length(self, max_audio_length):
|
|
"""Sets maximum input length and extends positional encodings if needed."""
|
|
self.max_audio_length = max_audio_length
|
|
if self.pos_enc is None: # "no_pos" mode has no buffer to extend
|
|
return
|
|
device = next(self.parameters()).device
|
|
dtype = next(self.parameters()).dtype
|
|
self.pos_enc.extend_pe(max_audio_length, device, dtype)
|
|
|
|
def freeze(self) -> None:
|
|
freeze(self)
|
|
|
|
def unfreeze(self, partial: bool = False) -> None:
|
|
unfreeze(self, partial=partial)
|