Fine tuning TensorFlow DeBERTa fails on TPU

[tmoroder]

System Info

Latest version of transformers, Colab TPU, tensorflow 2.

• Colab TPU
• transformers: 4.21.0
• tensorflow: 2.8.2 / 2.6.2
• Python 3.7

Who can help?

@LysandreJik, @Rocketknight1, @san

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

I am facing some issues while trying to fine-tune a TensorFlow DeBERTa model microsoft/deberta-v3-base on TPU.

I have created some Colab notebooks showing the errors. Note, the second and third notebooks already include some measures to circumvent previous errors.

• ValueError with partially known TensorShape with latest take_along_axis change: FineTuning_TF_DeBERTa_TPU_1
• Output shape mismatch of branches with custom dropout: FineTuning_TF_DeBERTa_TPU_2
• XLA compilation error because of dynamic/computed tensor shapes: FineTuning_TF_DeBERTa_TPU_3

I have seen similar issues when using microsoft/deberta-base.

I believe the following issues are related:

• TF2 DeBERTaV2 runs super slow on TPUs #18239
• Debertav2 debertav3 TPU : socket closed #18276. From this I used the fix on take_along_axis.

Thanks!

Expected behavior

Fine tuning is possible as it happens when using a GPU.

[gante]

Hi @tmoroder 👋 Thank you for adding all that information to the issue <3

If I got it right, the second notebook replaces the take_along_axis function, and the third notebook also replaces the custom dropout. Still, there are XLA exceptions.

Before diving into debugging, two questions:

1. Does it return the same error on a GPU?
2. I see that you prepare a dataset with static batch size and that the input is padded. Do you think that there is any additional source of shape variability in the inputs? (I don't think so, but asking doesn't hurt :D )

[tmoroder]

Hi @gante.

If I got it right, the second notebook replaces the take_along_axis function, and the third notebook also replaces the custom dropout. Still, there are XLA exceptions.

Correct. I think the XLA exceptions occur during gradient computation at these dynamic/computed tensor shape sizes.
The first collection seems to me being triggered within the TFDebertaV2DisentangledSelfAttention.disentangled_att_bias method, like at L735. I am not about other position TFDebertaV2DisentangledSelfAttention.call like L704.

1. Does it return the same error on a GPU?

It runs on GPU without errors if I use transformers==4.20.1, see FineTuning_TF_DeBERTa _GPU. With version 4.21.0 I get the same error ValueError.

2. I see that you prepare a dataset with static batch size and that the input is padded. Do you think that there is any additional source of shape variability in the inputs? (I don't think so, but asking doesn't hurt :D )

No further shape variability as fas as I can judge.

[Rocketknight1]

Hi @tmoroder, can you try on GPU with jit_compile=True in both 4.20 and 4.21? I believe the code had issues with XLA before 4.21, and TPU code is always compiled to XLA.

[gante]

Interesting. Since transformers==4.20.1, there are only two DeBERTa PRs:

1. fix bias keyword argument in TFDebertaEmbeddings #17940 (should have no impact at all here)
2. Change how take_along_axis is computed in DeBERTa to stop confusing XLA #18256 (what should have been a TPU-friendly take_along_axis)

As @Rocketknight1 said, that data would be interesting. If v4.21 works on GPU but not on TPU, we are up for an interesting challenge :D

[tmoroder]

Hi @tmoroder, can you try on GPU with jit_compile=True in both 4.20 and 4.21?

Using jit_compile=True while compiling the model gives an error for both 4.20.1 and 4.21, e.g., FineTuning_TF_DeBERTa _GPU_Tests for 4.21; with 4.20.1 it crashes in the last command.

As @Rocketknight1 said, that data would be interesting. If v4.21 works on GPU but not on TPU, we are up for an interesting challenge :D

Without jit_compile=True it also fails on GPU with 4.21; with 4.20.1 it works.

[Rocketknight1]

That makes sense - we made changes to the model to make it XLA-compatible in 4.21. XLA compatibility is necessary for TPU support, so the 4.20 model would never have run on TPU. However, we seem to have some other TPU-specific issues now - in my testing I was able to get DeBERTa to work with XLA on GPU in 4.21.

[WissamAntoun]

Weird!

During my TPU and GPU tests, i was using a custom training loop instead of keras's .fit(), which I'm not sure if it actually matters.

In my custom training code, I got deberta to train in an electra style training, with XLA enabled with jit_compile=True with non of the issues mentioned above.

I will be sharing my code asap once I finish the pretraining and validate the results. It is based on Nvidia BERT and Electra TF2 training code https://github.com/NVIDIA/DeepLearningExamples/tree/master/TensorFlow2/LanguageModeling/

[gante]

@tmoroder I can confirm that I can run your example with jit_compile=True (i.e. XLA compilation) on model.compile(), using a GPU, if the two changes you made in your third TPU notebook:

• replace take_along_axis by the tf.einsum version
• replace dropout by the standard dropout

If XLA compilation works, then it should run on TPU. I noticed that in your notebook you were using TF 2.6, which may explain the XLA failure. Are you able to bump your TPU TF version (to as high as possible)?

Meanwhile I'm opening a PR to reflect those two changes :)

[tmoroder]

@gante

Thanks a lot for your effort. Maybe I am doing something wrong... but using the code from your pull request it now runs on GPU (with jit_compile=True as additional argument during model compilation), while it still fails on TPU (without using jit_compile=True as an argument). I am using TF 2.8.2 in both cases which is the current default in the Colab environment. On TPU it seems again to have errors on the tile operation.

• Working GPU version: FineTuning_TF_DeBERTa_Propsed_Fix_GPU
• Failing TPU version: FineTuning_TF_DeBERTa_Propsed_Fix_TPU

[gante]

(linking issues -- the Tile issue is also present in the following unsolved issue: #14058)

[gante]

The cause is trivial (the multiple argument of tf.tile can't have dynamic shapes), but the fix may be not :D Will look into it

[gante]

@tmoroder the dynamic shape in question is the batch size. I may be able to get an alternative to tf.tile, but I highly doubt that it will make a difference -- it will be a dynamic shape no matter how I turn it around, as it is not set.

Alternatively, could you try setting the batch_size argument in the Input layers? It will make that shape static, which should be enough to squash the problem you're seeing :)

[tmoroder]

@gante

Great, setting the batch_size works 🥳. I only had to make sure that it divides the strategy.num_replicas_in_sync, FineTuning_TF_DeBERTa_Working_Fix_TPU. Thanks a lot, I will test the procedure now on my real use case at hand.