RFC: Replace custom TF embeddings by Keras embeddings

[gante]

What does this PR do?

This is an RFC with a code example in the PR -- my primary goal is not to get the PR approved, but rather to discuss an improvement to our TF codebase, with an example that passes all tests.

Context

In our TF implementation of models with embedding layers, we rely on two custom-made classes:

1. TFSharedEmbeddings -- a custom embedding layer whose added benefit is the ability to also use it as a dense layer;
2. TFWrappedEmbeddings -- used to manipulate the scope of the weights, which would normally depend on the layer where the weights are first used in an operation. Used with tied weight embeddings.

Problems with this setup include:

1. Users can't use the expected Keras tools to handle embeddings;
2. Relies on TF1 compatibility to set the right name to the weights (tf.compat.v1.variable_scope);
3. Resizing the embeddings, a major source of bugs atm, uses complex logic that consists in manipulating tf.Variable.

Proposed change

The proposal is straightforward: replace TFSharedEmbeddings by tf.keras.layer.Embedding, remove TFWrappedEmbeddings, and make the necessary adaptations. A few details to keep in mind (and that you can browse in the code):

1. There is a whole new code path for resizing the embeddings. Instead of if/else in the original functions, changed functions were rewritten with _v2 prepended to their name (which should also facilitate the transition). You can see that the new functions are simpler than the originals;
2. Giving the right name to the embeddings (so we can load existing weights) was the hardest part. TF had limited maneuverability here. To pull it off, I relied on UNDOCUMENTED behavior of tf.name_scope. Normally, tf.name_scope appends to the existing scope -- if the scope for the current layer is foo, weights are in the form of foo/weights:0; if we add a context manager tf.name_scope("bar"), weights will be in the form of foo/bar/weights:0. However, if the argument of tf.name_scope ends with /, then it will be a stand-alone name scope. Taking the previous example, with tf.name_scope("bar/"), weights will be in the form of bar/weights:0. This behavior has been in the TF codebase since its first commit (>7 yrs), and replacing TFWrappedEmbeddings relies on this behavior;
3. The existing TF Bart assumes the input/output embeddings are tied, which PT Bart does not assume. I've not changed this part, so the example you can see in this PR is for models with tied weights;
4. If you open PT Bart and compare side by side, you'll see that the implementations on the two frameworks are now more similar :)

I estimate about a 1-2 weeks worth of work to propagate the change, which includes:

1. Replace all TFSharedEmbeddings and TFWrappedEmbeddings;
2. Handle edge cases -- resize_token_embeddings is not implemented/is broken (and untested) in several recent TF models;
3. Remove/deprecate old code after 1 and 2 are done.

Pros/cons

(+) Simpler and smaller codebase, especially for the models;
(+) TF model code closer to PT's;
(+) Keras-native embeddings ( = users and contributors can be more productive);
(+) resize_token_embeddings usable in all models;
(-) Time spent refactoring is time not spent building new things;
(-) The solution still relies on named scopes for cross-framework weight matching, which is hacky.

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[gante]

(make fix-copies)

[gante]

The exact same test exists in test_modeling_tf_common, and is touched in this PR

[Rocketknight1]

My thoughts:

• I think the behaviour of tf.name_scope is intended and stable, even if it's not documented (TF documentation isn't always great). I think we can rely on that safely, and it's a lot better than using compatibility methods from v1.
• I agree that how we're doing this right now isn't great, and this code is a big improvement.
• I think how we use name_scope is still a little problematic. However, I don't want to make any big breaking changes there right now because the PT codebase will probably also change soon to use whatever new pickle-free state dict save format the PT devs come up with!

So overall, I think this is a good addition that cleans up a longstanding source of issues in the code, and shouldn't take too long to implement across the codebase.

[sgugger]

It does look cleaner this way, and handling the resize token embeddings for all new models is something that should be done either way. So I'm in favor of this switch.