[WIP] Flax BLOOM implementation + demo
Author: haileyschoelkopf
Pull request: yes

# What does this PR do?

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

This PR will add a Flax implementation of BLOOM, and also I'd be happy to help contribute a tutorial / showcase of how to fine-tune BLOOM as well as discussed in #17703 :)


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [x] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case. --> linked above
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation). --> documentation in progress 
- [ ] Did you write any new necessary tests? --> will add once code is closer to completion


## 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.

@patrickvonplaten 

and @sanchit-gandhi @patil-suraj I believe were interested in collaborating. happy to discuss how best to do this.


Comments:
A note on the initial status of this PR:

- This first commit contains much of the code and structure of the `modeling_flax_bloom.py` file, copied from the gpt-neo Flax implementation and edited in many places already to better match the PyTorch Bloom implementation.
- There are many TODOs I've left in this file that I still need to get to. The code is still not in a runnable/finished state, 

Next steps:
- Finish implementing all methods, in particular the FlaxBloomAttention `__call__` method, until code runs (see other TODOs in file for other things that need tweaking/fixing)
- Determine how to deal with alibi tensors and how to deal with Bloom not having any hardcoded max length  
- Once code is working, start testing whether the implementation is the same as PyTorch
- Make sure tensor parallelism is working correctly / accounted for properly (see issue #17653 , this still seems to be an open issue on how best to deal with it, but bigscience/bloom-350m has TP=1 so it can be used for testing at first without worrying about TP)


Later on:
- Add unit tests once the code is working at least reasonably well!
- Make sure all functions are stateless / code works fine with `jit` - I'm relatively new to Flax/Jax so I definitely need to confirm correctness of code on this end

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_17761). All of your documentation changes will be reflected on that endpoint.

Thanks for the helpful comments, @sanchit-gandhi ! I'll do another revision through the code fixing these and adding some more things as soon as I have time.

I think that the main thing that would need to be discussed soon is how to handle [AliBi](https://arxiv.org/abs/2108.12409) for position information, since it means that there is no specific max length for BLOOM inputs. I'm not too sure yet how to account for this given things like line 176, where the causal mask is made at the max length of the model and then sliced to get the mask for shorter sequences. (One idea I had was selecting a reasonable "starting max length", then if the model gets a longer input sequence the causal mask is extended either permanently or just for that forward pass).

Would there be any issues in implementing it in the first of the two ways proposed (set to `max_length`, slicing as required)?

The problem I envision with the latter of the two approaches is that once the function is jit'd, providing a new input length to the model would result in XLA having to recompile. Each time XLA sees a new input shape, it needs to trace out (compile) the function again. So if we provide a new input shape for each forward pass, XLA will recompile every time (very slow)! The performance benefits of jit'ing a function come when we re-use a function that has already been jit'd, meaning we should try and use fixed input shapes where possible.

Yeah, the recompilation is definitely something to try to avoid! 

But the issue is that [the bigscience/bloom config](https://huggingface.co/bigscience/bloom/blob/main/config.json) doesn't have any seq_length attribute ([but bigscience/bloom-1b3 does--4096](https://huggingface.co/bigscience/bloom-1b3/blob/main/config.json)) and we want BLOOM to be able to handle sequences as long as a user wants since AliBi allows generalization to longer sequences. We could maybe just choose a reasonable default `max_length`, and then if the user passes a sequence that's too long, permanently double the size of the causal mask--this would allow for fewer recompilations, hopefully.

But I think we should keep the possibility open to using the model on very long sequences without problems--I don't know if any other models in Transformers use AliBi embeddings yet so that's a unique benefit of this model.

Let's go with that to start - we can iterate and find an optimal solution as we progress. There's also the option of asking on one of the JAX/Flax Forums to see if the framework authors have any ideas if we're stuck!

You're right, this will be the first JAX/Flax model in Transformers to use AliBi embeddings! Will be very cool having a model with no theoretical `max_len`!

Actually, I don't see a big problem with computing the `position_ids` for the embeddings on the fly if they depend only on the input length of `input_ids`
In general whenever the user passes a different input length of `input_ids` to the model will have to be recompiled it anyways so I don't see an issue with generating the position_ids and the causal_mask from the `input_ids` either no? 

Or am I misunderstanding something here?

If just generating the causal mask at every forward pass is acceptable and wouldn't incur a speed penalty, then that should work fine!

And yes, I don't think that we need to pass position_ids into the model, and we can just compute the alibi embedding within the forward pass (the pytorch implementation does this.)

sorry for the delay on this--I'll work on it in the next 2 days.

> If just generating the causal mask at every forward pass is acceptable and wouldn't incur a speed penalty, then that should work fine!
> 
> And yes, I don't think that we need to pass position_ids into the model, and we can just compute the alibi embedding within the forward pass (the pytorch implementation does this.)
> 
> sorry for the delay on this--I'll work on it in the next 2 days.

Great! Yeah, I just talked to @sanchit-gandhi offline - I think what we want to do here to only recompile when the model has to be recompiled anyways which translates into doing the folowing:

Allow `position_ids` to be passed but default them to `None` . If `None` they will be computed on the fly depending on the shape of `input_ids` and the values of `attention_mask` (the same would hold true for the causal_mask). Let me know if this doesn't make sense @haileyschoelkopf or if you have any other questions, more than happy to help :-)


Hey @haileyschoelkopf! This looks good with regards to the fused key-query-value matmuls in https://github.com/huggingface/transformers/pull/17761/commits/faddb8d446bfc8db0c1c77f29d9847a25f70b5a2! Just as a heads-up, for gradient checkpointing, you can follow the PR at https://github.com/huggingface/transformers/pull/17843. Feel free to reach out if there's anything you wish to discuss, very happy to help with any questions!

Added gradient checkpointing, thanks for the pointer @sanchit-gandhi ! 

Sorry that I haven't been able to push things forward on this PR faster, ended up being busier the past few weeks than expected... EDIT: saw the other PR. @younesbelkada , FYI, there is gradient checkpointing code on this PR now if you need it.

Thank you @haileyschoelkopf for jumping on this so quickly and getting the structure for the model in place! This PR was completed in https://github.com/huggingface/transformers/pull/18022

Let me know if there's anything else you'd like to have a go at adding in JAX/Flax! Or if you'd like to have a go at porting another model to JAX/Flax I can make some suggestions!

Thanks so much for all the helpful comments @sanchit-gandhi on this PR and apologies I wasn't able to iterate quicker on it!

If I have more time to add another JAX model I'll ping you for sure :) 

Very sorry that we rushed this PR so much @haileyschoelkopf! Very much looking forward to other PRs if you'd like :-)

Of course, will ping you if so :)