Repository: rwth-i6/returnn

Issue:
NAdam different between TF1.15 and TF2.3
I am not posting the full information here yet, because I need to collect some more things, but this is more a starting point for this issue:

Summary of important facts:
 - Nearly lastest RETURNN (I started the tests at 9th November, so RETURNN is from 9th November)
 - Both versions run with CUDA 10.1 on the same machine(s)
 - 6 Layer BLSTM (nativelstm2) without any extras
 - No specaugment, only l2 and dropout
 - CTC loss from RASR process, 139 output labels
 - NAdam optimizer
 - Training is starting from a fixed initial state

What happens:
 - With TF1.15, the model starts to converge significantly around 500 steps
 - With TF2.3, no convergence happens in the first 3000 steps
 - With TF1.15 the LSTM gradients are low after a peak, but go upwards shortly after
 - With TF2.3, the LSTM gradients stay low
 - With TF1.15, the l2 constraint increases
 - With TF2.3, the l2 constraint decreases

I did many repetitions of this experiment with different learning rates, different batch sizes, tried step-based warmup, but it always resulted in the same behavior.

TF1:
![TF1_grads](https://user-images.githubusercontent.com/10742135/143285232-29166452-0629-4178-9092-34c53b0a3c55.png)
![TF1_loss](https://user-images.githubusercontent.com/10742135/143285234-7d35f8b3-99ff-4afc-ba56-2b93611f1788.png)

TF2:
![TF2_grads](https://user-images.githubusercontent.com/10742135/143285235-c1f1c24e-003c-44f1-9fef-1e55b496af58.png)
![TF2_loss](https://user-images.githubusercontent.com/10742135/143285237-8845d2b7-6660-4b80-8a9e-0d3b91cb0a39.png)





Discussion:
The gradient norms you mean?

When you make sure to start from the same model (e.g. via `task="initialize_model"`), you could verify whether it is really exactly the same or not (at least in the first batch).

> The gradient norms you mean?

If the norm is different the gradient is different. I just displayed the norm here, of course I looked at all values.

> When you make sure to start from the same model (e.g. via `task="initialize_model"`), you could verify whether it is really exactly the same or not (at least in the first batch).

Ah sorry, I did not mention this. Yes this is the case, I start from the same initialized model. But the reason was rather to exclude the possibility that the initialization changed between TF1.15 and TF2.3 for any reason.

But is the gradient still the same (within some threshold) for the first mini batch? I.e. the difference accumulates slowly over time? Or is there some specific later step where it is clearly different, and from there on it becomes different? Or is it already different in the first mini batch? And where exactly is it different? Already the error signal from the loss? Or at what layer does it become different in backprop?

> But is the gradient still the same (within some threshold) for the first mini batch? I.e. the difference accumulates slowly over time? Or is there some specific later step where it is clearly different, and from there on it becomes different? Or is it already different in the first mini batch? And where exactly is it different? Already the error signal from the loss? Or at what layer does it become different in backprop?

The training logs show that the error signal is identical (first minibatch shows the same loss, so before the first gradient it is identical), and for the second minibatch there is already a noticeable difference.

I will try to come up with a toy task (artificial data, small network) that reproduces this.

For the last question I would need to have produce shorter logs, you can not see the raw values properly in the Tensorboard as it is now, so I can not see if for the first step the gradient is only different in e.g. the last LSTM layer but not for the linear transformation before the Softmax.
The same loss does not mean that the gradients are identical. Did you directly check the gradients (not just the norm)?
I would not just look at TensorBoard. I would write a small script which dumps the gradients to some file, maybe after each mini batch, maybe also other information, and then investigate that by hand.
> The same loss does not mean that the gradients are identical. Did you directly check the gradients (not just the norm)?

It does not, and I did not say that. I say the first entry (first minibatch before any update) shows identical loss, which means that the loss computation is not flawed.  

The second entry shows diverging loss, so the first update was already different. For which location was different, I would indeed need to dump the gradients. 
How much different is the loss in the second step? Or how much do the gradients really differ in the first step?

I would assume that the first gradients are still all very similar (up to some threshold) and the differences just accumulate over time.

This is then maybe not even a bug. Just different behavior in different TF versions, and you have bad luck here that this different behavior leads to convergence in one case but not in the other.

Although, as this is now a problem for you, you need to investigate more about what exactly is different in the behavior, and better understand it. Maybe replicate the old behavior somehow, or find other ways.
If the first gradients are already different, that's good. Because that means that it will be easy to debug. Then you should check at what point it becomes different. E.g. already at the loss gradient, or at some layer?

>  E.g. already at the loss gradient, or at some layer?

I did not extract the gradients yet, but I opened the raw values of the summary file. There you can clearly see that in the first update step the gradients for the softmax bias, which directly depend on the error, are already off.

TF23:
```
tag: "grads/optimize/gradients/output/linear/add_bias_grad/tuple/summaries_grad_of_b/grad_of_b_mean"
simple_value: -0.0031194891780614853

tag: "grads/optimize/gradients/output/linear/add_bias_grad/tuple/summaries_grad_of_b/grad_of_b_stddev"
simple_value: 107.59998321533203

tag: "grads/optimize/gradients/output/linear/add_bias_grad/tuple/summaries_grad_of_b/grad_of_b_rms"
simple_value: 107.59998321533203

tag: "grads/optimize/gradients/output/linear/add_bias_grad/tuple/summaries_grad_of_b/grad_of_b_l2"
simple_value: 634.2925415039062

tag: "grads/optimize/gradients/output/linear/add_bias_grad/tuple/summaries_grad_of_b/grad_of_b_max"
simple_value: 22.030406951904297

tag: "grads/optimize/gradients/output/linear/add_bias_grad/tuple/summaries_grad_of_b/grad_of_b_min"
simple_value: -1243.468017578125
```

TF115
```
tag: "grads/optimize/gradients/output/linear/add_bias_grad/tuple/summaries_grad_of_b/grad_of_b_mean"
simple_value: -0.0035799972247332335

tag: "grads/optimize/gradients/output/linear/add_bias_grad/tuple/summaries_grad_of_b/grad_of_b_stddev"
simple_value: 106.58272552490234

tag: "grads/optimize/gradients/output/linear/add_bias_grad/tuple/summaries_grad_of_b/grad_of_b_rms"
simple_value: 106.58272552490234

tag: "grads/optimize/gradients/output/linear/add_bias_grad/tuple/summaries_grad_of_b/grad_of_b_l2"
simple_value: 628.2958984375

tag: "grads/optimize/gradients/output/linear/add_bias_grad/tuple/summaries_grad_of_b/grad_of_b_max"
simple_value: 22.547555923461914

tag: "grads/optimize/gradients/output/linear/add_bias_grad/tuple/summaries_grad_of_b/grad_of_b_min"
simple_value: -1227.8902587890625
```

I assume the gradient statistics for Adam are differently initialized in TF1.15 and TF2.3

EDIT: I already checked the default parameters, but they did not change except for the epsilon, which is given fixed by the RETURNN config.
Ah interesting.

Can you disable Adam to see that Adam or some Adam specific things are not causing this? Just use standard SGD.

I'm not sure if this is maybe still within numerical fluctuations...

You can easily directly check the gradient of the loss then (w.r.t. the logits).

Maybe TF accumulates them differently between TF 1.15 and TF 2.3. For the bias, it will accumulate them over time (hopefully correctly discarding the padded frames as well).

Okay, I did something wrong while managing my debug configs. The error is most likely that the NAdam implementation changed between TF1.15 and TF2.3
So yes, the model is converging nicely when switching to Adam also with TF2.3. I already saw the NAdam discrepancy yesterday but thought I switched to Adam to exclude this mismatch.

The difference is the following implementation difference:
Native code in TF1.15 (`tensorflow/tensorflow/core/kernels/training_ops_gpu.cu.cc`)
```cpp
template <typename T>
__global__ void ApplyAdamKernel(int32 data_dim, T* var, T* m, T* v,
                                const T* const beta1_power_,
                                const T* const beta2_power_, const T* const lr_,
                                const T* const beta1_, const T* const beta2_,
                                const T* const epsilon_, const T* grad,
                                bool use_nesterov) {
  eigen_assert(blockDim.y == 1);
  eigen_assert(blockDim.z == 1);
  eigen_assert(gridDim.y == 1);
  eigen_assert(gridDim.z == 1);

  const T mul_factor = (*lr_) * sqrt(static_cast<T>(1.0) - (*beta2_power_)) /
                       (static_cast<T>(1.0) - (*beta1_power_));
  const T epsilon = (*epsilon_);
  const T beta1 = (*beta1_);
  const T one_minus_beta1 = static_cast<T>(1.0) - (beta1);
  const T one_minus_beta2 = static_cast<T>(1.0) - (*beta2_);
  const int32 stripe = gridDim.x * blockDim.x;

  for (int32 i = blockIdx.x * blockDim.x + threadIdx.x; i < data_dim;
       i += stripe) {
    auto m_i = m[i];
    auto g_i = grad[i];
    auto v_i = v[i];

    m_i += one_minus_beta1 * (g_i - m_i);
    v_i += one_minus_beta2 * (g_i * g_i - v_i);
    if (use_nesterov) {
      var[i] -= mul_factor * (m_i * beta1 + one_minus_beta1 * g_i) /
                (epsilon + sqrt(v_i));
    } else {
      var[i] -= mul_factor * m_i / (epsilon + sqrt(v_i));
    }

    m[i] = m_i;
    v[i] = v_i;
  }
}
```
Code in TF2.3 / Keras (`keras/optimizer_v2/nadam.py`)
```python
  def _resource_apply_dense(self, grad, var, apply_state=None):
    var_device, var_dtype = var.device, var.dtype.base_dtype
    coefficients = ((apply_state or {}).get((var_device, var_dtype))
                    or self._fallback_apply_state(var_device, var_dtype))

    m = self.get_slot(var, 'm')
    v = self.get_slot(var, 'v')

    g_prime = grad / coefficients['one_minus_m_schedule_new']
    m_t = (coefficients['beta_1_t'] * m +
           coefficients['one_minus_beta_1_t'] * grad)
    m_t = tf.compat.v1.assign(m, m_t, use_locking=self._use_locking)
    m_t_prime = m_t / coefficients['one_minus_m_schedule_next']
    v_t = (coefficients['beta_2_t'] * v +
           coefficients['one_minus_beta_2_t'] * tf.square(grad))
    v_t = tf.compat.v1.assign(v, v_t, use_locking=self._use_locking)
    v_t_prime = v_t / coefficients['v_t_prime_denominator']
    m_t_bar = (coefficients['one_minus_m_t'] * g_prime +
               coefficients['m_t_1'] * m_t_prime)
    var_t = var - coefficients['lr_t'] * m_t_bar / (
        tf.sqrt(v_t_prime) + coefficients['epsilon'])
    return tf.compat.v1.assign(var, var_t, use_locking=self._use_locking).op
```

The differences are:
 - different computation of the LR/multiplication scale (`mul_factor`)
 - numerically different computation of `m_t` /  `m_i` and `v_t`/  `v_i`

I reported the problem for TF here: https://github.com/tensorflow/tensorflow/issues/53204

Resolving patch:
diff --git a/returnn/tf/updater.py b/returnn/tf/updater.py
--- a/returnn/tf/updater.py
+++ b/returnn/tf/updater.py
@@ -31,7 +31,12 @@ def _init_optimizer_classes_dict():
   if _OptimizerClassesDictInitialized:
     return
   _OptimizerClassesDictInitialized = True
-  potential_list = list(vars(tf_compat.v1.train).items())
+  # Build up all potential candidates of such classes.
+  # We will filter it below.
+  # The order of the list matters, and the first optimizer is used.
+  # E.g. our own NadamOptimizer here will be preferred over others.
+  potential_list = list(globals().items())
+  potential_list += list(vars(tf_compat.v1.train).items())
   if tf_version_tuple() >= (1, 2, 0):
     try:
       from tensorflow.contrib import opt  # noqa
@@ -42,7 +47,8 @@ def _init_optimizer_classes_dict():
   if KerasOptimizer:
     potential_list += list(vars(tf_compat.v2.keras.optimizers).items())
     allowed_types += (KerasOptimizer,)
-  potential_list += list(globals().items())
+    # Special name for Nadam in Keras, as we provide our own Nadam here.
+    potential_list += [("NadamKeras", tf_compat.v2.keras.optimizers.Nadam)]
   for name, v in potential_list:
     assert isinstance(name, str)
     # We might have duplicate names if TF v1 and v2 are mixed, etc.
@@ -1258,6 +1264,96 @@ def _apply(self, grad, var, indices=None):
     return tf.group(*[var_update, m_t])
 
 
+class NadamOptimizer(tf_compat.v1.train.AdamOptimizer):
+  """
+  Optimizer that implements the Nadam algorithm.
+  See [Dozat, T., 2015](http://cs229.stanford.edu/proj2015/054_report.pdf).
+
+  Copied from:
+  https://github.com/tensorflow/tensorflow/blob/v1.15.5/tensorflow/contrib/opt/python/training/nadam_optimizer.py
+
+  We have this here to have this Nadam variant available in TF 2
+  because the Keras Nadam behaves a bit different.
+  https://github.com/rwth-i6/returnn/issues/766
+  https://github.com/tensorflow/tensorflow/issues/53204
+
+  We can still use this old code because the underlying kernel still supports the ``use_nesterov`` option.
+  """
+
+  def _apply_dense(self, grad, var):
+    from tensorflow.python.training import training_ops
+    from tensorflow.python.ops import math_ops
+    m = self.get_slot(var, "m")
+    v = self.get_slot(var, "v")
+    beta1_power, beta2_power = self._get_beta_accumulators()
+    return training_ops.apply_adam(
+      var,
+      m,
+      v,
+      math_ops.cast(beta1_power, var.dtype.base_dtype),
+      math_ops.cast(beta2_power, var.dtype.base_dtype),
+      math_ops.cast(self._lr_t, var.dtype.base_dtype),
+      math_ops.cast(self._beta1_t, var.dtype.base_dtype),
+      math_ops.cast(self._beta2_t, var.dtype.base_dtype),
+      math_ops.cast(self._epsilon_t, var.dtype.base_dtype),
+      grad,
+      use_locking=self._use_locking,
+      use_nesterov=True).op
+
+  def _resource_apply_dense(self, grad, var):
+    from tensorflow.python.training import training_ops
+    from tensorflow.python.ops import math_ops
+    m = self.get_slot(var, "m")
+    v = self.get_slot(var, "v")
+    beta1_power, beta2_power = self._get_beta_accumulators()
+    return training_ops.resource_apply_adam(
+      var.handle,
+      m.handle,
+      v.handle,
+      math_ops.cast(beta1_power, grad.dtype.base_dtype),
+      math_ops.cast(beta2_power, grad.dtype.base_dtype),
+      math_ops.cast(self._lr_t, grad.dtype.base_dtype),
+      math_ops.cast(self._beta1_t, grad.dtype.base_dtype),
+      math_ops.cast(self._beta2_t, grad.dtype.base_dtype),
+      math_ops.cast(self._epsilon_t, grad.dtype.base_dtype),
+      grad,
+      use_locking=self._use_locking,
+      use_nesterov=True)
+
+  def _apply_sparse_shared(self, grad, var, indices, scatter_add):
+    from tensorflow.python.ops import math_ops
+    from tensorflow.python.ops import state_ops
+    from tensorflow.python.ops import array_ops
+    from tensorflow.python.framework import ops
+    from tensorflow.python.ops import control_flow_ops
+    beta1_power, beta2_power = self._get_beta_accumulators()
+    beta1_power = math_ops.cast(beta1_power, var.dtype.base_dtype)
+    beta2_power = math_ops.cast(beta2_power, var.dtype.base_dtype)
+    lr_t = math_ops.cast(self._lr_t, var.dtype.base_dtype)
+    beta1_t = math_ops.cast(self._beta1_t, var.dtype.base_dtype)
+    beta2_t = math_ops.cast(self._beta2_t, var.dtype.base_dtype)
+    epsilon_t = math_ops.cast(self._epsilon_t, var.dtype.base_dtype)
+    lr = (lr_t * math_ops.sqrt(1 - beta2_power) / (1 - beta1_power))
+    # m_t = beta1 * m + (1 - beta1) * g_t
+    m = self.get_slot(var, "m")
+    m_scaled_g_values = grad * (1 - beta1_t)
+    m_t = state_ops.assign(m, m * beta1_t, use_locking=self._use_locking)
+    with ops.control_dependencies([m_t]):
+      m_t = scatter_add(m, indices, m_scaled_g_values)
+      # m_bar = (1 - beta1) * g_t + beta1 * m_t
+      m_bar = m_scaled_g_values + beta1_t * array_ops.gather(m_t, indices)
+    # v_t = beta2 * v + (1 - beta2) * (g_t * g_t)
+    v = self.get_slot(var, "v")
+    v_scaled_g_values = (grad * grad) * (1 - beta2_t)
+    v_t = state_ops.assign(v, v * beta2_t, use_locking=self._use_locking)
+    with ops.control_dependencies([v_t]):
+      v_t = scatter_add(v, indices, v_scaled_g_values)
+    v_t_slice = array_ops.gather(v_t, indices)
+    v_sqrt = math_ops.sqrt(v_t_slice)
+    var_update = scatter_add(var, indices, -lr * m_bar / (v_sqrt + epsilon_t))
+    return control_flow_ops.group(*[var_update, m_bar, v_t])
+
+
 class CustomAdamOptimizer(BaseCustomOptimizer):
   """
   Reimplementation of Adam.


Test patch:
diff --git a/tests/test_TFEngine.py b/tests/test_TFEngine.py
--- a/tests/test_TFEngine.py
+++ b/tests/test_TFEngine.py
@@ -256,13 +256,21 @@ def test_engine_train_optimizer_class():
   test_engine_train(additional_config)
 
 
-def test_engine_train_keras_optimizer():
+def test_engine_train_nadam_optimizer():
   additional_config = {
     "optimizer": {"class": "nadam"},
   }
   test_engine_train(additional_config)
 
 
+@unittest.skipIf(tf.__version__.startswith("1."), "TF 1")
+def test_engine_train_keras_optimizer():
+  additional_config = {
+    "optimizer": {"class": "NadamKeras"},
+  }
+  test_engine_train(additional_config)
+
+
 def test_engine_train_new_dataset_pipeline():
   from returnn.datasets.generating import DummyDataset
   seq_len = 5