pre-training deberta model on TPUs
### System Info

I want to pre-train the DeBERTa model (TFDebertaV2ForMaskedLM) using TPUs.

However, I got errors when I set "relative_attention" as Tue.

Although I tried to pre-train using a small model (the number of layers: 1, hidden dimension: 256),
I got the same error messages.

Is there any problem in Huggingface's DeBERTa model related to Tensorflow or TPUs?
Is there any way to solve this problem?


### Who can help?

@LysandreJik 

### 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

The model config I setted is as below:

```
hidden_size: 256
intermediate_size: 1024
num_attention_heads: 8
num_hidden_layers: 2
vocab_size: 64000
attention_probs_dropout_prob: 0.1
hidden_dropout_prob: 0.1
initializer_range: 0.02
max_position_embeddings: 512
type_vocab_size: 0
layer_norm_eps: 1e-7
pos_att_type: ["p2c", "c2p"]
hidden_act: "gelu"

relative_attention: true
max_relative_positions: -1
position_buckets: 256
norm_rel_ebd: "layer_norm"
share_att_key: true
```

Then, I used TFDebertaV2ForMaskedLM class to pre-train the DeBERTa model.

### Expected behavior

The error messages I've encountered are as below:

```
2022-05-10 07:26:13.050675: I tensorflow/core/tpu/kernels/tpu_compile_op_common.cc:176] Compilation of 11977631266728880069 with session name  took 556.019199ms and failed
2022-05-10 07:26:13.050718: E tensorflow/core/tpu/kernels/tpu_compilation_cache_external.cc:113] Input 1 to node `while/tf_deberta_v2_for_masked_lm/deberta/encoder/layer_._0/attention/self/cond_1/transpose` with op Transpose 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.

	 [[{{node while/tf_deberta_v2_for_masked_lm/deberta/encoder/layer_._0/attention/self/cond_1/transpose}}]]

	 [[while/tf_deberta_v2_for_masked_lm/deberta/encoder/layer_._0/attention/self/cond_1]]

	 [[while]]
2022-05-10 07:26:13.050752: F tensorflow/core/tpu/kernels/tpu_program_group.cc:86] Check failed: xla_tpu_programs.size() > 0 (0 vs. 0)
https://symbolize.stripped_domain/r/?trace=7fc408fc003b,7fc408fc00bf,7fc35113be1b,7fc356b4faf9,7fc356bbc8c1,7fc356bbd4d9,7fc356bb3b3c,7fc356bb5b4c,7fc34d0614af,7fc34cff51d3,7fc34cfe7631,7fc356b3a474,7fc356b376d6,7fc34d7217ff,7fc408f62608&map=8803c190592e98180ccd06f60438606e92563a77:7fc34e143000-7fc362c20e00,d8d5a9b6e95c8bb11b4f7a7a6ac042e0e1f0a2ce:7fc34c504000-7fc34e089760
*** SIGABRT received by PID 337336 (TID 338092) on cpu 1 from PID 337336; stack trace: ***
PC: @     0x7fc408fc003b  (unknown)  raise
    @     0x7fc34b99862d        976  (unknown)
    @     0x7fc408fc00c0       3904  (unknown)
    @     0x7fc35113be1c        720  tensorflow::tpu::TpuProgramGroup::Initialize()
    @     0x7fc356b4fafa       1264  tensorflow::tpu::TpuCompilationCacheExternal::InitializeEntry()
    @     0x7fc356bbc8c2       1008  tensorflow::tpu::TpuCompilationCacheInterface::CompileIfKeyAbsentHelper()
    @     0x7fc356bbd4da        128  tensorflow::tpu::TpuCompilationCacheInterface::CompileIfKeyAbsent()
    @     0x7fc356bb3b3d       1056  tensorflow::tpu::TpuCompileOpKernelCommon::ComputeInternal()
    @     0x7fc356bb5b4d        544  tensorflow::tpu::TpuCompileOpKernelCommon::Compute()
    @     0x7fc34d0614b0        496  tensorflow::ThreadPoolDevice::Compute()
    @     0x7fc34cff51d4       2368  tensorflow::(anonymous namespace)::ExecutorState<>::Process()
    @     0x7fc34cfe7632         48  std::_Function_handler<>::_M_invoke()
    @     0x7fc356b3a475        160  Eigen::ThreadPoolTempl<>::WorkerLoop()
    @     0x7fc356b376d7         64  std::_Function_handler<>::_M_invoke()
    @     0x7fc34d721800         80  tensorflow::(anonymous namespace)::PThread::ThreadFn()
    @     0x7fc408f62609  (unknown)  start_thread
https://symbolize.stripped_domain/r/?trace=7fc408fc003b,7fc34b99862c,7fc408fc00bf,7fc35113be1b,7fc356b4faf9,7fc356bbc8c1,7fc356bbd4d9,7fc356bb3b3c,7fc356bb5b4c,7fc34d0614af,7fc34cff51d3,7fc34cfe7631,7fc356b3a474,7fc356b376d6,7fc34d7217ff,7fc408f62608&map=8803c190592e98180ccd06f60438606e92563a77:7fc34e143000-7fc362c20e00,d8d5a9b6e95c8bb11b4f7a7a6ac042e0e1f0a2ce:7fc34c504000-7fc34e089760,b6a2e09e144eca321e149e2a834bcad7:7fc33ded8000-7fc34bcea2b0
E0510 07:26:13.215024  338092 coredump_hook.cc:292] RAW: Remote crash data gathering hook invoked.
E0510 07:26:13.215058  338092 coredump_hook.cc:384] RAW: Skipping coredump since rlimit was 0 at process start.
E0510 07:26:13.215066  338092 client.cc:222] RAW: Coroner client retries enabled (b/136286901), will retry for up to 30 sec.
E0510 07:26:13.215075  338092 coredump_hook.cc:447] RAW: Sending fingerprint to remote end.
E0510 07:26:13.215084  338092 coredump_socket.cc:124] RAW: Stat failed errno=2 on socket /var/google/services/logmanagerd/remote_coredump.socket
E0510 07:26:13.215094  338092 coredump_hook.cc:451] RAW: Cannot send fingerprint to Coroner: [NOT_FOUND] Missing crash reporting socket. Is the listener running?
E0510 07:26:13.215100  338092 coredump_hook.cc:525] RAW: Discarding core.
E0510 07:26:13.863904  338092 process_state.cc:779] RAW: Raising signal 6 with default behavior
```


I found the error raised at the disentangled_att_bias function,
especially when trying to calculate the disentangled attention bias score.
https://github.com/huggingface/transformers/blob/v4.18.0/src/transformers/models/deberta/modeling_tf_deberta.py#L683<|||||>The suspect is the take_along_axis function
https://github.com/huggingface/transformers/blob/v4.18.0/src/transformers/models/deberta_v2/modeling_tf_deberta_v2.py#L525<|||||>The variable "permutation" used in take_along_axis, e.g. tf.transpose(x, perm=permutation), is not compile-time constant. Therefore, I think the XLA compiler raises an error.

https://github.com/huggingface/transformers/blob/v4.18.0/src/transformers/models/deberta_v2/modeling_tf_deberta_v2.py#L532<|||||>Hey @Kyeongpil do you have any updates on this issue?<|||||>@WissamAntoun To train DeBERTa using TPUs, tf.rank(x) in take_along_axis should be replaced with constants, that is 3.<|||||>This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.