# Copyright (c) 2022, 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, field
from typing import Any, Optional
import torch
from torch import nn as nn
from nemo.collections.asr.parts.submodules import multi_head_attention as mha
from nemo.collections.common.parts import adapter_modules
from nemo.core.classes.mixins import adapter_mixin_strategies
[docs]
class MHAResidualAddAdapterStrategy(adapter_mixin_strategies.ResidualAddAdapterStrategy):
"""
An implementation of residual addition of an adapter module with its input for the MHA Adapters.
"""
[docs]
def forward(self, input: dict, adapter: torch.nn.Module, *, module: 'AdapterModuleMixin'):
"""
A basic strategy, comprising of a residual connection over the input, after forward pass by
the underlying adapter. Additional work is done to pack and unpack the dictionary of inputs and outputs.
Note: The `value` tensor is added to the output of the attention adapter as the residual connection.
Args:
input: A dictionary of multiple input arguments for the adapter module.
`query`, `key`, `value`: Original output tensor of the module, or the output of the
previous adapter (if more than one adapters are enabled).
`mask`: Attention mask.
`pos_emb`: Optional positional embedding for relative encoding.
adapter: The adapter module that is currently required to perform the forward pass.
module: The calling module, in its entirety. It is a module that implements `AdapterModuleMixin`,
therefore the strategy can access all other adapters in this module via `module.adapter_layer`.
Returns:
The result tensor, after one of the active adapters has finished its forward passes.
"""
out = self.compute_output(input, adapter, module=module)
value_name = None
if 'value' in input:
value_name = 'value'
elif 'values' in input:
value_name = 'values'
else:
raise ValueError(
"Input dictionary must contain 'value' or 'values' key for residual connection. Input "
f"dictionary keys: {input.keys()}"
)
# If not in training mode, or probability of stochastic depth is 0, skip step.
p = self.stochastic_depth
if not module.training or p == 0.0:
pass
else:
out = self.apply_stochastic_depth(out, input[value_name], adapter, module=module)
# Return the residual connection output = input + adapter(input)
result = input[value_name] + out
# If l2_lambda is activated, register the loss value
self.compute_auxiliary_losses(result, input[value_name], adapter, module=module)
return result
[docs]
def compute_output(
self, input: torch.Tensor, adapter: torch.nn.Module, *, module: 'AdapterModuleMixin'
) -> torch.Tensor:
"""
Compute the output of a single adapter to some input.
Args:
input: Original output tensor of the module, or the output of the previous adapter (if more than
one adapters are enabled).
adapter: The adapter module that is currently required to perform the forward pass.
module: The calling module, in its entirety. It is a module that implements `AdapterModuleMixin`,
therefore the strategy can access all other adapters in this module via `module.adapter_layer`.
Returns:
The result tensor, after one of the active adapters has finished its forward passes.
"""
if isinstance(input, (list, tuple)):
out = adapter(*input)
elif isinstance(input, dict):
out = adapter(**input)
else:
out = adapter(input)
return out
@dataclass
class MHAResidualAddAdapterStrategyConfig(adapter_mixin_strategies.ResidualAddAdapterStrategyConfig):
_target_: str = "{0}.{1}".format(
MHAResidualAddAdapterStrategy.__module__, MHAResidualAddAdapterStrategy.__name__
) # mandatory field
[docs]
class MultiHeadAttentionAdapter(mha.MultiHeadAttention, adapter_modules.AdapterModuleUtil):
"""Multi-Head Attention layer of Transformer.
Args:
n_head (int): number of heads
n_feat (int): size of the features
dropout_rate (float): dropout rate
proj_dim (int, optional): Optional integer value for projection before computing attention.
If None, then there is no projection (equivalent to proj_dim = n_feat).
If > 0, then will project the n_feat to proj_dim before calculating attention.
If <0, then will equal n_head, so that each head has a projected dimension of 1.
adapter_strategy: By default, MHAResidualAddAdapterStrategyConfig. An adapter composition function object.
"""
def __init__(
self,
n_head: int,
n_feat: int,
dropout_rate: float,
proj_dim: Optional[int] = None,
adapter_strategy: MHAResidualAddAdapterStrategy = None,
use_pytorch_sdpa: bool = False,
use_pytorch_sdpa_backends: Optional[list] = None,
):
super().__init__(
n_head=n_head,
n_feat=n_feat,
dropout_rate=dropout_rate,
max_cache_len=0,
use_pytorch_sdpa=use_pytorch_sdpa,
use_pytorch_sdpa_backends=use_pytorch_sdpa_backends,
)
self.pre_norm = nn.LayerNorm(n_feat)
# Set the projection dim to number of heads automatically
if proj_dim is not None and proj_dim < 1:
proj_dim = n_head
self.proj_dim = proj_dim
# Recompute weights for projection dim
if self.proj_dim is not None:
if self.proj_dim % n_head != 0:
raise ValueError(f"proj_dim ({proj_dim}) is not divisible by n_head ({n_head})")
self.d_k = self.proj_dim // n_head
self.s_d_k = math.sqrt(self.d_k)
self.linear_q = nn.Linear(n_feat, self.proj_dim)
self.linear_k = nn.Linear(n_feat, self.proj_dim)
self.linear_v = nn.Linear(n_feat, self.proj_dim)
self.linear_out = nn.Linear(self.proj_dim, n_feat)
# Setup adapter strategy
self.setup_adapter_strategy(adapter_strategy)
# reset parameters for Q to be identity operation
self.reset_parameters()
[docs]
def forward(self, query, key, value, mask, pos_emb=None, cache=None):
"""Compute 'Scaled Dot Product Attention'.
Args:
query (torch.Tensor): (batch, time1, size)
key (torch.Tensor): (batch, time2, size)
value(torch.Tensor): (batch, time2, size)
mask (torch.Tensor): (batch, time1, time2)
cache (torch.Tensor) : (batch, time_cache, size)
returns:
output (torch.Tensor): transformed `value` (batch, time1, d_model) weighted by the query dot key attention
cache (torch.Tensor) : (batch, time_cache_next, size)
"""
# Need to perform duplicate computations as at this point the tensors have been
# separated by the adapter forward
query = self.pre_norm(query)
key = self.pre_norm(key)
value = self.pre_norm(value)
return super().forward(query, key, value, mask, pos_emb, cache=cache)
[docs]
def reset_parameters(self):
with torch.no_grad():
nn.init.zeros_(self.linear_out.weight)
nn.init.zeros_(self.linear_out.bias)
[docs]
def get_default_strategy_config(self) -> 'dataclass':
return MHAResidualAddAdapterStrategyConfig()
@dataclass
class MultiHeadAttentionAdapterConfig:
n_head: int
n_feat: int
dropout_rate: float = 0.0
proj_dim: Optional[int] = None
adapter_strategy: Optional[Any] = field(default_factory=lambda: MHAResidualAddAdapterStrategyConfig())
use_pytorch_sdpa: bool = False
use_pytorch_sdpa_backends: Optional[list] = None
_target_: str = "{0}.{1}".format(MultiHeadAttentionAdapter.__module__, MultiHeadAttentionAdapter.__name__)
[docs]
class RelPositionMultiHeadAttentionAdapter(mha.RelPositionMultiHeadAttention, adapter_modules.AdapterModuleUtil):
"""Multi-Head Attention layer of Transformer-XL with support of relative positional encoding.
Paper: https://arxiv.org/abs/1901.02860
Args:
n_head (int): number of heads
n_feat (int): size of the features
dropout_rate (float): dropout rate
proj_dim (int, optional): Optional integer value for projection before computing attention.
If None, then there is no projection (equivalent to proj_dim = n_feat).
If > 0, then will project the n_feat to proj_dim before calculating attention.
If <0, then will equal n_head, so that each head has a projected dimension of 1.
adapter_strategy: By default, MHAResidualAddAdapterStrategyConfig. An adapter composition function object.
"""
def __init__(
self,
n_head: int,
n_feat: int,
dropout_rate: float,
proj_dim: Optional[int] = None,
adapter_strategy: MHAResidualAddAdapterStrategyConfig = None,
use_pytorch_sdpa: bool = False,
use_pytorch_sdpa_backends: Optional[list] = None,
):
super().__init__(
n_head=n_head,
n_feat=n_feat,
dropout_rate=dropout_rate,
pos_bias_u=None,
pos_bias_v=None,
max_cache_len=0,
use_pytorch_sdpa=use_pytorch_sdpa,
use_pytorch_sdpa_backends=use_pytorch_sdpa_backends,
)
self.pre_norm = nn.LayerNorm(n_feat)
# Set the projection dim to number of heads automatically
if proj_dim is not None and proj_dim < 1:
proj_dim = n_head
self.proj_dim = proj_dim
# Recompute weights for projection dim
if self.proj_dim is not None:
if self.proj_dim % n_head != 0:
raise ValueError(f"proj_dim ({proj_dim}) is not divisible by n_head ({n_head})")
self.d_k = self.proj_dim // n_head
self.s_d_k = math.sqrt(self.d_k)
self.linear_q = nn.Linear(n_feat, self.proj_dim)
self.linear_k = nn.Linear(n_feat, self.proj_dim)
self.linear_v = nn.Linear(n_feat, self.proj_dim)
self.linear_out = nn.Linear(self.proj_dim, n_feat)
self.linear_pos = nn.Linear(n_feat, self.proj_dim, bias=False)
self.pos_bias_u = nn.Parameter(torch.FloatTensor(self.h, self.d_k))
self.pos_bias_v = nn.Parameter(torch.FloatTensor(self.h, self.d_k))
# Setup adapter strategy
self.setup_adapter_strategy(adapter_strategy)
# reset parameters for Q to be identity operation
self.reset_parameters()
[docs]
def forward(self, query, key, value, mask, pos_emb, cache=None):
"""Compute 'Scaled Dot Product Attention' with rel. positional encoding.
Args:
query (torch.Tensor): (batch, time1, size)
key (torch.Tensor): (batch, time2, size)
value(torch.Tensor): (batch, time2, size)
mask (torch.Tensor): (batch, time1, time2)
pos_emb (torch.Tensor) : (batch, time1, size)
cache (torch.Tensor) : (batch, time_cache, size)
Returns:
output (torch.Tensor): transformed `value` (batch, time1, d_model) weighted by the query dot key attention
cache_next (torch.Tensor) : (batch, time_cache_next, size)
"""
# Need to perform duplicate computations as at this point the tensors have been
# separated by the adapter forward
query = self.pre_norm(query)
key = self.pre_norm(key)
value = self.pre_norm(value)
return super().forward(query, key, value, mask, pos_emb, cache=cache)
[docs]
def reset_parameters(self):
with torch.no_grad():
nn.init.zeros_(self.linear_out.weight)
nn.init.zeros_(self.linear_out.bias)
# NOTE: This exact procedure apparently highly important.
# Above operation is safe to do as self.linear_out.weight *= 0.0 (similar for bias)
# However:
# DO NOT REPLACE BELOW WITH self.pos_bias_u *= 0.0 OR self.pos_bias_v *= 0.0
# For some reason at init sometimes it will cause the value of the tensor to become NaN
# All operations to compute matrix_ac and matrix_bd will then fail.
nn.init.zeros_(self.pos_bias_u)
nn.init.zeros_(self.pos_bias_v)
[docs]
def get_default_strategy_config(self) -> 'dataclass':
return MHAResidualAddAdapterStrategyConfig()
@dataclass
class RelPositionMultiHeadAttentionAdapterConfig:
n_head: int
n_feat: int
dropout_rate: float = 0.0
proj_dim: Optional[int] = None
adapter_strategy: Optional[Any] = field(default_factory=lambda: MHAResidualAddAdapterStrategyConfig())
use_pytorch_sdpa: bool = False
use_pytorch_sdpa_backends: Optional[list] = None
_target_: str = "{0}.{1}".format(
RelPositionMultiHeadAttentionAdapter.__module__, RelPositionMultiHeadAttentionAdapter.__name__
)
[docs]
class PositionalEncodingAdapter(mha.PositionalEncoding, adapter_modules.AdapterModuleUtil):
"""
Absolute positional embedding adapter.
.. note::
Absolute positional embedding value is added to the input tensor *without residual connection* !
Therefore, the input is changed, if you only require the positional embedding, drop the returned `x` !
Args:
d_model (int): The input dimension of x.
max_len (int): The max sequence length.
xscale (float): The input scaling factor. Defaults to 1.0.
adapter_strategy (AbstractAdapterStrategy): By default, ReturnResultAdapterStrategyConfig.
An adapter composition function object.
NOTE: Since this is a positional encoding, it will not add a residual !
"""
def __init__(
self,
d_model: int,
max_len: int = 5000,
xscale=1.0,
adapter_strategy: adapter_mixin_strategies.ReturnResultAdapterStrategyConfig = None,
):
super().__init__(
d_model=d_model,
dropout_rate=0.0,
max_len=max_len,
xscale=xscale,
dropout_rate_emb=0.0,
)
# Setup adapter strategy
self.setup_adapter_strategy(adapter_strategy)
[docs]
def get_default_strategy_config(self) -> 'dataclass':
return adapter_mixin_strategies.ReturnResultAdapterStrategyConfig()
@dataclass
class PositionalEncodingAdapterConfig:
d_model: int
max_len: int = 5000
xscale: float = 1.0
adapter_strategy: Optional[Any] = field(
default_factory=lambda: adapter_mixin_strategies.ResidualAddAdapterStrategyConfig()
)
_target_: str = "{0}.{1}".format(PositionalEncodingAdapter.__module__, PositionalEncodingAdapter.__name__)
[docs]
class RelPositionalEncodingAdapter(mha.RelPositionalEncoding, adapter_modules.AdapterModuleUtil):
"""
Relative positional encoding for TransformerXL's layers
See : Appendix B in https://arxiv.org/abs/1901.02860
.. note::
Relative positional embedding value is **not** added to the input tensor !
Therefore, the input should be updated changed, if you only require the positional embedding, drop the returned `x` !
Args:
d_model (int): embedding dim
max_len (int): maximum input length
xscale (bool): whether to scale the input by sqrt(d_model)
adapter_strategy: By default, ReturnResultAdapterStrategyConfig. An adapter composition function object.
"""
def __init__(
self,
d_model: int,
max_len: int = 5000,
xscale=1.0,
adapter_strategy: adapter_mixin_strategies.ReturnResultAdapterStrategyConfig = None,
):
super().__init__(d_model=d_model, dropout_rate=0.0, max_len=max_len, xscale=xscale, dropout_rate_emb=0.0)
# Setup adapter strategy
self.setup_adapter_strategy(adapter_strategy)
[docs]
def get_default_strategy_config(self) -> 'dataclass':
return adapter_mixin_strategies.ReturnResultAdapterStrategyConfig()
@dataclass
class RelPositionalEncodingAdapterConfig:
d_model: int
max_len: int = 5000
xscale: float = 1.0
adapter_strategy: Optional[Any] = field(
default_factory=lambda: adapter_mixin_strategies.ResidualAddAdapterStrategyConfig()
)
_target_: str = "{0}.{1}".format(RelPositionalEncodingAdapter.__module__, RelPositionalEncodingAdapter.__name__)
