Repository: OpenNMT/OpenNMT-py

Issue:
Issues with Average Attention
Training a Transformer model using Average Attention, gelu, shared vocab, shared embeddings, contect gate, using aan_ffn, bridge, and full context alignment and upon first validation it errors and cancels the training. 

```python
Traceback (most recent call last):
  File "C:\ProgramData\anaconda3\lib\runpy.py", line 196, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "C:\ProgramData\anaconda3\lib\runpy.py", line 86, in _run_code
    exec(code, run_globals)
  File "C:\ProgramData\anaconda3\Scripts\onmt_train.exe\__main__.py", line 7, in <module>
  File "C:\ProgramData\anaconda3\lib\site-packages\onmt\bin\train.py", line 65, in main
    train(opt)
  File "C:\ProgramData\anaconda3\lib\site-packages\onmt\bin\train.py", line 50, in train
    train_process(opt, device_id=0)
  File "C:\ProgramData\anaconda3\lib\site-packages\onmt\train_single.py", line 190, in main
    trainer.train(
  File "C:\ProgramData\anaconda3\lib\site-packages\onmt\trainer.py", line 294, in train
    self._gradient_accumulation(
  File "C:\ProgramData\anaconda3\lib\site-packages\onmt\trainer.py", line 525, in _gradient_accumulation
    raise exc
  File "C:\ProgramData\anaconda3\lib\site-packages\onmt\trainer.py", line 486, in _gradient_accumulation
    preds, texts_ref = self.scoring_preparator.translate(
  File "C:\ProgramData\anaconda3\lib\site-packages\onmt\utils\scoring_utils.py", line 179, in translate
    _, preds_ = translator._translate(
  File "C:\ProgramData\anaconda3\lib\site-packages\onmt\translate\translator.py", line 353, in _translate
    batch_data = self.translate_batch(
  File "C:\ProgramData\anaconda3\lib\site-packages\onmt\translate\translator.py", line 801, in translate_batch
    return self._translate_batch_with_strategy(
  File "C:\ProgramData\anaconda3\lib\site-packages\onmt\translate\translator.py", line 873, in _translate_batch_with_strategy
    self.model.decoder.map_state(fn_map_state)
  File "C:\ProgramData\anaconda3\lib\site-packages\onmt\decoders\transformer.py", line 338, in map_state
    if layer.self_attn.layer_cache[1]['prev_g'].numel() != 0:
KeyError: 'prev_g'
```

I'm using BLEU as the train_metrics and upon first train_eval_steps it throws this error.


Discussion:
Using OpenNMT-py v3.1.1 on Pypi

Is this just a bug or is average attention not compatible with some of the options used? Thanks for any help in advance
Can you please attach your config file ?
```
data:
   ...
share_vocab: True
src_vocab_size: 64000
tgt_vocab_size: 64000
word_vec_size: 512

self_attn_type: 'average'

transformer_ff: 2048
max_relative_positions: 32
hidden_size: 1024
bridge: True
full_context_alignment: True
context_gate: 'both'
encoder_type: 'transformer'
decoder_type: 'transformer'
model_dtype: 'fp16'
pos_ffn_activation_fn: 'gelu'

enc_layers: 12
dec_layers: 1
heads: 8

optim: "adam"
share_embedding: True
share_decoder_embeddings: True

src_subword_model: './sentencepiece.model'
tgt_subword_model: './sentencepiece.model'

param_init: 0
param_init_glorot: True

decay _method: 'noam'
```
Might be a bit off as I wrote it from memory (dont have access to config file atm)
Can you git pull and test again, maybe set the valid steps very low to make sure it works first.

Bear in mind that you are setting options that are unrelated with the new version (or transformer):
context_gate: unrelated with transformer
bridge: same
on the forum you posted a config with max_generator_batches this is deprecated, no longer used.



@vince62s Thanks for the suggestion, I updated to master and there are still issues (but prior issue resolved), now there is no 'prev_g' key error but instead errors like `RuntimeError: The size of tensor a (232) must match the size of tensor b (600) at non-singleton dimension 0` for example

When turning individual things on (valid steps, saving checkpoint, train_eval_steps) at early steps like you said, I realized `train_eval_steps` with `train_metrics` set to "BLEU" is what causes the issue; I forgot to link it above. Below is the configuration I was testing with:

```yaml
save_data: "run"


data:
    ...
    valid:
        path_src: validation.src
        path_tgt: validation.tgt

# Vocabulary files, generated by onmt_build_vocab
src_vocab: run/source.vocab
tgt_vocab: run/source.vocab

tensorboard: true
tensorboard_log_dir: ./tensorboard_logs

model_dtype: "fp16"

src_vocab_size: 32000
tgt_vocab_size: 32000
share_vocab: true


aan_useffn: true
save_model: "./saves/model"


src_subword_model: "general_multi.model"
tgt_subword_model: "general_multi.model"


# testing values 
save_checkpoint_steps: 5
valid_steps: 5

# value which it breaks upon when reached
train_eval_steps: 5
train_eval_steps: "BLEU"


warmup_steps: 8000
report_every: 1


decoder_type: "transformer"
encoder_type: "transformer"

word_vec_size: 1024
hidden_size: 1024

self_attn_type: "average"
transformer_ff: 4096
enc_layers: 8
dec_layers: 8
heads: 16


accum_count: [4]
optim: "adam"
adam_beta1: 0.9
adam_beta2: 0.98
decay_method: "noam"
learning_rate: 4.0
max_grad_norm: 0.0


pos_ffn_activation_fn: "gelu"
full_context_alignment: true

batch_size: 8192
valid_batch_size: 8192
batch_type: "tokens"
normalization: "tokens"
dropout: [0.1]
dropout_steps: [0]
label_smoothing: 0.1


num_workers: 0

batch_size_multiple: 8
vocab_size_multiple: 8

param_init: 0.0
param_init_glorot: true

max_relative_positions: 20
share_embeddings: true
share_decoder_embeddings: true
```

Also worth it to mention the acc and ppl calculations are definitely strange when average attention is enabled. Starting from scratch (I just turned off train_metrics for now) - the accuracy is in the 90s - and the perplexity decreases super rapidly (from 1671 to 1.8 in 150 steps at 8192 tok/batch)

This seems to affect not just average attention
<img width="565" alt="image" src="https://user-images.githubusercontent.com/52015929/236589840-dad1d1b3-dcb3-4c54-b721-378b0d9c6bc8.png">


Resolving patch:
diff --git a/onmt/decoders/transformer.py b/onmt/decoders/transformer.py
--- a/onmt/decoders/transformer.py
+++ b/onmt/decoders/transformer.py
@@ -475,9 +475,9 @@ def forward(self, tgt, enc_out=None, step=None, **kwargs):
                     layer.self_attn.layer_cache = (
                         False, {'keys': torch.tensor([]),
                                 'values': torch.tensor([])})
-                    layer.context_attn.layer_cache = (
-                        False, {'keys': torch.tensor([]),
-                                'values': torch.tensor([])})
+                layer.context_attn.layer_cache = (
+                    False, {'keys': torch.tensor([]),
+                            'values': torch.tensor([])})
 
         tgt_words = tgt[:, :, 0]
 


Test patch: