[WIP] Making modifications Open source that are live for `BLOOM` inference.
Author: Narsil
Pull request: yes

# What does this PR do?

Many things but the biggest are:

- TP enabled model (need to pass around a `ProcessGroup` everywhere.
  There were some discussions back&forth but currently this is not breaking
  anything, and at least it makes everything quite explicit to load the model
  in a sharded fashion

- Adding 1 custom kernel. Followed roughly DeformatDETR way of distributing this
  Exception is that the custom kernel is OPT-in, it's not loaded by default
  (The custom kernel does not support backward for instance).

<!--
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.
-->

<!-- Remove if not applicable -->

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](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [ ] 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.
- [ ] 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).
- [ ] 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.

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- albert, bert, xlm: @LysandreJik
- blenderbot, bart, marian, pegasus, encoderdecoder,  t5: @patrickvonplaten, @patil-suraj
- longformer, reformer, transfoxl, xlnet: @patrickvonplaten
- fsmt: @stas00
- funnel: @sgugger
- gpt2: @patrickvonplaten, @LysandreJik
- rag: @patrickvonplaten, @lhoestq
- tensorflow: @LysandreJik

Library:

- benchmarks: @patrickvonplaten
- deepspeed: @stas00
- ray/raytune: @richardliaw, @amogkam
- text generation: @patrickvonplaten
- tokenizers: @n1t0, @LysandreJik
- trainer: @sgugger
- pipelines: @LysandreJik

Documentation: @sgugger

HF projects:

- datasets: [different repo](https://github.com/huggingface/datasets)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Examples:

- maintained examples (not research project or legacy): @sgugger, @patil-suraj
- research_projects/bert-loses-patience: @JetRunner
- research_projects/distillation: @VictorSanh

 -->

Comments:
@thomasw21 I tagged you to maybe get advice on the `generation` modifications.

Should they get merged back into `main`  ? (They do seem necessary).

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

> @thomasw21 I tagged you to maybe get advice on the generation modifications.
Should they get merged back into main ? (They do seem necessary).

If I remember correctly the only thing that should be needed is the logit post processor.

@sgugger @LysandreJik 

We'd like some input as we don't necessarily agree with @thomasw21.
Bascially this PR intends to be:
- No breaking change
- Enable users to use TP with bloom with much less effort, even if in a slightly more contrived way
- Enable the same performance we get on production (be as open source as possible). This is where we disagree.

@thomasw21 thinks, that adding the custom kernel to transformer `main` forces us to be backward compatible for eternity, creates potential headaches because the kernel has only been tested on A100 and might have some unforeseen caveats (1 we know is that it's limited to 4096 tokens)

My point, is that we should strive to be open, and so making everything we did accessible should be a goal. Now I agree that maintaining that kernel should NOT be in scope, but I argue that:
- Enabling it actively by users (User has to write ` bloom.use_custom_kernel()` in order to use it)
- Making a proper warning when using that function
- Eventually marking this function as private.
Would enable us to merge it and still allow us to break it whenever we want. We obviously write that as `unstable`/`beta`.

One added benefit in the back of my mind, is that allowing us to ship custom kernels could enable us (in all our ecosystem) to boost some performance where needed. Making it core `transformers` not necessarily a goal, but knowing how to enable it seems OK.

`torch.fx` seems unhappy seems not playing well with `torch.jit.script`.

First things first, I strongly disagree with most of the modifications done in the modeling code, which make the model code less readable. I think of:
- paths with a test using `if fused_bloom_attention_cuda`
- model gaining a `process_group` argument at init
- use of `tp_rank` attribute

Keep in mind that we do not let users add a flag to select if the layernorm should be applied at the beginning or the end of the block and request a new architecture instead for instance, or Thom's comments on why the Mistral code in GPT-2 should never have been merged.

This should either:
- be a "fast" modeling file of its own in the same bloom folder
- a research project of its own (which would have the advantage of not being constrained by any BC considerations)

> be a "fast" modeling file of its own in the same bloom folder

That seems reasonable. It's not `fast` vs `non-fast` to be clear. It's TP vs single GPU code difference.
And since TP is the best way to get latency optimizations it would feel quite nice if it was included in `transformers`.
Actually anywhere in the HF ecosystem would be nice, as long as we could point out to our code and say it's "there".

But since it's been 3 months and we still haven't figured it out, I decided to go in that fashion. Keeping it in a remote branch is ok, but I feel like a fork (https://github.com/huggingface/transformers_bloom_parallel/issues/8) 
 
**Separate file is all good for me !**

Btw, we don't know of anyway to enable TP without modifying the code. 
@thomasw21 checked out torch.fx but it seemed to be a pain (and it was the most successful approach). 

> paths with a test using if fused_bloom_attention_cuda

Would you have 3 files (regular, TP, TP + custom kernel) ? seems like a stretch but fine to me.
What is so bad about this particular if ? It's very similar to this:
https://github.com/huggingface/transformers/blob/main/src/transformers/models/deformable_detr/modeling_deformable_detr.py#L662-L675

The custom kernel can go. It's a shame imo, and it would be nice if it was easily pluggable, but I can understand the logic.

> model gaining a process_group argument at init

New modeling file would solve that for sure.

I think two files is good, one for the "generic" model code and one with custom improvements for TP/TP+custom CUDA kernel. I mainly want to avoid a researcher coming to take the BLOOM code and being annoyed at all the special paths for TP.

> What is so bad about this particular if ? It's very similar to this:
https://github.com/huggingface/transformers/blob/main/src/transformers/models/deformable_detr/modeling_deformable_detr.py#L662-L675

We want to avoid those kinds of path that hurt readability, at least in mainstream models. Deformable DETR didn't use to have those paths (at first it was custom CUDA kernel only) and is not a mainstream model (maybe in the future?).

In any case, most of the code can be copied in the new modeling file and the copied from statements will insure they stay up to date with the original.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.