enable mixed-precision for layernorm

modules/aux_decoder/convnext.py

@@ -3,7 +3,7 @@
 import torch
 import torch.nn as nn
 
-from modules.commons.common_layers import AdamWConv1d
+from modules.commons.common_layers import AdamWConv1d, MixedPrecisionLayerNorm
 
 
 class ConvNeXtBlock(nn.Module):
@@ -26,7 +26,7 @@ def __init__(
  super().__init__()
  self.dwconv = nn.Conv1d(dim, dim, kernel_size=7, padding=3, groups=dim) # depthwise conv
 
- self.norm = nn.LayerNorm(dim, eps=1e-6)
+ self.norm = MixedPrecisionLayerNorm(dim, eps=1e-6)
  self.pwconv1 = nn.Linear(dim, intermediate_dim) # pointwise/1x1 convs, implemented with linear layers
  self.act = nn.GELU()
  self.pwconv2 = nn.Linear(intermediate_dim, dim)

modules/backbones/lynxnet.py

@@ -7,6 +7,7 @@
 
 from modules.commons.common_layers import SinusoidalPosEmb, SwiGLU, Transpose, AdamWConv1d
 from modules.commons.common_layers import KaimingNormalConv1d as Conv1d
+from modules.commons.common_layers import MixedPrecisionLayerNorm as LayerNorm
 from utils.hparams import hparams
 
 
@@ -34,7 +35,7 @@ def __init__(self, dim, expansion_factor, kernel_size=31, activation='PReLU', dr
  else:
  _dropout = nn.Identity()
  self.net = nn.Sequential(
-nn.LayerNorm(dim),
+ LayerNorm(dim),
  Transpose((1, 2)),
  nn.Conv1d(dim, inner_dim * 2, 1),
  SwiGLU(dim=1),
@@ -104,7 +105,7 @@ def __init__(self, in_dims, n_feats, *, num_layers=6, num_channels=512, expansio
  for _ in range(num_layers)
  ]
  )
- self.norm = nn.LayerNorm(num_channels)
+ self.norm = LayerNorm(num_channels)
  self.output_projection = AdamWConv1d(num_channels, in_dims * n_feats, kernel_size=1)
  self.strong_cond = strong_cond
  nn.init.zeros_(self.output_projection.weight)

modules/backbones/lynxnet2.py

@@ -5,6 +5,7 @@
 from modules.commons.common_layers import (
  SinusoidalPosEmb, SwiGLU, ATanGLU, SoftSignGLU, Transpose, AdamWLinear
 )
+from modules.commons.common_layers import MixedPrecisionLayerNorm as LayerNorm
 from utils.hparams import hparams
 
 
@@ -25,7 +26,7 @@ def __init__(self, dim, expansion_factor, kernel_size=31, dropout=0., glu_type='
  else:
  _dropout = nn.Identity()
  self.net = nn.Sequential(
-nn.LayerNorm(dim),
+ LayerNorm(dim),
  Transpose((1, 2)),
  nn.Conv1d(dim, dim, kernel_size=kernel_size, padding=kernel_size // 2, groups=dim),
  Transpose((1, 2)),
@@ -75,7 +76,7 @@ def __init__(self, in_dims, n_feats, *, num_layers=6, num_channels=512, expansio
  for _ in range(num_layers)
  ]
  )
- self.norm = nn.LayerNorm(num_channels)
+ self.norm = LayerNorm(num_channels)
  self.output_projection = AdamWLinear(num_channels, in_dims * n_feats)
  nn.init.kaiming_normal_(self.input_projection.weight)
  nn.init.kaiming_normal_(self.conditioner_projection.weight)

modules/commons/common_layers.py

@@ -287,6 +287,22 @@ def forward(
  return mixed_gammas * x + mixed_betas
 
 
+class MixedPrecisionLayerNorm(nn.LayerNorm):
+ """LayerNorm that keeps fp16/bf16 activations under AMP autocast"""
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ with torch.autocast(device_type=x.device.type, enabled=False):
+ weight = self.weight
+ bias = self.bias
+ if weight is not None and weight.dtype != x.dtype:
+ weight = weight.to(x.dtype)
+ if bias is not None and bias.dtype != x.dtype:
+ bias = bias.to(x.dtype)
+ return F.layer_norm(
+ x, self.normalized_shape, weight, bias, self.eps
+ )
+
+
 class TransformerFFNLayer(nn.Module):
  def __init__(self, hidden_size, filter_size, kernel_size=1, dropout=0., act='gelu'):
  super().__init__()