Conformer workload variants

algorithmic_efficiency/workloads/librispeech_conformer/librispeech_jax/models.py

@@ -12,6 +12,7 @@
   y = layer_norm(x)
 """
 
+import functools
 import math
 from typing import Any, List, Optional
 
@@ -57,6 +58,9 @@ class ConformerConfig:
   batch_norm_momentum: float = 0.999
   batch_norm_epsilon: float = 0.001
   use_specaug: bool = True
+  attention_temperature: float = 1.0
+  activation_function_name: str = 'swish'
+  use_post_layer_norm: bool = True
 
 
 class LayerNorm(nn.Module):
@@ -204,7 +208,15 @@ def __call__(self, inputs, padding_mask=None, train=False):
         use_bias=True,
         kernel_init=nn.initializers.xavier_uniform())(
             inputs)
-    inputs = nn.swish(inputs)
+    if config.activation_function_name == 'swish':
+      activation_fn = nn.swish
+    elif config.activation_function_name == 'gelu':
+      activation_fn = nn.gelu
+    else:
+      raise ValueError('Only "swish" and "gelu" are supported '
+                       'config.activation_function_name values, recieved '
+                       f'{config.activation_function_name}')
+    inputs = activation_fn(inputs)
     inputs = nn.Dropout(rate=config.feed_forward_dropout_rate)(
         inputs, deterministic=not train)
 
@@ -300,7 +312,8 @@ def dot_product_attention(query,
                           dropout_rate=0.,
                           deterministic=False,
                           dtype=jnp.float32,
-                          precision=None):
+                          precision=None,
+                          temperature=1.0):
   """Computes dot-product attention given query, key, and value.
 
   This is the core function for applying attention based on
@@ -360,7 +373,8 @@ def dot_product_attention(query,
 
   # return weighted sum over values for each query position
   return jnp.einsum(
-      '...hqk,...khd->...qhd', attn_weights, value, precision=precision)
+      '...hqk,...khd->...qhd', attn_weights, value,
+      precision=precision) * temperature
 
 
 class MultiHeadedSelfAttention(nn.Module):
@@ -383,6 +397,8 @@ def __call__(self, inputs, paddings, train):
 
     inputs = LayerNorm(dim=config.encoder_dim)(inputs)
 
+    attention_fn = functools.partial(
+        dot_product_attention, temperature=config.attention_temperature)
     result = nn.SelfAttention(
         num_heads=config.num_attention_heads,
         qkv_features=config.encoder_dim,
@@ -392,7 +408,7 @@ def __call__(self, inputs, paddings, train):
         bias_init=nn.initializers.zeros,
         use_bias=True,
         broadcast_dropout=False,
-        attention_fn=dot_product_attention,
+        attention_fn=attention_fn,
         dropout_rate=config.attention_dropout_rate,
         deterministic=not train)(inputs, attention_mask)
 
@@ -531,8 +547,15 @@ def __call__(self, inputs, input_paddings, train):
             inputs)
 
     inputs = BatchNorm(config)(inputs, input_paddings, train)
-
-    inputs = nn.swish(inputs)
+    if config.activation_function_name == 'swish':
+      activation_fn = nn.swish
+    elif config.activation_function_name == 'gelu':
+      activation_fn = nn.gelu
+    else:
+      raise ValueError('Only "swish" and "gelu" are supported '
+                       'config.activation_function_name values, recieved '
+                       f'{config.activation_function_name}')
+    inputs = activation_fn(inputs)
     inputs = nn.Dense(
         config.encoder_dim, kernel_init=nn.initializers.xavier_uniform())(
             inputs)
@@ -579,7 +602,8 @@ def __call__(self, inputs, input_paddings, train):
     inputs = inputs + 0.5 * FeedForwardModule(config=self.config)(
         inputs, padding_mask, train)
 
-    inputs = LayerNorm(dim=config.encoder_dim)(inputs)
+    if config.use_post_layer_norm:
+      inputs = LayerNorm(dim=config.encoder_dim)(inputs)
 
     return inputs
 

