Fix for Pre-LN Transformers
State: closed | Pull request: yes

This PR fixes bugs in the Pre-LN implementation in transformer encoders and decoders.

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.