Add Flax Dinov2
Author: MHRDYN7
Pull request: yes
Labels: name: run-slow

This PR adds the Flax implementation of Dinov2, which seems to have been due since #25579. 

All the components of the pytorch Dinov2 model can be converted to flax except "interpolate_pos_encoding" which uses torch.nn.functional.interpolate. The closest jax function to replicate this is jax.image.scale_and_translate, however there seems to be a slight difference between these functions in the "Bicubic" mode (https://github.com/google/jax/issues/15768). 

In Dinov2, the pretrained weights of the position encoding are for the image size of 512, but we load images of size 224 into the model. The interpolate function acts to convert the shapes of the position encoding according to the size of the input images. The ViT model does have this interpolate function, but it's not there in the FlaxViT implementation as the config and input image sizes are the same. 

For now, I have directly loaded the pos_encoding weights from the pt model to flax, right after interpolation (which is saved in a safetensors file). This passes all the tests (including the two new integration tests added on top of the FlaxViT tests). Surely, this brute force approach to loading the original interpolated pos_encodings will not work, but otherwise, the slight deviations from jax scale_and _translate will fail the tests. @amyeroberts @sanchit-gandhi 

Other Remaining Tasks:
1. ~~Add the flax weights in .msgpack files to hub~~
2. ~~Test the SwiGLUFFN dense layer for vit giant~~


Comments:
Hi @MHRDYN7, thanks for working on this conversion! 

The interpolate logic is tricky. If loading the weights directly at the moment means this model passes, then that's a good guide the conversion is OK. We might have to do something where we remove this before merge and skip the equivalence tests. 

In the meantime, the first thing to do is get the other tests passing. Some of the failing tests are unrelated and have fixes upstream. Could you rebase on main to include these? For the quality checks, running `make fixup` should resolve 

Hi @amyeroberts, I did try make fixup and I'm not really sure why the two tests are still failing. Moreover, what should be done for skipping the equivalence test? Should I just remove the directly loaded tensors after interpolate and change the "expected_slice" tensor in the integration tests accordingly to make them pass? 

@MHRDYN7 For the quality checks, you'll need to run `make fix-copies` and then possibly `make fixup` afterwards (I would just to be safe). You can see this [from the CI logs](https://app.circleci.com/pipelines/github/huggingface/transformers/98743/workflows/232454de-524e-4c17-9967-e3de23c246fb/jobs/1311459). 

 > Moreover, what should be done for skipping the equivalence test? 

For this, in general, we would add a `@unittest.skip` decorator on the test. Typically, this means overriding the test in the specific model's tests and then adding the decorator e.g. [like here](https://github.com/huggingface/transformers/blob/1c7ebf1d6eaf0ed0fd4101fd6eb7e64601429cfe/tests/models/vipllava/test_modeling_vipllava.py#L170) 

For your proposal re weights, this might be a good option as we'd still be checking the rest of the model. When you say "remove", I'm guessing you mean from the respective state dicts?

@amyeroberts thank you. I have tried to solve all the issues. To summarize, I have finally decided to keep the ~~jax.image.resize~~ jax.image.scale_and_translate layer as it is, even though it results in slightly different output tensors compared to the pytorch model. This is because: 1. the outputs are still close enough to correctly predict the class of the input image and 2. the implementation will be exact once jax updates the resize layer to follow the pytorch conventions in the future. As a result, there was no need to skip any tests. In addition, a PR has been opened on the [hub](https://huggingface.co/facebook/dinov2-base/discussions/12) to add the flax weights. Please let me know if anything is out of order.

@MHRDYN7 Great! 

> 2. the implementation will be exact once jax updates the resize layer to follow the pytorch conventions in the future.

Is there a plan for this to happen in the future? 

> In addition, a PR has been opened on the [hub](https://huggingface.co/facebook/dinov2-base/discussions/12) to add the flax weights.

This shouldn't be necessary, all the model frameworks: TF, PyTorch, Flax should be able to load the safetensors file. 

> Is there a plan for this to happen in the future?

The [issue](https://github.com/google/jax/issues/15768) from the jax repo mentioned on my first comment, suggests that they did come up with the fix but no steps were taken and also there are no PRs related to the issue. I might just open a PR there if I can; shouldn't be hard to solve. 
>This shouldn't be necessary, all the model frameworks: TF, PyTorch, Flax should be able to load the safetensors file.

It's good to hear that. Indeed, the models don't necessarily need the .msgpack weights. I observed the flax weights on the hub for many models and thought it was a convention to add those. 

@amyeroberts thanks a lot for the review. All the tests are passing again.

@sanchit-gandhi thanks a lot for the careful review. 

A summary of the updates

- Almost all of your suggestions have been incorporated
- Unable to use `# Copied from`  Beit for the PreTrainedModel class
- ~~Facing some issues trying to use `# Copied from` in the tests, and that's why they are kept as they were~~ Now, this has been solved, everything lgtm.

> Regarding your PR description: I didn't fully understand what the issue was with the position embedding weights - you've defined them as a standard self.param, and the keys look to match those from PyTorch? Let me know if I'm missing something here!

The position embedding weights can be loaded perfectly. However, these weights are later modified according to the number of patches using F.interpolate (with bicubic mode) in torch. We can replicate this behavior with jax.image.scale_and_translate (or with image.resize), but it [seems](https://github.com/google/jax/issues/15768) that this function is slight different from torch interpolate only in case of the bicubic mode, resulting slightly different output hidden_states.

Please Note:
the default config is image_size = 224, patch_size = 16 (and interpolation of the pos_embeds is not needed in this case); 
but the pretrained weights config on hub is image_size = 518, patch_size = 14

Thanks for the detailed explanation and iterating with us @MHRDYN7! @sanchit-gandhi is off at the moment, but I can see you've addressed his comments, so I think we're OK to merge without his second review.

Final step is running the slow tests for the model before merge. Could you push an empty commit with the message `[run_slow] dinov2`? 

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_31960). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

@amyeroberts, I've pushed the required commit. Now I guess, it requires your approval for running the slow tests

@amyeroberts slow tests passed! Ready to be merged