Repository: NVIDIA/NeMo

Issue:
Fix for Pre-LN Transformers
This PR fixes bugs in the Pre-LN implementation in transformer encoders and decoders.


Discussion:
if the encoders & decoders are basically the same thing, would it make sense to have a single file that they import blocks from?
> if the encoders & decoders are basically the same thing, would it make sense to have a single file that they import blocks from?

That is already happening to a good extent - both import MultiHeadAttention and PositionWiseFF blocks and that logic lies outside the Encoder/Decoder layers.
This pull request **introduces 1 alert** when merging 21a598eb581a78f1dfa171dc1b41d2cb7ca21732 into a33d2794cba9cff5b29044f787e41f654325424a - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-afa1324fdce133834d87c84752630c214b1d11cb)

**new alerts:**

* 1 for Unused import
LGTM
EDIT: it's the same return issue i pointed out.


dont merge yet; i am getting some sort of error:

```
Traceback (most recent call last):
  File "examples/nlp/machine_translation/enc_dec_nmt.py", line 127, in main
    trainer.fit(mt_model)
  File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/trainer/trainer.py", line 513, in fit
    self.dispatch()
  File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/trainer/trainer.py", line 553, in dispatch
    self.accelerator.start_training(self)
  File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/accelerators/accelerator.py", line 74, in start_training
    self.training_type_plugin.start_training(trainer)
  File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/plugins/training_type/training_type_plugin.py", line 111, in start_training
    self._results = trainer.run_train()
  File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/trainer/trainer.py", line 614, in run_train
    self.run_sanity_check(self.lightning_module)
  File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/trainer/trainer.py", line 863, in run_sanity_check
    _, eval_results = self.run_evaluation(max_batches=self.num_sanity_val_batches)
  File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/trainer/trainer.py", line 732, in run_evaluation
    output = self.evaluation_loop.evaluation_step(batch, batch_idx, dataloader_idx)
  File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/trainer/evaluation_loop.py", line 164, in evaluation_step
    output = self.trainer.accelerator.validation_step(args)
  File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/accelerators/accelerator.py", line 178, in validation_step
    return self.training_type_plugin.validation_step(*args)
  File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/plugins/training_type/ddp.py", line 290, in validation_step
    return self.model(*args, **kwargs)
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/parallel/distributed.py", line 671, in forward
    output = self.module(*inputs[0], **kwargs[0])
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/overrides/base.py", line 63, in forward
    output = self.module.validation_step(*inputs, **kwargs)
  File "/home/dcg-adlr-mchrzanowski-source/NeMo2/nemo/collections/nlp/models/machine_translation/mt_enc_dec_model.py", line 234, in validation_step
    return self.eval_step(batch, batch_idx, 'val')
  File "/home/dcg-adlr-mchrzanowski-source/NeMo2/nemo/collections/nlp/models/machine_translation/mt_enc_dec_model.py", line 199, in eval_step
    log_probs = self(src_ids, src_mask, tgt_ids, tgt_mask)
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/dcg-adlr-mchrzanowski-source/NeMo2/nemo/core/classes/common.py", line 721, in __call__
    return wrapped(*args, **kwargs)
  File "/home/dcg-adlr-mchrzanowski-source/NeMo2/nemo/collections/nlp/models/machine_translation/mt_enc_dec_model.py", line 167, in forward
    src_hiddens = self.encoder(src, src_mask)
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/dcg-adlr-mchrzanowski-source/NeMo2/nemo/collections/nlp/modules/common/transformer/transformer.py", line 103, in forward
    encoder_hidden_states = self._encoder(embeddings, encoder_mask)
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/dcg-adlr-mchrzanowski-source/NeMo2/nemo/collections/nlp/modules/common/transformer/transformer_encoders.py", line 158, in forward
    encoder_states = layer(encoder_states, encoder_attn_mask, memory_states)
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/dcg-adlr-mchrzanowski-source/NeMo2/nemo/collections/nlp/modules/common/transformer/transformer_encoders.py", line 99, in forward
    self.forward_preln(encoder_query, encoder_mask, encoder_keys)
  File "/home/dcg-adlr-mchrzanowski-source/NeMo2/nemo/collections/nlp/modules/common/transformer/transformer_encoders.py", line 70, in forward_preln
    encoder_query = self.layer_norm_1(encoder_query)
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/normalization.py", line 170, in forward
    input, self.normalized_shape, self.weight, self.bias, self.eps)
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/functional.py", line 2072, in layer_norm
    torch.backends.cudnn.enabled)
TypeError: layer_norm(): argument 'input' (position 1) must be Tensor, not NoneType
```

somewhere in between 3676721 and 9efb1e7
Good catch @mchrzanowski .  I would have thought one of our tests would have failed for something like this.
Nice catch, our CI runs with`pre_ln=false` and this was happening only when `pre_ln=true`. Fixed now.

Resolving patch:
diff --git a/nemo/collections/nlp/modules/common/transformer/transformer_decoders.py b/nemo/collections/nlp/modules/common/transformer/transformer_decoders.py
--- a/nemo/collections/nlp/modules/common/transformer/transformer_decoders.py
+++ b/nemo/collections/nlp/modules/common/transformer/transformer_decoders.py
@@ -66,34 +66,51 @@ def __init__(
         self.layer_norm_3 = nn.LayerNorm(hidden_size, eps=1e-5)
         self.third_sub_layer = PositionWiseFF(hidden_size, inner_size, ffn_dropout, hidden_act)
 
-    # TODO: add Neural Types
-    def forward(self, decoder_query, decoder_mask, decoder_keys, encoder_states, encoder_mask):
+    def forward_preln(self, decoder_query, decoder_mask, decoder_keys, encoder_states, encoder_mask):
+        """
+        Pre-LayerNorm block
+        Order of operations: LN -> Self-Attn -> Residual -> LN -> Cross-Attn -> Residual -> LN -> FFN
+        """
+        residual = decoder_query
+        decoder_query = self.layer_norm_1(decoder_query)
+        decoder_keys = self.layer_norm_1(decoder_keys)
+        self_attn_output = self.first_sub_layer(decoder_query, decoder_keys, decoder_keys, decoder_mask)
+        self_attn_output += residual
 
-        # Pre-LN: LN -> Self-Attn -> Drop -> Residual -> LN -> Cross-Attn -> Drop -> Residual -> LN -> FFN
-        # Post-LN: Self-Attn -> Drop -> Residual -> LN -> Cross-Attn -> Drop -> Residual -> LN -> FFN -> Residual -> LN
-        if self.pre_ln:
-            # Share same LN params for query, key (self-attn)
-            decoder_query = self.layer_norm_1(decoder_query)
-            decoder_keys = self.layer_norm_1(decoder_keys)
+        residual = self_attn_output
+        self_attn_output = self.layer_norm_2(self_attn_output)
+        enc_dec_attn_output = self.second_sub_layer(self_attn_output, encoder_states, encoder_states, encoder_mask)
+        enc_dec_attn_output += residual
+
+        residual = enc_dec_attn_output
+        enc_dec_attn_output = self.layer_norm_3(enc_dec_attn_output)
+        output_states = self.third_sub_layer(enc_dec_attn_output)
+        output_states += residual
+
+        return output_states
 
+    def forward_postln(self, decoder_query, decoder_mask, decoder_keys, encoder_states, encoder_mask):
+        """
+        Post-LayerNorm block
+        Order of operations: Self-Attn -> Residual -> LN -> Cross-Attn -> Residual -> LN -> FFN -> Residual -> LN
+        """
         self_attn_output = self.first_sub_layer(decoder_query, decoder_keys, decoder_keys, decoder_mask)
         self_attn_output += decoder_query
-
-        self_attn_output = self.layer_norm_2(self_attn_output) if self.pre_ln else self.layer_norm_1(self_attn_output)
+        self_attn_output = self.layer_norm_1(self_attn_output)
 
         enc_dec_attn_output = self.second_sub_layer(self_attn_output, encoder_states, encoder_states, encoder_mask)
         enc_dec_attn_output += self_attn_output
-
-        enc_dec_attn_output = (
-            self.layer_norm_3(enc_dec_attn_output) if self.pre_ln else self.layer_norm_2(enc_dec_attn_output)
-        )
+        enc_dec_attn_output = self.layer_norm_2(enc_dec_attn_output)
 
         output_states = self.third_sub_layer(enc_dec_attn_output)
+        output_states += enc_dec_attn_output
+        return self.layer_norm_3(output_states)
 
-        if not self.pre_ln:
-            output_states = self.layer_norm_3(output_states + enc_dec_attn_output)
-
-        return output_states
+    def forward(self, decoder_query, decoder_mask, decoder_keys, encoder_states, encoder_mask):
+        if self.pre_ln:
+            return self.forward_preln(decoder_query, decoder_mask, decoder_keys, encoder_states, encoder_mask)
+        else:
+            return self.forward_postln(decoder_query, decoder_mask, decoder_keys, encoder_states, encoder_mask)
 
 
 class TransformerDecoder(nn.Module):
diff --git a/nemo/collections/nlp/modules/common/transformer/transformer_encoders.py b/nemo/collections/nlp/modules/common/transformer/transformer_encoders.py
--- a/nemo/collections/nlp/modules/common/transformer/transformer_encoders.py
+++ b/nemo/collections/nlp/modules/common/transformer/transformer_encoders.py
@@ -61,26 +61,45 @@ def __init__(
         self.layer_norm_2 = nn.LayerNorm(hidden_size, eps=1e-5)
         self.second_sub_layer = PositionWiseFF(hidden_size, inner_size, ffn_dropout, hidden_act)
 
-    def forward(self, encoder_query, encoder_mask, encoder_keys):
+    def forward_preln(self, encoder_query, encoder_mask, encoder_keys):
+        """
+        Pre-LayerNorm block
+        Order of operations: LN -> Self-Attn -> Residual -> LN -> Cross-Attn -> Residual -> LN -> FFN
+        """
+        residual = encoder_query
+        encoder_query = self.layer_norm_1(encoder_query)
+        encoder_keys = self.layer_norm_1(encoder_keys)
+        self_attn_output = self.first_sub_layer(encoder_query, encoder_keys, encoder_keys, encoder_mask)
+        self_attn_output += residual
 
-        # Pre-LN: LN -> Attn -> Drop -> Residual -> LN -> FFN
-        # Post-LN: Attn -> Drop -> Residual -> LN -> FFN -> Residual -> LN
-        if self.pre_ln:
-            # Share same LN params for query, key (self-attn)
-            encoder_query = self.layer_norm_1(encoder_query)
-            encoder_keys = self.layer_norm_1(encoder_keys)
+        residual = self_attn_output
+        self_attn_output = self.layer_norm_2(self_attn_output)
+        output_states = self.second_sub_layer(self_attn_output)
+        output_states += residual
+
+        return output_states
 
+    def forward_postln(self, encoder_query, encoder_mask, encoder_keys):
+        """
+        Post-LayerNorm block
+        Order of operations: Self-Attn -> Residual -> LN -> Cross-Attn -> Residual -> LN -> FFN -> Residual -> LN
+        """
         self_attn_output = self.first_sub_layer(encoder_query, encoder_keys, encoder_keys, encoder_mask)
         self_attn_output += encoder_query
-
-        self_attn_output = self.layer_norm_2(self_attn_output) if self.pre_ln else self.layer_norm_1(self_attn_output)
+        self_attn_output = self.layer_norm_1(self_attn_output)
 
         output_states = self.second_sub_layer(self_attn_output)
+        output_states += self_attn_output
+        output_states = self.layer_norm_2(output_states)
 
-        if not self.pre_ln:
-            output_states = self.layer_norm_2(output_states + self_attn_output)
         return output_states
 
+    def forward(self, encoder_query, encoder_mask, encoder_keys):
+        if self.pre_ln:
+            return self.forward_preln(encoder_query, encoder_mask, encoder_keys)
+        else:
+            return self.forward_postln(encoder_query, encoder_mask, encoder_keys)
+
 
 class TransformerEncoder(nn.Module):
     def __init__(


Test patch: