Repository: OpenNMT/OpenNMT-tf

Issue:
Training fails with --mixed_precision and guided alignments
This occurs with OpenNMT-TF 2.28.00. I have a base mode that is trained without --mixed_precision (as my models fail to train with FP16 quite consistently after about 20k iterations resulting NaN) and without guided alignments. When starting a fine tunning process with guided alignments and --mixed_precision following error occurs. Training starts just fine without --mixed_precision flag:
```
2022-11-05 07:09:26.242000: I main.py:318] Accumulate gradients of 8 iterations to reach effective batch size of 64
2022-11-05 07:09:27.756000: I dataset_ops.py:2270] Training on 1190598 examples
Traceback (most recent call last):
  File "opennmt/bin/main.py", line 361, in <module>
    main()
  File "opennmt/bin/main.py", line 318, in main
    runner.train(
  File "/ai/onmt/onmt-2280/OpenNMT-tf/opennmt/runner.py", line 281, in train
    summary = trainer(
  File "/ai/onmt/onmt-2280/OpenNMT-tf/opennmt/training.py", line 109, in __call__
    for i, loss in enumerate(
  File "/ai/onmt/onmt-2280/OpenNMT-tf/opennmt/training.py", line 248, in _steps
    loss = forward_fn()
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/util/traceback_utils.py", line 153, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/tmp/__autograph_generated_fileoze0_yzx.py", line 16, in tf___forward
    retval_ = ag__.converted_call(ag__.ld(self)._forward, (ag__.converted_call(ag__.ld(next), (ag__.ld(iterator),), None, fscope),), dict(accum_steps=ag__.ld(accum_steps), report_steps=ag__.ld(report_steps)), fscope)
  File "/tmp/__autograph_generated_file0y79rrw1.py", line 11, in tf___forward
    (loss, gradients) = ag__.converted_call(ag__.ld(self)._compute_gradients, (ag__.ld(batch), ag__.ld(accum_steps), ag__.ld(report_steps)), None, fscope)
  File "/tmp/__autograph_generated_filex26rbucc.py", line 14, in tf___compute_gradients
    (reported_loss, gradients) = ag__.converted_call(ag__.ld(self)._model.compute_gradients, (ag__.ld(features), ag__.ld(labels), ag__.ld(self)._optimizer), dict(loss_scale=(ag__.ld(accum_steps) * ag__.ld(self).num_replicas)), fscope)
  File "/tmp/__autograph_generated_file_m3r7w1w.py", line 98, in tf__compute_gradients
    ag__.if_stmt(ag__.converted_call(ag__.ld(tf).executing_eagerly, (), None, fscope), if_body_3, else_body_3, get_state_3, set_state_3, ('gradients', 'report_loss'), 2)
  File "/tmp/__autograph_generated_file_m3r7w1w.py", line 92, in else_body_3
    (train_loss, report_loss) = ag__.converted_call(ag__.ld(_compute_loss), (), None, fscope)
  File "/tmp/__autograph_generated_file_m3r7w1w.py", line 17, in _compute_loss
    (train_loss, report_loss) = ag__.converted_call(ag__.ld(self).compute_training_loss, (ag__.ld(features), ag__.ld(labels)), dict(step=ag__.ld(optimizer).iterations), fscope_1)
  File "/tmp/__autograph_generated_file7d1kbnh_.py", line 12, in tf__compute_training_loss
    loss = ag__.converted_call(ag__.ld(self).compute_loss, (ag__.ld(outputs), ag__.ld(labels)), dict(training=True), fscope)
  File "/tmp/__autograph_generated_filegf82ikml.py", line 112, in tf__compute_loss
    ag__.if_stmt(ag__.and_((lambda : (ag__.ld(noisy_logits) is not None)), (lambda : ag__.converted_call(ag__.ld(params).get, ('contrastive_learning',), None, fscope))), if_body_4, else_body_4, get_state_4, set_state_4, ('do_return', 'retval_'), 2)
  File "/tmp/__autograph_generated_filegf82ikml.py", line 100, in else_body_4
    ag__.if_stmt(ag__.ld(training), if_body_3, else_body_3, get_state_3, set_state_3, ('loss',), 1)
  File "/tmp/__autograph_generated_filegf82ikml.py", line 93, in if_body_3
    ag__.if_stmt(ag__.and_((lambda : (ag__.ld(gold_alignments) is not None)), (lambda : (ag__.ld(guided_alignment_type) is not None))), if_body_2, else_body_2, get_state_2, set_state_2, ('loss',), 1)
  File "/tmp/__autograph_generated_filegf82ikml.py", line 88, in if_body_2
    ag__.if_stmt((ag__.ld(attention) is None), if_body_1, else_body_1, get_state_1, set_state_1, ('loss',), 1)
  File "/tmp/__autograph_generated_filegf82ikml.py", line 87, in else_body_1
    loss += ag__.converted_call(losses.guided_alignment_cost, (attention[:, :(- 1)], gold_alignments), dict(sequence_length=ag__.converted_call(self.labels_inputter.get_length, (labels,), dict(ignore_special_tokens=True), fscope), cost_type=guided_alignment_type, weight=ag__.converted_call(params.get, ('guided_alignment_weight', 1), None, fscope)), fscope)
TypeError: in user code:

    File "/ai/onmt/onmt-2280/OpenNMT-tf/opennmt/training.py", line 233, in _forward  *
        next(iterator),
    File "/ai/onmt/onmt-2280/OpenNMT-tf/opennmt/training.py", line 316, in _forward  *
        loss, gradients = self._compute_gradients(
    File "/ai/onmt/onmt-2280/OpenNMT-tf/opennmt/training.py", line 298, in _compute_gradients  *
        reported_loss, gradients = self._model.compute_gradients(
    File "/ai/onmt/onmt-2280/OpenNMT-tf/opennmt/models/model.py", line 223, in _compute_loss  *
        train_loss, report_loss = self.compute_training_loss(
    File "/ai/onmt/onmt-2280/OpenNMT-tf/opennmt/models/model.py", line 264, in compute_training_loss  *
        loss = self.compute_loss(outputs, labels, training=True)
    File "/ai/onmt/onmt-2280/OpenNMT-tf/opennmt/models/sequence_to_sequence.py", line 457, in compute_loss  *
        loss += losses.guided_alignment_cost(

    TypeError: Input 'y' of 'AddV2' Op has type float16 that does not match type float32 of argument 'x'.

```


