Issue in layer-drop implementation in TensorFlow models in graph mode
Author: thevasudevgupta
Pull request: no

## Environment info

- `transformers` version: 4.8.1
- Platform: Linux-5.4.104+-x86_64-with-Ubuntu-18.04-bionic
- Python version: 3.7.10
- PyTorch version (GPU?): 1.9.0+cu102 (False)
- Tensorflow version (GPU?): 2.5.0 (False)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: No

### Who can help

@Rocketknight1

## Information

Model I am using: TFBartForConditionalGeneration

The problem arises when using:
* [ ] the official example scripts: (give details below)
* [x] my own modified scripts: (give details below)

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [x] my own task or dataset: (give details below)

## To reproduce

```python
from transformers import TFBartForConditionalGeneration, BartConfig

# keeping layerdrop to be very high value for demonstration error
model = TFBartForConditionalGeneration(BartConfig(encoder_layerdrop=0.5))

import tensorflow as tf
import numpy as np

array = np.random.randint(1, 300, size=(4, 256))
dataset = tf.constant(array, dtype=tf.int32)

# following cell works perfectly when `tf.function(...)` is removed
@tf.function
def train_step(tensor):
  return model(tensor, training=True)

from tqdm.auto import tqdm

for tensor in tqdm(dataset, total=len(dataset)):
  tensor = tf.expand_dims(tensor, 0)
  output = train_step(tensor)
```

You can checkout this [small Colab notebook](https://colab.research.google.com/drive/1ACfyQcSUtv0pQGD_mZEvhjEO0tOzS2zR?usp=sharing) also for reproducing the error.

<!-- If you have code snippets, error messages, stack traces please provide them here as well.
     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->
```python
ValueError: in user code:

    <ipython-input-5-ca2e97b30313>:4 train_step  *
        return model(tensor, training=True)
    /usr/local/lib/python3.7/dist-packages/transformers/models/bart/modeling_tf_bart.py:1393 call  *
        outputs = self.model(
    /usr/local/lib/python3.7/dist-packages/transformers/models/bart/modeling_tf_bart.py:1125 call  *
        inputs["encoder_outputs"] = self.encoder(
    /usr/local/lib/python3.7/dist-packages/transformers/models/bart/modeling_tf_bart.py:764 call  *
        hidden_states, attn = encoder_layer(
    /usr/local/lib/python3.7/dist-packages/transformers/models/bart/modeling_tf_bart.py:305 call  *
        hidden_states, self_attn_weights, _ = self.self_attn(
    /usr/local/lib/python3.7/dist-packages/transformers/models/bart/modeling_tf_bart.py:178 call  *
        query_states = self.q_proj(hidden_states) * self.scaling
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/engine/base_layer.py:1023 __call__  **
        self._maybe_build(inputs)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/engine/base_layer.py:2625 _maybe_build
        self.build(input_shapes)  # pylint:disable=not-callable
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/layers/core.py:1198 build
        trainable=True)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/engine/base_layer.py:655 add_weight
        caching_device=caching_device)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/training/tracking/base.py:815 _add_variable_with_custom_getter
        **kwargs_for_getter)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/engine/base_layer_utils.py:139 make_variable
        shape=variable_shape if variable_shape else None)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/ops/variables.py:260 __call__
        return cls._variable_v1_call(*args, **kwargs)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/ops/variables.py:221 _variable_v1_call
        shape=shape)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/ops/variables.py:67 getter
        return captured_getter(captured_previous, **kwargs)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/def_function.py:769 invalid_creator_scope
        "tf.function-decorated function tried to create "

    ValueError: tf.function-decorated function tried to create variables on non-first call.

```

Side note: I have checked this same thing for TFWav2Vec2 also, but same issue is happening. So, possibly all TF model using layer-drop needs to be fixed.

## Expected behavior

layer drop should work perfectly in graph mode.


Comments:
Confirmed that the issue is reproducible at my end, we're investigating!

On investigation, I'm pretty sure the issue is caused by the way we're doing layerdrop: https://github.com/huggingface/transformers/blob/master/src/transformers/models/bart/modeling_tf_bart.py#L755-L772

This code is correct for eager execution, but I suspect in graph mode that this leads to the creation of new variables and graph edges whenever a layer is skipped for the first time. I can see some workarounds, but unfortunately no perfect ones - this seems like a fundamental limitation of the way graph mode works in TF.

You're welcome to investigate and try to find a solution if you like, but we're probably just going to explicitly disable layer drop in graph mode for now.

Yeah sure, I am also looking for some solution on this. Will keep you updated (or will make a PR) if I get some solution.

Cool! We'll hold off on disabling it for now - if you find a solution, let us know, and don't panic if it turns out to be impossible - just say so and we'll close this issue and disable layerdrop in graph mode instead. Thanks for your help!

@Rocketknight1,

I think I got a solution to this:
```python
# we will define this in the end of __init__(...)
self.step_0 = True

# then we will replace layer-drop condition with this:
if (not self.step_0) and inputs["training"] and (dropout_probability < self.layerdrop):  # skip the layer
    continue

# in the end of layer (just before return), we will do this
self.step_0 = False
```

Code works without any error after adding above stuff with `layer-drop > 0`. Checkout this for complete code: https://github.com/vasudevgupta7/transformers/commit/acf69cea945ebe97293621ba8730a7d988f2c2aa


@Rocketknight1, do think it's correct?? Like I am not sure but if graph is built with all the layers in the first step then will `continue` in the next steps work???

Thanks!

Hi, firstly I'm extremely sorry for the slow response! I was working on another project and had to drop my Github issues for a while.

I'm not sure this works, though - I *think* the value of `self.step_0` will just be treated as a constant at compilation time. As a result, this code won't cause errors, but it will never skip any layers either. Can you test it with a high value for the layerdrop probability and see if you get different answers when you run the same batch multiple times?

@Rocketknight1 I will test it the way you suggested. Thanks for your reply!!

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/master/CONTRIBUTING.md) are likely to be ignored.