algorithmic_efficiency/workloads/librispeech_conformer/librispeech_jax/workload.py

@@ -31,6 +31,18 @@ def __init__(self,
     self.metrics_bundle = metrics.get_metrics_bundle(tokenizer_vocab_path)
     self.use_specaug = use_specaug
 
+  @property
+  def use_post_layer_norm(self) -> bool:
+    return True
+
+  @property
+  def use_gelu(self) -> bool:
+    return False
+
+  @property
+  def attention_temperature(self) -> float:
+    return 1.0
+
   def init_model_fn(
       self,
       rng: spec.RandomState,
@@ -41,11 +53,18 @@ def init_model_fn(
     Here we use dropout_rate as *_residual_dropout_rate, and aux_dropout_rate as
     input_dropout_rate.
     """
+    if self.use_gelu:
+      activation_function_name = 'gelu'
+    else:
+      activation_function_name = 'swish'
     model_config = models.ConformerConfig(
         attention_residual_dropout_rate=dropout_rate,
         feed_forward_residual_dropout_rate=dropout_rate,
         input_dropout_rate=aux_dropout_rate,
-        use_specaug=self.use_specaug)
+        use_specaug=self.use_specaug,
+        attention_temperature=self.attention_temperature,
+        use_post_layer_norm=self.use_post_layer_norm,
+        activation_function_name=activation_function_name)
     self._model = models.Conformer(model_config)
     input_shape = [(320000,), (320000,)]
     fake_input_batch = [np.zeros((2, *x), jnp.float32) for x in input_shape]
@@ -344,3 +363,25 @@ def sync_batch_stats(
     new_model_state = model_state.copy(
         {'batch_stats': avg_fn(model_state['batch_stats'])})
     return new_model_state
+
+
+class LibriSpeechConformerAttentionTemperatureWorkload(
+    LibriSpeechConformerWorkload):
+
+  @property
+  def attention_temperature(self) -> float:
+    return 1.6
+
+
+class LibriSpeechConformerLayerNormWorkload(LibriSpeechConformerWorkload):
+
+  @property
+  def use_post_layer_norm(self) -> bool:
+    return False
+
+
+class LibriSpeechConformerGeluWorkload(LibriSpeechConformerWorkload):
+
+  @property
+  def use_gelu(self) -> bool:
+    return True

algorithmic_efficiency/workloads/librispeech_conformer/librispeech_pytorch/models.py

@@ -3,6 +3,7 @@
 """
 
 from dataclasses import dataclass
+from functools import partial
 import math
 from typing import Tuple
 
@@ -42,6 +43,9 @@ class ConformerConfig:
   batch_norm_momentum: float = 0.999
   batch_norm_epsilon: float = 0.001
   use_specaug: bool = True
+  attention_temperature: float = 1.0
+  activation_function_name: str = 'swish'
+  use_post_layer_norm: bool = True
 
 
 def initialize(m):
@@ -207,7 +211,16 @@ def __init__(self, config: ConformerConfig):
   def forward(self, inputs, padding_mask):
     inputs = self.ln(inputs)
     inputs = self.linear1(inputs)
-    inputs = F.silu(inputs)
+    if self.config.activation_function_name == 'swish':
+      activation_fn = F.silu
+    elif self.config.activation_function_name == 'gelu':
+      # Use tanh approximation of GELU which is default for jax
+      activation_fn = partial(F.gelu, approximate='tanh')
+    else:
+      raise ValueError('Only "swish" and "gelu" are supported '
+                       'config.activation_function_name values, recieved '
+                       f'{self.config.activation_function_name}')
+    inputs = activation_fn(inputs)
     inputs = self.dropout1(inputs)
     inputs = inputs * padding_mask
     inputs = self.linear2(inputs)
@@ -270,6 +283,7 @@ def __init__(self, config: ConformerConfig):
     self.in_proj = nn.Linear(config.encoder_dim, 3 * config.encoder_dim)
     self.out_proj = nn.Linear(config.encoder_dim, config.encoder_dim)
     self.qs = QueryScaler(dim=config.encoder_dim // config.num_attention_heads)
+    self.attention_temperature = config.attention_temperature
 
   def forward(self, inputs, key_padding_mask=None):
     batch_size, seq_len, embed_dim = inputs.shape
@@ -284,6 +298,7 @@ def forward(self, inputs, key_padding_mask=None):
         attn_mask=~key_padding_mask[:, None, None],
         dropout_p=self.dropout,
     ).transpose(1, 2).reshape(batch_size, seq_len, embed_dim)
+    out = out * self.attention_temperature
     out = self.out_proj(out)
     return out
 
@@ -405,7 +420,15 @@ def forward(self, inputs, input_paddings):
     inputs = inputs.permute(0, 2, 1)
 
     inputs = self.bn(inputs, input_paddings)
-    inputs = F.silu(inputs)
+    if self.config.activation_function_name == 'swish':
+      activation_fn = F.silu
+    elif self.config.activation_function_name == 'gelu':
+      activation_fn = F.gelu
+    else:
+      raise ValueError('Only "swish" and "gelu" are supported '
+                       'config.activation_function_name values, recieved '
+                       f'{self.config.activation_function_name}')
+    inputs = activation_fn(inputs)
     inputs = self.lin3(inputs)
 
     inputs = self.dropout(inputs)
@@ -421,15 +444,18 @@ def __init__(self, config: ConformerConfig):
     self.mhsa = MultiHeadedSelfAttention(config)
     self.conv = ConvolutionBlock(config)
     self.ff2 = FeedForwardModule(config)
-    self.ln = LayerNorm(dim=config.encoder_dim)
+    self.ln = None
+    if config.use_post_layer_norm:
+      self.ln = LayerNorm(dim=config.encoder_dim)
 
   def forward(self, inputs, input_paddings):
     padding_mask = 1 - input_paddings[:, :, None]
     inputs = inputs + 0.5 * self.ff1(inputs, padding_mask)
     inputs = inputs + self.mhsa(inputs, input_paddings)
     inputs = inputs + self.conv(inputs, input_paddings)
     inputs = inputs + 0.5 * self.ff2(inputs, padding_mask)
-    inputs = self.ln(inputs)
+    if self.ln:
+      inputs = self.ln(inputs)
     return inputs
 
 

algorithmic_efficiency/workloads/librispeech_conformer/librispeech_pytorch/workload.py

@@ -20,7 +20,7 @@
 from algorithmic_efficiency.workloads.librispeech_conformer.input_pipeline import \
     LibriSpeechDataset
 from algorithmic_efficiency.workloads.librispeech_conformer.librispeech_pytorch import \
-    models as conformer_model
+    models
 
 USE_PYTORCH_DDP, RANK, DEVICE, N_GPUS = pytorch_utils.pytorch_setup()
 
@@ -36,6 +36,18 @@ def __init__(self,
     self.tokenizer = metrics.load_tokenizer(tokenizer_vocab_path)
     self.use_specaug = use_specaug
 
+  @property
+  def use_post_layer_norm(self) -> bool:
+    return True
+
+  @property
+  def use_gelu(self) -> bool:
+    return False
+
+  @property
+  def attention_temperature(self) -> float:
+    return 1.0
+
   def init_model_fn(
       self,
       rng: spec.RandomState,
@@ -52,15 +64,22 @@ def init_model_fn(
     torch.backends.cuda.enable_flash_sdp(False)
     torch.backends.cuda.enable_mem_efficient_sdp(False)
     torch.backends.cuda.enable_math_sdp(True)
-    model = conformer_model.ConformerEncoderDecoder(
-        conformer_model.ConformerConfig(
+    if self.use_gelu:
+      activation_function_name = 'gelu'
+    else:
+      activation_function_name = 'swish'
+    model = models.ConformerEncoderDecoder(
+        models.ConformerConfig(
             attention_residual_dropout_rate=dropout_rate,
             feed_forward_residual_dropout_rate=dropout_rate,
             conv_residual_dropout_rate=dropout_rate,
             input_dropout_rate=aux_dropout_rate,
-            use_specaug=self.use_specaug))
+            use_specaug=self.use_specaug,
+            attention_temperature=self.attention_temperature,
+            use_post_layer_norm=self.use_post_layer_norm,
+            activation_function_name=activation_function_name))
     self.ctc_loss = torch.nn.CTCLoss(blank=0, reduction='none')
-    conformer_model.initialize(model)
+    models.initialize(model)
     self._param_shapes = param_utils.pytorch_param_shapes(model)
     self._param_types = param_utils.pytorch_param_types(self._param_shapes)
     model.to(DEVICE)
@@ -310,3 +329,25 @@ def _eval_model_on_split(self,
             float(total_metrics['word_errors'].item() /
                   total_metrics['num_words'].item()),
     }
+
+
+class LibriSpeechConformerAttentionTemperatureWorkload(
+    LibriSpeechConformerWorkload):
+
+  @property
+  def attention_temperature(self) -> float:
+    return 1.6
+
+
+class LibriSpeechConformerLayerNormWorkload(LibriSpeechConformerWorkload):
+
+  @property
+  def use_post_layer_norm(self) -> bool:
+    return False
+
+
+class LibriSpeechConformerGeluWorkload(LibriSpeechConformerWorkload):
+
+  @property
+  def use_gelu(self) -> bool:
+    return True

algorithmic_efficiency/workloads/librispeech_conformer/workload.py

@@ -13,6 +13,18 @@ def target_metric_name(self) -> str:
     """The name of the target metric (useful for scoring/processing code)."""
     return 'wer'
 
+  @property
+  def use_post_layer_norm(self) -> bool:
+    raise NotImplementedError
+
+  @property
+  def use_gelu(self) -> bool:
+    raise NotImplementedError
+
+  @property
+  def attention_temperature(self) -> float:
+    raise NotImplementedError
+
   def has_reached_validation_target(self, eval_result: Dict[str,
                                                             float]) -> bool:
     return eval_result['validation/wer'] < self.validation_target_value

algorithmic_efficiency/workloads/workloads.py

@@ -48,6 +48,20 @@
         'workload_path': 'librispeech_conformer/librispeech',
         'workload_class_name': 'LibriSpeechConformerWorkload',
     },
+    'librispeech_conformer_attention_temperature': {
+        'workload_path':
+            'librispeech_conformer/librispeech',
+        'workload_class_name':
+            'LibriSpeechConformerAttentionTemperatureWorkload',
+    },
+    'librispeech_conformer_layernorm': {
+        'workload_path': 'librispeech_conformer/librispeech',
+        'workload_class_name': 'LibriSpeechConformerLayerNormWorkload',
+    },
+    'librispeech_conformer_gelu': {
+        'workload_path': 'librispeech_conformer/librispeech',
+        'workload_class_name': 'LibriSpeechConformerGeluWorkload',
+    },
     'librispeech_deepspeech': {
         'workload_path': 'librispeech_deepspeech/librispeech',
         'workload_class_name': 'LibriSpeechDeepSpeechWorkload',

docker/scripts/startup.sh

@@ -114,7 +114,9 @@ VALID_DATASETS=("criteo1tb" "imagenet"  "fastmri" "ogbg" "librispeech" \
                 "wmt" "mnist")
 VALID_WORKLOADS=("criteo1tb" "imagenet_resnet" "imagenet_vit" "fastmri" "ogbg" \
                  "wmt" "librispeech_deepspeech" "librispeech_conformer" "mnist" \
-                 "criteo1tb_resnet" "criteo1tb_layernorm" "criteo1tb_embed_init")
+                 "criteo1tb_resnet" "criteo1tb_layernorm" "criteo_embed_init" \
+                 "conformer_layernorm" "conformer_attention_temperature" \
+                 "conformer_gelu")
 
 
 # Set data and experiment paths