Discussion:

Resolving patch:
diff --git a/opennmt/utils/losses.py b/opennmt/utils/losses.py
--- a/opennmt/utils/losses.py
+++ b/opennmt/utils/losses.py
@@ -164,6 +164,7 @@ def guided_alignment_cost(
         sample_weight = None
         normalizer = tf.size(attention_probs)
 
+    attention_probs = tf.cast(attention_probs, tf.float32)
     cost = loss(gold_alignment, attention_probs, sample_weight=sample_weight)
     cost /= tf.cast(normalizer, cost.dtype)
     return weight * cost


Test patch:
diff --git a/opennmt/tests/model_test.py b/opennmt/tests/model_test.py
--- a/opennmt/tests/model_test.py
+++ b/opennmt/tests/model_test.py
@@ -300,6 +300,22 @@ def testSequenceToSequenceWithGuidedAlignment(self, ga_type):
         loss = model.compute_loss(outputs, labels, training=True)
         loss = loss[0] / loss[1]
 
+    @test_util.run_with_mixed_precision
+    def testSequenceToSequenceWithGuidedAlignmentMixedPrecision(self):
+        model, params = _seq2seq_model(training=True)
+        params["guided_alignment_type"] = "ce"
+        features_file, labels_file, data_config = self._makeToyEnDeData(
+            with_alignments=True
+        )
+        model.initialize(data_config, params=params)
+        model.create_variables()
+        dataset = model.examples_inputter.make_training_dataset(
+            features_file, labels_file, 16
+        )
+        features, labels = next(iter(dataset))
+        outputs, _ = model(features, labels=labels, training=True)
+        model.compute_loss(outputs, labels, training=True)
+
     def testSequenceToSequenceWithGuidedAlignmentAndWeightedDataset(self):
         model, _ = _seq2seq_model()
         features_file, labels_file, data_config = self._makeToyEnDeData(