Reproducibility issue with transformers (BERT) and tf2.2

[MFreidank]

Dear @duncanriach,
Thank you for your contributions, work and guidance towards making tensorflow deterministic in the recent releases.
Unfortunately, for popular keras NLP models (BERT) some problems seem to remain (see also related issue in this repository #14).

In spite of combining learnings from:

• the "complete recipe" in your slides from gputechconf
• your recently suggested workaround for issues with crossentropy loss

... I am still arriving at the following short, non-deterministic colab notebook example.

My results for the sum of model weights (as computed with a function you had suggested) after training for only 5 steps is (differences are highlighted below):

	Device	Before training	After training
Run 1	GPU	-641227.5609667897224	-641237.442 5159916282
Run 2	GPU	-641227.5609667897224	-641237.442 3093758523
Run 1	CPU	-641227.5609667301178	-641238.1506845243275
Run 2	CPU	-641227.5609667301178	-641238.1506845243275

This variance gets increasingly more pronounced when the model is trained for longer periods of time.

Could you please help identify the source of non-determinism and provide guidance on how we can resolve it?

As transformers is a very popular package (29.1k Github stars), I am expecting that many other people are silently impacted by this very phenomenon.

Note: As shown above, I have observed that the same code becomes fully deterministic when running on the colab CPU runtime.

[duncanriach]

Beautifully presented. Thanks, @MFreidank. I made a copy of your colab code and have been looking at it. The primary issue right now is that the trainable variables are not matching between runs:

### Before training: ###
Summary of weights: -641227.5609667897224
### Before training: ###
Summary of weights: -641227.7293046712875

I can see that you have them matching, and I don't understand that would be different for me. Have you changed the colab code in some way since you ran it?

The second issue I see is that you're setting from_logits=True in the constructor of tf.keras.losses.SparseCategoricalCrossentropy. As your notes suggest, this argument should be excluded (or set to False).

[duncanriach]

Oh, I see. You have to restart the runtime to get the same initial trainable variables. I can hopefully provide a work-around for that too.

[duncanriach]

So, the solution for getting the same initial trainable variables every time you run the block of code that starts with the definition of summarize_keras_weights is to call tf.random.set_seed at the beginning of that block. This will reset the pseudorandom number generator that is used to initialize the trainable variables of the model.

[duncanriach]

And ... solved. By removing from_logits=True from the constructor of tf.keras.losses.SparseCategoricalCrossentropy() I was able to get the same trainable variables after both runs.

### Before training: ###
Summary of weights: -641227.5609667897224
5/5 [==============================] - 7s 1s/step - loss: 0.7225 - accuracy: 0.4000
### After training: ###
Summary of weights: -641238.1517347339541

### Before training: ###
Summary of weights: -641227.5609667897224
5/5 [==============================] - 7s 1s/step - loss: 0.7225 - accuracy: 0.4000
### After training: ###
Summary of weights: -641238.1517347339541

You were so close. If you only you coded exactly what your notes required. :-)

[duncanriach]

Please confirm that your issue has been solved. Train your model for much longer, at least for one whole epoch, and confirm that it's getting the accuracy you expect while also getting the perfect, bit-exact reproducibility.

[MFreidank]

@duncanriach Thank you! I can reproduce the resolution and things are now deterministic in the scenario above - should have taken my own advice from the notes based on your workaround in the tensorflow issue thread ;)

There is an issue remaining though: changing epochs=1 to epochs=2 reintroduces non-determinism (even when keeping steps_per_epoch at only 5).
Note that training for the same 10 steps by using epochs=1, steps_per_epoch=10 is deterministic.

Could you have a look at this? I updated my colab notebook to reflect the current state and expose the issue mentioned above.

Almost looks like keras is doing some non-deterministic operations in between epochs.
For my purposes, I may be able to simply artificially stretch the epoch I am training for (to multiple passes over the dataset) and get things running deterministic this way; I'll investigate this.
Nevertheless, I believe this warrants some further investigation, happy to help in any way I can.

Update: For epochs=2, steps_per_epoch=10, I found it to be reproducible on the CPU.
So the issue must occur on something that does relate to GPU.

[duncanriach]

Could you have a look at this?

Will do.

Almost looks like Keras is doing some non-deterministic operations in between epochs.

These between-epoch issues are common and there as several possible sources. Let's see if we can get determinism without you needing to limit the training to one epoch ...

[duncanriach]

Running in colab, with my old copy of your code (with the fixes), I'm now no longer seeing reproducibility on 5 steps in one epoch on GPU. This is very concerning and I have not yet figured out what the issue is. Also, looking at your updated colab code and notes, it seems that one epoch with 10 steps on the GPU is not operating reproducibly, which does not match what you wrote above.

[duncanriach]

Just to recap where we're at and the solutions we have:

1. Using tf.random.set_seed, reset TensorFlow's PRNG before initializing trainable variables.
2. Set TF_DETERMINISTIC_OPS=1 to enable all deterministic ops in the model.
3. Replace non-deterministic fused softmax/cross-entropy with a deterministic version. You and I are also working on adding a fix for this, which will also be under the control of TF_DETERMINISTIC_OPS.

With these three adjustments, there is still some non-determinism. However, rather than just being totally different on every run, the final state of the trainable variables is now one of a discrete set of values. The number of possible values seems to increase with the number of steps per epoch.

With steps_per_epoch=1 and steps_per_epoch=2, I got the same final value after several runs.

With steps_per_epoch=5, over seven runs, I got only four different results. One result was repeated three times and another was repeated twice.

What this suggests to me is that there may be some non-determinism in the interaction between the data-loader (based on tf.data) and model.fit. I've seen things like this before, but not exactly like this, and nothing jumps out at me from your code that could be causing this (such a multiple data-loader workers or an unseeded shuffle).

I'll investigate more tomorrow.

[MFreidank]

@duncanriach Thank you a lot for your work and drive on this and for the conclusive summary of where we stand and what we know.
I agree with all your points, but was not able to pinpoint the exact source of the problem (I tried setting workers=0 to make it run on the same thread as the main training loop, but to no avail).
Looking forward to your further investigation and happy to help from my side in any way I can.