TF2 DeBERTaV2 runs super slow on TPUs

[WissamAntoun]

System Info

latest version of transformers, Colab TPU, tensorflow 2

Who can help?

@kamalkraj @Rocketknight1 @BigBird01

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

It's currently hard to share code and access to the google bucket. But I believe any TF2 DeBERTaV2 code running on TPUs will have this issue

Expected behavior

I've been trying to train a deberta v3 model on GPU and TPUs. I got it to work on multi-node and multi-gpus using Nvidia deeplearning examples libraries https://github.com/NVIDIA/DeepLearningExamples/blob/master/TensorFlow2/LanguageModeling/
I basically used the training setup and loop from the BERT code, the dataset utils from the ELECTRA code, and the model from Huggingface transformers with some changes in order to share embeddings.

On 6xA40 45gb gpus i get around 1370 sentences per seconds during training (which is lower than what Nvidia gets for Electra but it's fine).

Ok, now the problem.... on TPU i get 20 sentences per second

I traced the issue back to the tf.gather function here https://github.com/huggingface/transformers/blob/main/src/transformers/models/deberta_v2/modeling_tf_deberta_v2.py#L525

I ran TPU profiling and this is the output:

GatherV2 takes most of the time:

zoomed in pictures of the fast ops

Also, I'm not sure if this is TPU specific since on GPUs the training ~30% slower compared to regular ELECTRA.

[Rocketknight1]

Hi @WissamAntoun, this is an interesting issue! I honestly have no idea what the cause could be, but the fact that it highlights that function is interesting. The reason is that the DeBERTa code was ported from PyTorch, and so we wrote our own implementation of take_along_axis because TF didn't have one. One thing to try would be to edit the code to use tf.experimental.numpy.take_along_axis instead of that function. If that doesn't work then we might have to see if we can do things in a different, more performant way.

Also, just in case XLA compilation is the issue, have you tried using jit_compile=True in compile() when running DeBERTa on GPU? If that also causes performance degradation then the problem is caused by XLA and not TPUs, and we can investigate from there.

[Rocketknight1]

Also cc @sanchit-gandhi because I'm not a TPU expert - don't worry about investigating this deeply, but if anything comes to mind when you read it, let me know!

[WissamAntoun]

@Rocketknight1 I read all the discussions that you had with Kamal about the torch.gather and take_along_axis .

On GPUs I already enabled XLA via tf.config.optimizer.set_jit and via TF_XLA_FLAGS="--tf_xla_auto_jit=2 --tf_xla_cpu_global_jit" but I was reading that this isn't the optimal way to do it, so I'm now trying the jit_compile=True and will report back.

Also I just finished testing tf.experimental.numpy.take_along_axis, on GPUs it improved performance by ~10% yet on TPUs I still have the same issue. I will also test the jit_compile on TPUs but I don't think it will solve anything.

Thanks a lot for the replies and for the effort you put in convert the pytorch code into TF

[WissamAntoun]

runnig the training with jit_compile=True on GPU revealed a new bug. Then it is now an XLA/JIT issue not a TPU one

View log dump

2022-07-21 23:36:18.107830: W tensorflow/core/framework/op_kernel.cc:1745] OP_REQUIRES failed at bcast_ops.cc:50 : 
INVALID_ARGUMENT: 
Input 0 to node `pretraining_model/tf_deberta_v2_for_masked_lm/deberta/encoder/layer_._0/attention/self/BroadcastArgs`
with op BroadcastArgs must be a compile-time constant.

XLA compilation requires that operator arguments that represent shapes or dimensions be evaluated to concrete values at compile time. 
This error means that a shape or dimension argument could not be evaluated at compile time, 
usually because the value of the argument depends on a parameter to the computation, on a variable, 
or on a stateful operation such as a random number generator.

Stack trace for op definition: 
File "run_pretraining.py", line 204, in <module>
  config = main(start_time)
File "run_pretraining.py", line 184, in main
  trained_model = run_customized_training_loop(
File "/workspaces/nv-deberta-tf2/electra/model_training_utils.py", line 675, in run_customized_training_loop
  train_steps_strategy(
File "/workspaces/nv-deberta-tf2/electra/model_training_utils.py", line 407, in train_steps_strategy
  if num_grad_accumulates != 1:
File "/workspaces/nv-deberta-tf2/electra/model_training_utils.py", line 408, in train_steps_strategy
  for step_idx in tf.range(steps * num_grad_accumulates):
File "/workspaces/nv-deberta-tf2/electra/model_training_utils.py", line 410, in train_steps_strategy
  strategy.run(_forward, args=(next(iterator),))
File "/workspaces/nv-deberta-tf2/electra/model_training_utils.py", line 324, in _forward
  loss, model_outputs = model(inputs, is_training=True)
File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 64, in error_handler
  return fn(*args, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/keras/engine/base_layer.py", line 1096, in __call__
  outputs = call_fn(inputs, *args, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 92, in error_handler
  return fn(*args, **kwargs)
File "/workspaces/nv-deberta-tf2/electra/modeling_tf_deberta_v2.py", line 2491, in call
  if config.uniform_generator:
File "/workspaces/nv-deberta-tf2/electra/modeling_tf_deberta_v2.py", line 2496, in call
  mlm_output = self._get_masked_lm_output(
File "/workspaces/nv-deberta-tf2/electra/modeling_tf_deberta_v2.py", line 2541, in _get_masked_lm_output
  if self._config.uniform_generator:
File "/workspaces/nv-deberta-tf2/electra/modeling_tf_deberta_v2.py", line 2550, in _get_masked_lm_output
  outputs = generator(
File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 64, in error_handler
  return fn(*args, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/keras/engine/base_layer.py", line 1096, in __call__
  outputs = call_fn(inputs, *args, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 92, in error_handler
  return fn(*args, **kwargs)
File "/workspaces/nv-deberta-tf2/electra/modeling_tf_utils.py", line 1872, in run_call_with_unpacked_inputs
  )
File "/workspaces/nv-deberta-tf2/electra/modeling_tf_deberta_v2.py", line 1880, in call
  outputs = self.deberta(
File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 64, in error_handler
  return fn(*args, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/keras/engine/base_layer.py", line 1096, in __call__
  outputs = call_fn(inputs, *args, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 92, in error_handler
  return fn(*args, **kwargs)
File "/workspaces/nv-deberta-tf2/electra/modeling_tf_utils.py", line 1872, in run_call_with_unpacked_inputs
  )
File "/workspaces/nv-deberta-tf2/electra/modeling_tf_deberta_v2.py", line 1617, in call
  encoder_outputs = self.encoder(
File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 64, in error_handler
  return fn(*args, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/keras/engine/base_layer.py", line 1096, in __call__
  outputs = call_fn(inputs, *args, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 92, in error_handler
  return fn(*args, **kwargs)
File "/workspaces/nv-deberta-tf2/electra/modeling_tf_deberta_v2.py", line 527, in call
  for i, layer_module in enumerate(self.layer):
File "/workspaces/nv-deberta-tf2/electra/modeling_tf_deberta_v2.py", line 532, in call
  layer_outputs = layer_module(
File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 64, in error_handler
  return fn(*args, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/keras/engine/base_layer.py", line 1096, in __call__
  outputs = call_fn(inputs, *args, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 92, in error_handler
  return fn(*args, **kwargs)
File "/workspaces/nv-deberta-tf2/electra/modeling_tf_deberta_v2.py", line 317, in call
  attention_outputs = self.attention(
File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 64, in error_handler
  return fn(*args, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/keras/engine/base_layer.py", line 1096, in __call__
  outputs = call_fn(inputs, *args, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 92, in error_handler
  return fn(*args, **kwargs)
File "/workspaces/nv-deberta-tf2/electra/modeling_tf_deberta_v2.py", line 226, in call
  self_outputs = self.self(
File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 64, in error_handler
  return fn(*args, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/keras/engine/base_layer.py", line 1096, in __call__
  outputs = call_fn(inputs, *args, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 92, in error_handler
  return fn(*args, **kwargs)
File "/workspaces/nv-deberta-tf2/electra/modeling_tf_deberta_v2.py", line 876, in call
  if self.relative_attention:
File "/workspaces/nv-deberta-tf2/electra/modeling_tf_deberta_v2.py", line 878, in call
  rel_att = self.disentangled_att_bias(
File "/workspaces/nv-deberta-tf2/electra/modeling_tf_deberta_v2.py", line 991, in disentangled_att_bias
  if "c2p" in self.pos_att_type:
File "/workspaces/nv-deberta-tf2/electra/modeling_tf_deberta_v2.py", line 1012, in disentangled_att_bias
  c2p_att = tnp.take_along_axis(

2022-07-21 23:36:18.184105: W tensorflow/core/framework/op_kernel.cc:1745] OP_REQUIRES failed at xla_ops.cc:248 : 
INVALID_ARGUMENT: 
Input 0 to node `pretraining_model/tf_deberta_v2_for_masked_lm/deberta/encoder/layer_._0/attention/self/BroadcastArgs` 
with op BroadcastArgs must be a compile-time constant.

XLA compilation requires that operator arguments that represent shapes or dimensions be evaluated to concrete values at compile time. 
This error means that a shape or dimension argument could not be evaluated at compile time, 
usually because the value of the argument depends on a parameter to the computation, 
on a variable, or on a stateful operation such as a random number generator.

Stack trace for op definition: 
File "run_pretraining.py", line 204, in <module>
  config = main(start_time)
File "run_pretraining.py", line 184, in main
  trained_model = run_customized_training_loop(
File "/workspaces/nv-deberta-tf2/electra/model_training_utils.py", line 675, in run_customized_training_loop
  train_steps_strategy(
File "/workspaces/nv-deberta-tf2/electra/model_training_utils.py", line 407, in train_steps_strategy
  if num_grad_accumulates != 1:
File "/workspaces/nv-deberta-tf2/electra/model_training_utils.py", line 408, in train_steps_strategy
  for step_idx in tf.range(steps * num_grad_accumulates):
File "/workspaces/nv-deberta-tf2/electra/model_training_utils.py", line 410, in train_steps_strategy
  strategy.run(_forward, args=(next(iterator),))
File "/workspaces/nv-deberta-tf2/electra/model_training_utils.py", line 324, in _forward
  loss, model_outputs = model(inputs, is_training=True)
File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 64, in error_handler
  return fn(*args, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/keras/engine/base_layer.py", line 1096, in __call__
  outputs = call_fn(inputs, *args, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 92, in error_handler
  return fn(*args, **kwargs)
File "/workspaces/nv-deberta-tf2/electra/modeling_tf_deberta_v2.py", line 2491, in call
  if config.uniform_generator:
File "/workspaces/nv-deberta-tf2/electra/modeling_tf_deberta_v2.py", line 2496, in call
  mlm_output = self._get_masked_lm_output(
File "/workspaces/nv-deberta-tf2/electra/modeling_tf_deberta_v2.py", line 2541, in _get_masked_lm_output
  if self._config.uniform_generator:
File "/workspaces/nv-deberta-tf2/electra/modeling_tf_deberta_v2.py", line 2550, in _get_masked_lm_output
  outputs = generator(
File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 64, in error_handler
  return fn(*args, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/keras/engine/base_layer.py", line 1096, in __call__
  outputs = call_fn(inputs, *args, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 92, in error_handler
  return fn(*args, **kwargs)
File "/workspaces/nv-deberta-tf2/electra/modeling_tf_utils.py", line 1872, in run_call_with_unpacked_inputs
  )
File "/workspaces/nv-deberta-tf2/electra/modeling_tf_deberta_v2.py", line 1880, in call
  outputs = self.deberta(
File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 64, in error_handler
  return fn(*args, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/keras/engine/base_layer.py", line 1096, in __call__
  outputs = call_fn(inputs, *args, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 92, in error_handler
  return fn(*args, **kwargs)
File "/workspaces/nv-deberta-tf2/electra/modeling_tf_utils.py", line 1872, in run_call_with_unpacked_inputs
  )
File "/workspaces/nv-deberta-tf2/electra/modeling_tf_deberta_v2.py", line 1617, in call
  encoder_outputs = self.encoder(
File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 64, in error_handler
  return fn(*args, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/keras/engine/base_layer.py", line 1096, in __call__
  outputs = call_fn(inputs, *args, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 92, in error_handler
  return fn(*args, **kwargs)
File "/workspaces/nv-deberta-tf2/electra/modeling_tf_deberta_v2.py", line 527, in call
  for i, layer_module in enumerate(self.layer):
File "/workspaces/nv-deberta-tf2/electra/modeling_tf_deberta_v2.py", line 532, in call
  layer_outputs = layer_module(
File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 64, in error_handler
  return fn(*args, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/keras/engine/base_layer.py", line 1096, in __call__
  outputs = call_fn(inputs, *args, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 92, in error_handler
  return fn(*args, **kwargs)
File "/workspaces/nv-deberta-tf2/electra/modeling_tf_deberta_v2.py", line 317, in call
  attention_outputs = self.attention(
File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 64, in error_handler
  return fn(*args, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/keras/engine/base_layer.py", line 1096, in __call__
  outputs = call_fn(inputs, *args, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 92, in error_handler
  return fn(*args, **kwargs)
File "/workspaces/nv-deberta-tf2/electra/modeling_tf_deberta_v2.py", line 226, in call
  self_outputs = self.self(
File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 64, in error_handler
  return fn(*args, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/keras/engine/base_layer.py", line 1096, in __call__
  outputs = call_fn(inputs, *args, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 92, in error_handler
  return fn(*args, **kwargs)
File "/workspaces/nv-deberta-tf2/electra/modeling_tf_deberta_v2.py", line 876, in call
  if self.relative_attention:
File "/workspaces/nv-deberta-tf2/electra/modeling_tf_deberta_v2.py", line 878, in call
  rel_att = self.disentangled_att_bias(
File "/workspaces/nv-deberta-tf2/electra/modeling_tf_deberta_v2.py", line 991, in disentangled_att_bias
  if "c2p" in self.pos_att_type:
File "/workspaces/nv-deberta-tf2/electra/modeling_tf_deberta_v2.py", line 1012, in disentangled_att_bias
  c2p_att = tnp.take_along_axis(

         [[{{node pretraining_model/tf_deberta_v2_for_masked_lm/deberta/encoder/layer_._0/attention/self/BroadcastArgs}}]]
Traceback (most recent call last):
  File "run_pretraining.py", line 204, in <module>
    config = main(start_time)
  File "run_pretraining.py", line 184, in main
    trained_model = run_customized_training_loop(
  File "/workspaces/nv-deberta-tf2/electra/model_training_utils.py", line 675, in run_customized_training_loop
    train_steps_strategy(
  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 "/usr/local/lib/python3.8/dist-packages/tensorflow/python/eager/execute.py", line 54, in quick_execute
    tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
tensorflow.python.framework.errors_impl.InvalidArgumentError: Graph execution error:

Input 0 to node `pretraining_model/tf_deberta_v2_for_masked_lm/deberta/encoder/layer_._0/attention/self/BroadcastArgs` 
with op BroadcastArgs must be a compile-time constant.

XLA compilation requires that operator arguments that represent shapes or dimensions be evaluated to concrete values at compile time. 
This error means that a shape or dimension argument could not be evaluated at compile time, 
usually because the value of the argument depends on a parameter to the computation, 
on a variable, or on a stateful operation such as a random number generator.

Stack trace for op definition: 
File "run_pretraining.py", line 204, in <module>
  config = main(start_time)
File "run_pretraining.py", line 184, in main
  trained_model = run_customized_training_loop(
File "/workspaces/nv-deberta-tf2/electra/model_training_utils.py", line 675, in run_customized_training_loop
  train_steps_strategy(
File "/workspaces/nv-deberta-tf2/electra/model_training_utils.py", line 407, in train_steps_strategy
  if num_grad_accumulates != 1:
File "/workspaces/nv-deberta-tf2/electra/model_training_utils.py", line 408, in train_steps_strategy
  for step_idx in tf.range(steps * num_grad_accumulates):
File "/workspaces/nv-deberta-tf2/electra/model_training_utils.py", line 410, in train_steps_strategy
  strategy.run(_forward, args=(next(iterator),))
File "/workspaces/nv-deberta-tf2/electra/model_training_utils.py", line 324, in _forward
  loss, model_outputs = model(inputs, is_training=True)
File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 64, in error_handler
  return fn(*args, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/keras/engine/base_layer.py", line 1096, in __call__
  outputs = call_fn(inputs, *args, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 92, in error_handler
  return fn(*args, **kwargs)
File "/workspaces/nv-deberta-tf2/electra/modeling_tf_deberta_v2.py", line 2491, in call
  if config.uniform_generator:
File "/workspaces/nv-deberta-tf2/electra/modeling_tf_deberta_v2.py", line 2496, in call
  mlm_output = self._get_masked_lm_output(
File "/workspaces/nv-deberta-tf2/electra/modeling_tf_deberta_v2.py", line 2541, in _get_masked_lm_output
  if self._config.uniform_generator:
File "/workspaces/nv-deberta-tf2/electra/modeling_tf_deberta_v2.py", line 2550, in _get_masked_lm_output
  outputs = generator(
File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 64, in error_handler
  return fn(*args, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/keras/engine/base_layer.py", line 1096, in __call__
  outputs = call_fn(inputs, *args, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 92, in error_handler
  return fn(*args, **kwargs)
File "/workspaces/nv-deberta-tf2/electra/modeling_tf_utils.py", line 1872, in run_call_with_unpacked_inputs
  )
File "/workspaces/nv-deberta-tf2/electra/modeling_tf_deberta_v2.py", line 1880, in call
  outputs = self.deberta(
File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 64, in error_handler
  return fn(*args, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/keras/engine/base_layer.py", line 1096, in __call__
  outputs = call_fn(inputs, *args, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 92, in error_handler
  return fn(*args, **kwargs)
File "/workspaces/nv-deberta-tf2/electra/modeling_tf_utils.py", line 1872, in run_call_with_unpacked_inputs
  )
File "/workspaces/nv-deberta-tf2/electra/modeling_tf_deberta_v2.py", line 1617, in call
  encoder_outputs = self.encoder(
File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 64, in error_handler
  return fn(*args, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/keras/engine/base_layer.py", line 1096, in __call__
  outputs = call_fn(inputs, *args, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 92, in error_handler
  return fn(*args, **kwargs)
File "/workspaces/nv-deberta-tf2/electra/modeling_tf_deberta_v2.py", line 527, in call
  for i, layer_module in enumerate(self.layer):
File "/workspaces/nv-deberta-tf2/electra/modeling_tf_deberta_v2.py", line 532, in call
  layer_outputs = layer_module(
File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 64, in error_handler
  return fn(*args, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/keras/engine/base_layer.py", line 1096, in __call__
  outputs = call_fn(inputs, *args, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 92, in error_handler
  return fn(*args, **kwargs)
File "/workspaces/nv-deberta-tf2/electra/modeling_tf_deberta_v2.py", line 317, in call
  attention_outputs = self.attention(
File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 64, in error_handler
  return fn(*args, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/keras/engine/base_layer.py", line 1096, in __call__
  outputs = call_fn(inputs, *args, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 92, in error_handler
  return fn(*args, **kwargs)
File "/workspaces/nv-deberta-tf2/electra/modeling_tf_deberta_v2.py", line 226, in call
  self_outputs = self.self(
File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 64, in error_handler
  return fn(*args, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/keras/engine/base_layer.py", line 1096, in __call__
  outputs = call_fn(inputs, *args, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 92, in error_handler
  return fn(*args, **kwargs)
File "/workspaces/nv-deberta-tf2/electra/modeling_tf_deberta_v2.py", line 876, in call
  if self.relative_attention:
File "/workspaces/nv-deberta-tf2/electra/modeling_tf_deberta_v2.py", line 878, in call
  rel_att = self.disentangled_att_bias(
File "/workspaces/nv-deberta-tf2/electra/modeling_tf_deberta_v2.py", line 991, in disentangled_att_bias
  if "c2p" in self.pos_att_type:
File "/workspaces/nv-deberta-tf2/electra/modeling_tf_deberta_v2.py", line 1012, in disentangled_att_bias
  c2p_att = tnp.take_along_axis(

         [[{{node pretraining_model/tf_deberta_v2_for_masked_lm/deberta/encoder/layer_._0/attention/self/BroadcastArgs}}]]
         [[while/body/_1/while/StatefulPartitionedCall]] [Op:__inference_train_steps_strategy_177980]

[Rocketknight1]

@WissamAntoun Confirmed reproduction of the issue here. Our TF DeBERTa implementation seems to have issues with XLA - I'm investigating now.

[Rocketknight1]

@WissamAntoun We have a potential fix - I've confirmed that I can compile microsoft/deberta-v3-small with XLA on my local machine. Can you try installing this branch and let me know if this fixes the problem for you? You can use pip install git+https://github.com/huggingface/transformers.git@deberta-xla-fixes

[WissamAntoun]

I confirm it works on GPUs with XLA, and I got ~20% improved speedup.
I'm still testing now on TPUs, will let you know ASAP

[WissamAntoun]

Weirdly enough TPUs didn't seem to care about the changes 😅 even after we removed all the if branches

[Rocketknight1]

Hmm. Can you check that you don't get the slowdown if you switch the model to another model, like BERT or ELECTRA, while keeping all of the other code the same (especially data loading)? I know the profiling indicates that the GatherV2 is the problem, but I'm a little suspicious!

[WissamAntoun]

I tried disabling relative_attention in deberta, which makes the model a regular BERT, and the performance improved 40x 😅

[Rocketknight1]

@WissamAntoun So the issue really is in that gather! That's extremely interesting - with the simplified code, it's just a single call to tf.gather, but perhaps the batch_dims argument is not handled elegantly on TPU, or XLA converts it in a way that doesn't run well on TPU.

Is it possible that some kind of memory spill is occurring? Can you try lowering your batch size and increasing steps_per_execution?

If that isn't it, then I have no idea - maybe there's some way to rewrite the gather, but I don't really know what to try!

[WissamAntoun]

@Rocketknight1 I tried your suggestions without any success, sadly!

Then I tried replacing the whole take_along_axis function with tf.gather(..,...,batch_dims=2) which is equivalent, according to this test I made. GPU still runs fine, TPU still has the same issue 😔.

I also ran out of ideas to try, now I'm just waiting for the TPU gods 😅

View code

#%%
import tensorflow as tf

#%%
x_shape = [32, 128, 512]
indices_shape = [32, 128, 128]
x = tf.random.uniform(shape=x_shape)
indices = tf.random.uniform(shape=indices_shape, minval=1, maxval=128, dtype=tf.int32)
#%%
flat_x = tf.reshape(x, (-1, x_shape[-1]))
print(flat_x.shape)  # (4096, 512)
flat_indices = tf.reshape(indices, (-1, indices_shape[-1]))
print(flat_indices.shape)  # (4096, 128)

#%%
gathered = tf.gather(
    params=flat_x, indices=flat_indices, batch_dims=1, validate_indices=None
)
print(gathered.shape)  # (4096, 128)
gathered_reshaped = tf.reshape(gathered, indices.shape)
print(gathered_reshaped.shape)  # ( 32, 128, 128)

# %%
gathered2 = tf.gather(params=x, indices=indices, batch_dims=2, validate_indices=None)
print(gathered2.shape)  # (32, 128, 128)
# %%
tf.assert_equal(gathered2, gathered_reshaped)  # passes

# %%

[Rocketknight1]

I'm clueless in that case - @patrickvonplaten @sanchit-gandhi do you have any idea why a gather or take_along_axis op which is performant on GPU and compiles with XLA would become a huge bottleneck on TPU?