Unable to resume checkpoints with TFBertModel using tf.distribute.Strategy and a custom LM head that shares the underlying TFBertEmbeddings layer
Author: marhlder
Pull request: no

## Environment info

- `transformers` version: 4.5.0
- Platform: Linux (Ubuntu 18.04, 20.04 + CentOs)
- Python version: 3.7.4
- PyTorch version (GPU?): N/A
- Tensorflow version (GPU?): 2.4.1 GPU
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: Yes, using MirroredStrategy

Models:

- BERT: @LysandreJik

## Information

Model I am using (Bert, XLNet ...):

The problem arises when using:
* TFBertModel
* tf.distribute.MirroredStrategy
* A custom LM head that shares a TFBertEmbeddings layer with TFBertModel

The tasks I am working on is:
* [x] my own task and dataset: (not relevant)

## To reproduce

Steps to reproduce the behavior:

1. Train a BERT model, with MirroredStrategy, based on the TFBertModel class using using a custom head with a shared TFBertEmbeddings layer
2. Stop training
3. Attempt to resume using checkpoints

<!-- 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.-->

## Expected behavior
I should be able to re-load checkpoints trained with shared layers under a MirroredStrategy

## Problem 
Someone has used tf.name_scope() and assumed that would change anything about the variable names in tf 2.x (It does not, it only modifies the names of ops)
See the build method of TFBertEmbeddings: https://github.com/huggingface/transformers/blob/cd56f3fe7eae4a53a9880e3f5e8f91877a78271c/src/transformers/models/bert/modeling_tf_bert.py#L155

In the above referenced code the variable name "embeddings" is used for both of the variables created with the property names "token_type_embeddings" and "position_embeddings". This does not matter in most cases as TensorFlow 2.x will use the property names of the variables, i.e. the object hierarchy path, and not the variable name given to the add_weight member function of a given Keras layer.
But it does matter in this case, it would seem, as the Distribution Strategy has issues resolving where to assign the saved variable from a checkpoint given that they have the same name. 

