albert flax

[kamalkraj]

What does this PR do?

ALBERT Flax Model

Fixes # (issue)

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@LysandreJik @patrickvonplaten

[patrickvonplaten]

Really cool PR @kamalkraj !!!

I think we can merge it soon. The main points to add are:

• A slow integration test making sure that Flax gives the same results as PT for Albert
• Add # Copied from ... statements to make sure the code is kept maintainable

[kamalkraj]

@patrickvonplaten
Done changes according to your suggestions.
Thanks for the review

[stefan-it]

Hey @kamalkraj , thanks for that PR! Can't wait to road-test it 🤗

I tried training several ALBERT models with the official implementation, but had no luck in training a good performing model.

[kamalkraj]

I tried training several ALBERT models with the official implementation, but had no luck in training a good performing model.

pre-training with TF 1 code?

[stefan-it]

Yeah, it was the TF 1 code base, and I've also trained various model sizes. Let's see if I have more luck using the FLAX implementation 😅

[patrickvonplaten]

Awesome addition @kamalkraj - thanks a lot :-) Left a couple of final additions

[kamalkraj]

@patrickvonplaten
Done changes according to review.

[patrickvonplaten]

Slow tests are passing on CPU - thanks for the model addition @kamalkraj !