distilbert-flax

[kamalkraj]

What does this PR do?

DistilBert Flax

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.

@VictorSanh @patrickvonplaten

[patrickvonplaten]

Great Great job @kamalkraj!

Think the only major thing to update is to docs in the modeling file (at the moment it looks like it's the PyTorch docs, but should be Flax :-))

[kamalkraj]

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

[stefan-it]

Hi @kamalkraj , I'm also really interested in that PR - thanks for adding it 🤗

Do you also plan to add a script for the distillation process (like it is done in the "old" script), as I would like to re-distillate some of my previous DistilBERT models (I don't have access to multi GPU setups, only to TPUs at the moment).

[kamalkraj]

Hi @stefan-it,

I will go through the scripts and pings you.
I have multi-GPU access. Which TPU do you use? v3-8 ?

[patrickvonplaten]

JAX_PLATFORM_NAME=cpu RUN_SLOW=1 pytest tests/test_modeling_flax_distilbert.py::FlaxDistilBertModelIntegrationTest::test_inference_no_head_absolute_embedding

passes and the code looks good :-) Ready to merge IMO 🎉 !

@patil-suraj the slow test doesn't pass on TPU since distilbert has pretty extreme activations in the forward pass like a couple of other models. We need to think a bit how to adapt the slow test depending on whether they're run on TPU or not in general...

[patrickvonplaten]

Great work @kamalkraj !