##  Proposed solution 
Stop the meaningless use of tf.name_scope() (See: https://www.tensorflow.org/api_docs/python/tf/name_scope)
Give variables different names with the add_weight member function

## Stack trace

`ValueError: in user code:    /opt/ml/code/unsilo_ml/python/supervisors/TF2custom_keras_loop_supervisor.py:395 train_epoch  *        loss, local_global_step = distributed_train_step(x)    /opt/ml/code/unsilo_ml/python/supervisors/TF2custom_keras_loop_supervisor.py:359 distributed_train_step  *        per_replica_losses, per_replica_global_step = self.dist_strategy.run(    /opt/ml/code/unsilo_ml/python/supervisors/TF2custom_keras_loop_supervisor.py:313 train_step  *        predictions = keras_model(features, training=True)    /opt/ml/code/unsilo_ml/python/models/base_models/base_model.py:72 call  *        return self.build_forward_pass(training=training, inputs=inputs)    /opt/ml/code/unsilo_ml/python/models/multitask_model.py:103 build_forward_pass  *        inputs_with_encoder_output = self.prepare_inputs_with_encoder_output(    /opt/ml/code/unsilo_ml/python/models/multitask_model.py:147 prepare_inputs_with_encoder_output  *        encoder_outputs = self.encoder(inputs, training=training)    /opt/ml/code/unsilo_ml/python/modules/encoders/util_encoders/pipe_encoder.py:18 call  *        encoder_output = self.resolve_tensor_dict(    /opt/ml/code/unsilo_ml/python/modules/encoders/bert_encoder.py:83 call  *        hidden_states = self.bert_model(    /opt/conda/lib/python3.7/site-packages/transformers/models/bert/modeling_tf_bert.py:887 call  *        outputs = self.bert(    /opt/conda/lib/python3.7/site-packages/transformers/models/bert/modeling_tf_bert.py:645 call  *        embedding_output = self.embeddings(    /opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/engine/base_layer.py:1008 __call__  **        self._maybe_build(inputs)    /opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/engine/base_layer.py:2710 _maybe_build        self.build(input_shapes)  # pylint:disable=not-callable    /opt/conda/lib/python3.7/site-packages/transformers/models/bert/modeling_tf_bert.py:159 build        initializer=get_initializer(self.initializer_range),    /opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/engine/base_layer.py:639 add_weight        caching_device=caching_device)    /opt/conda/lib/python3.7/site-packages/tensorflow/python/training/tracking/base.py:810 _add_variable_with_custom_getter        **kwargs_for_getter)    /opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/engine/base_layer_utils.py:142 make_variable        shape=variable_shape if variable_shape else None)    /opt/conda/lib/python3.7/site-packages/tensorflow/python/ops/variables.py:260 __call__        return cls._variable_v1_call(*args, **kwargs)    /opt/conda/lib/python3.7/site-packages/tensorflow/python/ops/variables.py:221 _variable_v1_call        shape=shape)    /opt/conda/lib/python3.7/site-packages/tensorflow/python/ops/variables.py:67 getter        return captured_getter(captured_previous, **kwargs)    /opt/conda/lib/python3.7/site-packages/tensorflow/python/distribute/shared_variable_creator.py:69 create_new_variable        v = next_creator(**kwargs)    /opt/conda/lib/python3.7/site-packages/tensorflow/python/ops/variables.py:67 getter        return captured_getter(captured_previous, **kwargs)    /opt/conda/lib/python3.7/site-packages/tensorflow/python/distribute/distribute_lib.py:2083 creator_with_resource_vars        created = self._create_variable(next_creator, **kwargs)    /opt/conda/lib/python3.7/site-packages/tensorflow/python/distribute/mirrored_strategy.py:489 _create_variable        distribute_utils.VARIABLE_POLICY_MAPPING, **kwargs)    /opt/conda/lib/python3.7/site-packages/tensorflow/python/distribute/distribute_utils.py:311 create_mirrored_variable        value_list = real_mirrored_creator(**kwargs)    /opt/conda/lib/python3.7/site-packages/tensorflow/python/distribute/mirrored_strategy.py:481 _real_mirrored_creator        v = next_creator(**kwargs)    /opt/conda/lib/python3.7/site-packages/tensorflow/python/ops/variables.py:67 getter        return captured_getter(captured_previous, **kwargs)    /opt/conda/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py:714 variable_capturing_scope        lifted_initializer_graph=lifted_initializer_graph, **kwds)    /opt/conda/lib/python3.7/site-packages/tensorflow/python/ops/variables.py:264 __call__        return super(VariableMetaclass, cls).__call__(*args, **kwargs)    /opt/conda/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py:227 __init__        initial_value = initial_value()    /opt/conda/lib/python3.7/site-packages/tensorflow/python/training/tracking/base.py:82 __call__        self._checkpoint_position, shape, shard_info=shard_info)    /opt/conda/lib/python3.7/site-packages/tensorflow/python/training/tracking/base.py:117 __init__        self.wrapped_value.set_shape(shape)    /opt/conda/lib/python3.7/site-packages/tensorflow/python/framework/ops.py:1217 set_shape        (self.shape, shape))    ValueError: Tensor's shape (512, 768) is not compatible with supplied shape [2, 768] | ValueError: in user code: /opt/ml/code/unsilo_ml/python/supervisors/TF2custom_keras_loop_supervisor.py:395 train_epoch * loss, local_global_step = distributed_train_step(x) /opt/ml/code/unsilo_ml/python/supervisors/TF2custom_keras_loop_supervisor.py:359 distributed_train_step * per_replica_losses, per_replica_global_step = self.dist_strategy.run( /opt/ml/code/unsilo_ml/python/supervisors/TF2custom_keras_loop_supervisor.py:313 train_step * predictions = keras_model(features, training=True) /opt/ml/code/unsilo_ml/python/models/base_models/base_model.py:72 call * return self.build_forward_pass(training=training, inputs=inputs) /opt/ml/code/unsilo_ml/python/models/multitask_model.py:103 build_forward_pass * inputs_with_encoder_output = self.prepare_inputs_with_encoder_output( /opt/ml/code/unsilo_ml/python/models/multitask_model.py:147 prepare_inputs_with_encoder_output * encoder_outputs = self.encoder(inputs, training=training) /opt/ml/code/unsilo_ml/python/modules/encoders/util_encoders/pipe_encoder.py:18 call * encoder_output = self.resolve_tensor_dict( /opt/ml/code/unsilo_ml/python/modules/encoders/bert_encoder.py:83 call * hidden_states = self.bert_model( /opt/conda/lib/python3.7/site-packages/transformers/models/bert/modeling_tf_bert.py:887 call * outputs = self.bert( /opt/conda/lib/python3.7/site-packages/transformers/models/bert/modeling_tf_bert.py:645 call * embedding_output = self.embeddings( /opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/engine/base_layer.py:1008 __call__ ** self._maybe_build(inputs) /opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/engine/base_layer.py:2710 _maybe_build self.build(input_shapes) # pylint:disable=not-callable /opt/conda/lib/python3.7/site-packages/transformers/models/bert/modeling_tf_bert.py:159 build initializer=get_initializer(self.initializer_range), /opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/engine/base_layer.py:639 add_weight caching_device=caching_device) /opt/conda/lib/python3.7/site-packages/tensorflow/python/training/tracking/base.py:810 _add_variable_with_custom_getter **kwargs_for_getter) /opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/engine/base_layer_utils.py:142 make_variable shape=variable_shape if variable_shape else None) /opt/conda/lib/python3.7/site-packages/tensorflow/python/ops/variables.py:260 __call__ return cls._variable_v1_call(*args, **kwargs) /opt/conda/lib/python3.7/site-packages/tensorflow/python/ops/variables.py:221 _variable_v1_call shape=shape) /opt/conda/lib/python3.7/site-packages/tensorflow/python/ops/variables.py:67 getter return captured_getter(captured_previous, **kwargs) /opt/conda/lib/python3.7/site-packages/tensorflow/python/distribute/shared_variable_creator.py:69 create_new_variable v = next_creator(**kwargs) /opt/conda/lib/python3.7/site-packages/tensorflow/python/ops/variables.py:67 getter return captured_getter(captured_previous, **kwargs) /opt/conda/lib/python3.7/site-packages/tensorflow/python/distribute/distribute_lib.py:2083 creator_with_resource_vars created = self._create_variable(next_creator, **kwargs) /opt/conda/lib/python3.7/site-packages/tensorflow/python/distribute/mirrored_strategy.py:489 _create_variable distribute_utils.VARIABLE_POLICY_MAPPING, **kwargs) /opt/conda/lib/python3.7/site-packages/tensorflow/python/distribute/distribute_utils.py:311 create_mirrored_variable value_list = real_mirrored_creator(**kwargs) /opt/conda/lib/python3.7/site-packages/tensorflow/python/distribute/mirrored_strategy.py:481 _real_mirrored_creator v = next_creator(**kwargs) /opt/conda/lib/python3.7/site-packages/tensorflow/python/ops/variables.py:67 getter return captured_getter(captured_previous, **kwargs) /opt/conda/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py:714 variable_capturing_scope lifted_initializer_graph=lifted_initializer_graph, **kwds) /opt/conda/lib/python3.7/site-packages/tensorflow/python/ops/variables.py:264 __call__ return super(VariableMetaclass, cls).__call__(*args, **kwargs) /opt/conda/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py:227 __init__ initial_value = initial_value() /opt/conda/lib/python3.7/site-packages/tensorflow/python/training/tracking/base.py:82 __call__ self._checkpoint_position, shape, shard_info=shard_info) /opt/conda/lib/python3.7/site-packages/tensorflow/python/training/tracking/base.py:117 __init__ self.wrapped_value.set_shape(shape) /opt/conda/lib/python3.7/site-packages/tensorflow/python/framework/ops.py:1217 set_shape (self.shape, shape)) ValueError: Tensor's shape (512, 768) is not compatible with supplied shape [2, 768]`

Comments:
I have tested it locally under tf.distribute.OneDeviceStrategy and the problem seems to be the same.

```
Traceback (most recent call last):
  File "/home/marhlder/unsilo/unsilo-ml/unsilo_ml/python/run.py", line 73, in <module>
    cli()
  File "/home/marhlder/anaconda3/envs/unsilo-ml/lib/python3.7/site-packages/click/core.py", line 829, in __call__
    return self.main(*args, **kwargs)
  File "/home/marhlder/anaconda3/envs/unsilo-ml/lib/python3.7/site-packages/click/core.py", line 782, in main
    rv = self.invoke(ctx)
  File "/home/marhlder/anaconda3/envs/unsilo-ml/lib/python3.7/site-packages/click/core.py", line 1259, in invoke
    return _process_result(sub_ctx.command.invoke(sub_ctx))
  File "/home/marhlder/anaconda3/envs/unsilo-ml/lib/python3.7/site-packages/click/core.py", line 1066, in invoke
    return ctx.invoke(self.callback, **ctx.params)
  File "/home/marhlder/anaconda3/envs/unsilo-ml/lib/python3.7/site-packages/click/core.py", line 610, in invoke
    return callback(*args, **kwargs)
  File "/home/marhlder/unsilo/unsilo-ml/unsilo_ml/python/run.py", line 31, in train
    provisioner.train()
  File "/home/marhlder/unsilo/unsilo-ml/mlops/provisioning/local/local.py", line 72, in train
    **self.entry_point_parameters,
  File "/home/marhlder/unsilo/unsilo-ml//unsilo_ml/python/provisioner_entry_point.py", line 78, in run
    tc.train(tracker=tracker)
  File "/home/marhlder/unsilo/unsilo-ml/unsilo_ml/python/module_composition/train_composer.py", line 83, in train
    return self.model_supervisor.train(**kwargs)
  File "/home/marhlder/unsilo/unsilo-ml/unsilo_ml/python/supervisors/TF2custom_keras_loop_supervisor.py", line 463, in train
    train_epoch()
  File "/home/marhlder/anaconda3/envs/unsilo-ml/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 828, in __call__
    result = self._call(*args, **kwds)
  File "/home/marhlder/anaconda3/envs/unsilo-ml/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 871, in _call
    self._initialize(args, kwds, add_initializers_to=initializers)
  File "/home/marhlder/anaconda3/envs/unsilo-ml/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 726, in _initialize
    *args, **kwds))
  File "/home/marhlder/anaconda3/envs/unsilo-ml/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 2969, in _get_concrete_function_internal_garbage_collected
    graph_function, _ = self._maybe_define_function(args, kwargs)
  File "/home/marhlder/anaconda3/envs/unsilo-ml/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 3361, in _maybe_define_function
    graph_function = self._create_graph_function(args, kwargs)
  File "/home/marhlder/anaconda3/envs/unsilo-ml/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 3206, in _create_graph_function
    capture_by_value=self._capture_by_value),
  File "/home/marhlder/anaconda3/envs/unsilo-ml/lib/python3.7/site-packages/tensorflow/python/framework/func_graph.py", line 990, in func_graph_from_py_func
    func_outputs = python_func(*func_args, **func_kwargs)
  File "/home/marhlder/anaconda3/envs/unsilo-ml/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 634, in wrapped_fn
    out = weak_wrapped_fn().__wrapped__(*args, **kwds)
  File "/home/marhlder/anaconda3/envs/unsilo-ml/lib/python3.7/site-packages/tensorflow/python/framework/func_graph.py", line 977, in wrapper
    raise e.ag_error_metadata.to_exception(e)
ValueError: in user code:

    /home/marhlder/unsilo/unsilo-ml/unsilo_ml/python/supervisors/TF2custom_keras_loop_supervisor.py:395 train_epoch  *
        loss, local_global_step = distributed_train_step(x)
    /home/marhlder/unsilo/unsilo-ml/unsilo_ml/python/supervisors/TF2custom_keras_loop_supervisor.py:359 distributed_train_step  *
        per_replica_losses, per_replica_global_step = self.dist_strategy.run(
    /home/marhlder/unsilo/unsilo-ml/unsilo_ml/python/supervisors/TF2custom_keras_loop_supervisor.py:313 train_step  *
        predictions = keras_model(features, training=True)
    /home/marhlder/unsilo/unsilo-ml/unsilo_ml/python/models/base_models/base_model.py:72 call  *
        return self.build_forward_pass(training=training, inputs=inputs)
    /home/marhlder/unsilo/unsilo-ml/unsilo_ml/python/models/multitask_model.py:103 build_forward_pass  *
        inputs_with_encoder_output = self.prepare_inputs_with_encoder_output(
    /home/marhlder/unsilo/unsilo-ml/unsilo_ml/python/models/multitask_model.py:147 prepare_inputs_with_encoder_output  *
        encoder_outputs = self.encoder(inputs, training=training)
    /home/marhlder/unsilo/unsilo-ml/unsilo_ml/python/modules/encoders/util_encoders/pipe_encoder.py:18 call  *
        encoder_output = self.resolve_tensor_dict(
    /home/marhlder/unsilo/unsilo-ml/unsilo_ml/python/modules/encoders/bert_encoder.py:83 call  *
        hidden_states = self.bert_model(
    /home/marhlder/anaconda3/envs/unsilo-ml/lib/python3.7/site-packages/transformers/models/bert/modeling_tf_bert.py:887 call  *
        outputs = self.bert(
    /home/marhlder/anaconda3/envs/unsilo-ml/lib/python3.7/site-packages/transformers/models/bert/modeling_tf_bert.py:645 call  *
        embedding_output = self.embeddings(
    /home/marhlder/anaconda3/envs/unsilo-ml/lib/python3.7/site-packages/tensorflow/python/keras/engine/base_layer.py:1008 __call__  **
        self._maybe_build(inputs)
    /home/marhlder/anaconda3/envs/unsilo-ml/lib/python3.7/site-packages/tensorflow/python/keras/engine/base_layer.py:2710 _maybe_build
        self.build(input_shapes)  # pylint:disable=not-callable
    /home/marhlder/anaconda3/envs/unsilo-ml/lib/python3.7/site-packages/transformers/models/bert/modeling_tf_bert.py:159 build
        initializer=get_initializer(self.initializer_range),
    /home/marhlder/anaconda3/envs/unsilo-ml/lib/python3.7/site-packages/tensorflow/python/keras/engine/base_layer.py:639 add_weight
        caching_device=caching_device)
    /home/marhlder/anaconda3/envs/unsilo-ml/lib/python3.7/site-packages/tensorflow/python/training/tracking/base.py:810 _add_variable_with_custom_getter
        **kwargs_for_getter)
    /home/marhlder/anaconda3/envs/unsilo-ml/lib/python3.7/site-packages/tensorflow/python/keras/engine/base_layer_utils.py:142 make_variable
        shape=variable_shape if variable_shape else None)
    /home/marhlder/anaconda3/envs/unsilo-ml/lib/python3.7/site-packages/tensorflow/python/ops/variables.py:260 __call__
        return cls._variable_v1_call(*args, **kwargs)
    /home/marhlder/anaconda3/envs/unsilo-ml/lib/python3.7/site-packages/tensorflow/python/ops/variables.py:221 _variable_v1_call
        shape=shape)
    /home/marhlder/anaconda3/envs/unsilo-ml/lib/python3.7/site-packages/tensorflow/python/ops/variables.py:67 getter
        return captured_getter(captured_previous, **kwargs)
    /home/marhlder/anaconda3/envs/unsilo-ml/lib/python3.7/site-packages/tensorflow/python/distribute/distribute_lib.py:2083 creator_with_resource_vars
        created = self._create_variable(next_creator, **kwargs)
    /home/marhlder/anaconda3/envs/unsilo-ml/lib/python3.7/site-packages/tensorflow/python/distribute/one_device_strategy.py:278 _create_variable
        return next_creator(**kwargs)
    /home/marhlder/anaconda3/envs/unsilo-ml/lib/python3.7/site-packages/tensorflow/python/ops/variables.py:67 getter
        return captured_getter(captured_previous, **kwargs)
    /home/marhlder/anaconda3/envs/unsilo-ml/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py:714 variable_capturing_scope
        lifted_initializer_graph=lifted_initializer_graph, **kwds)
    /home/marhlder/anaconda3/envs/unsilo-ml/lib/python3.7/site-packages/tensorflow/python/ops/variables.py:264 __call__
        return super(VariableMetaclass, cls).__call__(*args, **kwargs)
    /home/marhlder/anaconda3/envs/unsilo-ml/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py:227 __init__
        initial_value = initial_value()
    /home/marhlder/anaconda3/envs/unsilo-ml/lib/python3.7/site-packages/tensorflow/python/training/tracking/base.py:82 __call__
        self._checkpoint_position, shape, shard_info=shard_info)
    /home/marhlder/anaconda3/envs/unsilo-ml/lib/python3.7/site-packages/tensorflow/python/training/tracking/base.py:117 __init__
        self.wrapped_value.set_shape(shape)
    /home/marhlder/anaconda3/envs/unsilo-ml/lib/python3.7/site-packages/tensorflow/python/framework/ops.py:1217 set_shape
        (self.shape, shape))

    ValueError: Tensor's shape (512, 768) is not compatible with supplied shape [2, 768]


Process finished with exit code 1
```


#11202 

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.