Add support for DeepSpeed sequence parallelism (Ulysses)
Author: ronald-d-rogers
Pull request: yes

# What does this PR do?

Adds support for sequence parallel to Tranformers.

**❗ The sister PR to `accelerate` (https://github.com/huggingface/accelerate/pull/3299) must be merged first ❗**

The PR mostly pulls in and updates changes made in @zeyugao's very thorough PR here:
https://github.com/huggingface/accelerate/pull/2877

And pulls in @samedjacob's PR here:
https://github.com/huggingface/transformers/pull/32305

To support changes made by @samejacobs in DeepSpeed, to make sequence parallel integration with HF cleaner here:
https://github.com/microsoft/DeepSpeed/pull/5774

And to respond to comments to the his PR to `transformers` here #32305.

Note that my testing has revealed that all of @zeyugao's changes are indeed required for this integration to work.

I've added two decorators:
- `deepspeed_ulysses_attention`: Instead of modifying the global `_flash_attention_forward` we wrap it instead as suggested by @ArthurZucker.
- `support_deepspeed_ulysses`: An extension of the concept above to add support for sequence parallel to modules by injecting required variables like `sp_group_size`.

I am not married to any of these changes, and am completely open to suggestions.

The only major addition made in my PR is to find the right place to shard the inputs for sequence parallelism to work.

I added a method to `Trainer` called `_finalize_inputs` which is meant to run right before the inputs are passed to the model. I toyed around with sharding the sequences in various parts of the data loading pipeline (see commits if you are curious), but ultimately determined that it had to be done in trainers, right before the inputs are passed to the model, otherwise libraries with custom trainers (i.e. `trl`) would have to refactor the way their trainers prepare data.

As things currently stand, **the sister PR to `accelerate` (https://github.com/huggingface/accelerate/pull/3299) must be merged first**, so that the new flags added to `HFDeepSpeedConfig` for sequence parallelism are available, but if desired we can likely make these PRs orthogonal by caching the same flags in `transformers` for now. 

So far everything appears to work (forward/backward pass tested on 4xA10s and 8xA100s). I will be doing more testing soon (i.e. testing loss is same, tuning a model & evaluating, etc.).

~~Note that we should also update `loss_utils.py`'s `fixed_cross_entropy` method to use 
[`vocab_sequence_parallel_cross_entropy`](https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/sequence/cross_entropy.py#L59) when sequence parallel is enabled, but this method currently does not allow you to propagate the args `ignore_index` and `reduction`. I will create a PR for this in DeepSpeed shortly.~~

I have also integrated DeepSpeed's Ulysses loss ([`vocab_sequence_parallel_cross_entropy`](https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/sequence/cross_entropy.py#L59)) into `loss_utils.py` and my tests are currently showing good loss curves.

I am also beginning to think we should use the `DistributedAttention` that is in `DeepSpeed` here instead of decorating the global flash attention function, similar to the way @zeyugao had it, as it seems to have a lot of functionality:
https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/sequence/layer.py#L300

I will likely test this soon as well.

<!--
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 -->

## 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#create-a-pull-request),
      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.
- [x] 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).
- [x] 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:

- text models: @ArthurZucker
- vision models: @amyeroberts, @qubvel
- speech models: @ylacombe, @eustlb
- graph models: @clefourrier

Library:

- flax: @sanchit-gandhi
- generate: @zucchini-nlp (visual-language models) or @gante (all others)
- pipelines: @Rocketknight1
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @muellerzr and @SunMarc
- chat templates: @Rocketknight1

Integrations:

- deepspeed: HF Trainer/Accelerate: @muellerzr
- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization (bitsandbytes, autogpt): @SunMarc @MekkCyber

Documentation: @stevhliu

HF projects:

- accelerate: [different repo](https://github.com/huggingface/accelerate)
- datasets: [different repo](https://github.com/huggingface/datasets)
- diffusers: [different repo](https://github.com/huggingface/diffusers)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Maintained examples (not research project or legacy):

- Flax: @sanchit-gandhi
- PyTorch: See Models above and tag the person corresponding to the modality of the example.
- TensorFlow: @Rocketknight1

 -->


Comments:
cc @sunmarc @muellerzr

cc @XuehaiPan 

Hello all, just to provide a progress update, I am currently tuning and evaluating a few models to validate that this is all working correctly.

One thing I have noticed is that the loss for sequence parallel seems to be slightly different, but in the same ballpark. I am still trying to figure out exactly why or if this matters. There seems to be a single token which is being clipped and I have yet to determine what that token is and how to preserve it.

Also, though this works with TRL's SFT without issue the memory usage for DPO seems to be way too high. I am trying to investigate what the cause of this is.

Hello all, I was busy with something else and just got back to work on this.

I have tuned and evaluated several Llama 3.1 8B models on several different tasks (classification/summarization) and am currently comparing the results of all of the runs.

So far I am definitely noticing some oddities between DeepSpeed ZeRO-3 and DeepSpeed Ulysses.

The `vocab_sequence_parallel_cross_entropy` loss as it is currently implemented/integrated doesn't appear to work at all. Hugging Face's default loss aggregation across GPUs seems to work for classification (i.e. evaluation improves over base model) but I am currently running to issues with summarization. I am currently tuning a few more runs to see if the issues goes away (with more epochs), but if not I think I will have to start digging into how to loss is implemented.

Currently from what I can tell, we aren't doing `softmax` across the whole sequence in calculating cross-entropy loss, which I think could be causing issues, but am not sure.

Either way, I will post charts and updates within a day or two.

Hi, 
thanks for the great work! I tried to get your half-finished PR to work but failed. Instead, I got this (https://github.com/princeton-nlp/ProLong) Sequence Parallel implementation to work, which integrates Ulysses-style SP into the HF trainer. Maybe you can get some inspiration there to finish the PR. 
I'm glad to try it once it is merged!


@Falko1: Thank you very much for this link. Yes, I am completely stuck trying to get other integrations to work. Maybe I will pivot to this instead.

Hope it won't be too inconvenient for you! 🤗 

@ArthurZucker: Yes I have noticed :). I think at this point, since this PR isn't actually working, I will go ahead and close this. I will likely redo all of the work in the updated Transformers, perhaps without any dependency on DeepSpeed (like ProLong).

So far I have been able to reproduce the same per vocab logits and loss with sequence parallelism turned as turned off (with slight differences) using both Tranformers' loss mechanics and DeepSpeed's provided sequence parallel loss (modified to work w/ ignore indices), but the model seems to drift after back-propagation... If I run into the same problem I guess I will create an issue and see if I can get help from the community.

We are looking into RingAttention as well could be good!

@ArthurZucker: [Picotron](https://github.com/huggingface/picotron)?

Yeah, with the attention refactor we think things should be easier now!