HF Trainer: ALST/Ulysses sequence parallelism integration via HF Accelerate
Author: stas00
Pull request: yes

Integrates HF Accelerate's support for ALST/Ulysses sequence parallelism https://github.com/huggingface/accelerate/pull/3817 into HF Trainer

TODO:

- [x] docs - no idea where? the FSDP/CP is not documented, or any parallelism for that matter.
- [x] tests
- [x] need to merge https://github.com/huggingface/accelerate/pull/3817
- [x] need to wait for a new HF Accelerate release and to use its version in compatibility checks in the code 1.11.1 most likely
- [x] need to wait for a merge of https://github.com/deepspeedai/DeepSpeed/pull/7649
- [x] need to wait for a new Deepspeed release after above is merged



Comments:
I have a hard time finding where I can add the documentation for this new backend. Context parallelism isn't documented anywhere in HF Trainer - how can users discover it? If I'm missing it, could you please point me to where I should extend the documentation? Thank you.

Same story with CP tests - there are none :( so need to figure out how to write some.

OK, the next issue with the existing integration of CP/FSDP. What does the following mean?
```
 $ sometrainerscript.py --help
 [...]
 --parallelism_config PARALLELISM_CONFIG, --parallelism-config PARALLELISM_CONFIG
```
this arg tells users absolutely nothing about what value(s) to pass to `--parallelism_config`. If it's not meant to be in CLI args and can only be used explicitly by writing code perhaps it shouldn't be listed in `--help` or at least say that it has to be coded?

cc @sunmarc

@SunMarc, this is ready for a review. Tests fail because they need the accelerate PR https://github.com/huggingface/accelerate/pull/3817

I just didn't know where to update docs since parallelism doesn't seem to be documented here at all. Please correct me if I'm wrong.

Thanks to @kashif with the test.

Thanks @SunMarc I did an initial run with default settings in TRL and these branches, and it all worked nicely (apart from model saving due to not updating deepspeed i think). I will check the config that uses a bespoke `compute_loss` in the sft trainer, thanks for the heads up!

> @SunMarc, this is ready for a review. Tests fail because they need the accelerate PR https://github.com/huggingface/accelerate/pull/3817
> 
> I just didn't know where to update docs since parallelism doesn't seem to be documented here at all. Please correct me if I'm wrong.
> 
> Thanks to @kashif with the test.

We indeed do not have docs related to that. @kashif as you added cp support in trainer, would you be willing to add some docs around that.
@stas00, I think we can either update deepspeed docs https://huggingface.co/docs/transformers/main/en/deepspeed and/or create a new docs called contextparallel like in accelerate. 

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

@SunMarc docs on TRL side: https://github.com/huggingface/trl/pull/4420

i can look at the docs on the accelerate side as well... 

> @stas00, I think we can either update deepspeed docs https://huggingface.co/docs/transformers/main/en/deepspeed and/or create a new docs called contextparallel like in accelerate.

Ideally we would have a dedicated doc like you suggested, which could then link into deepspeed for nuances as one way to do that. The key is for the user to quickly understand what's possible, thus a single context parallel entry point doc would be very useful to users.

@stas00 @SunMarc added initial docs https://github.com/stas00/transformers/pull/3

Thanks a lot, Kashif - merged your docs.

deepspeed  1.18.2 is out BTW

Thanks Kashif, yes, I already updated the OP. I need to find time to do a massive adaptation to `s/cp/sp/` since it's wrong to call SP CP and it will prevent future combination of SP+CP, since the 2 are unrelated to each other, while seemingly doing the same thing. Marc has already given a green light, I was just busy with other work. 

@SunMarc, `s/cp/sp/` is done - this is ready for your review. Also made some corresponding tweaks on the Accelerate size.

We would need to merge the Accelerate PR first, then make a new release there, then anchor this PR on that new release.

Can I quickly mention 2 issues I discovered while trying out the new sequence-parallelism:
1. Trainer.get_cp_size() currently doesn't take pc.sp_size into account, and hence total_train_batch_size is not correct;
2. After accelerator.deepspeed_ulysses_dl_adapter() gets called, len_dataloader is not updated and hence steps_in_epoch is not correct.
@stas00 Hope this helps and thanks for the great work!

Thanks a lot for trying it out and spotting problems, @zhangwj618. I should have time to look at it at the end of the day, unless you already have the fixes and would like to propose them. 

@stas00  I have fixes for the issue from @zhangwj618 as suggestions above for your consideration

@zhangwj618, could you please re-test with the current HEAD? Kashif fixed both issues.

@kashif, how do we communicate that the DL should still feed unique data streams on each rank. That's why I was trying to say with `DP=SP=world_size`, since the sample sharding happens in the DL Adapter, unlike TP where each rank has to feed the same sample - which would be very wrong here.



@stas00 good point! Let me see where to add the important difference.

> @zhangwj618, could you please re-test with the current HEAD? Kashif fixed both issues.

I tested the current code, and realized that max_steps (and state.max_steps) also has the same problem of not taking sp_size into account, affecting ProgressCallback.training_bar and potentially a few other things. So maybe the best place to update len_dataloader with sp_size is inside set_initial_training_values(), not after?

@zhangwj618, just thinking, should the SP affect the number of training/optimizer steps?  The SP adapter multiplies dataloader length because each rank needs `sp_world_size` iterations to process all sequence chunks. But does this mean we should take `sp_world_size` times more optimizer steps? What do you think?

@kashif imo yes, specifically in the epoch-based max_steps calculation. Because with sp the dp_world_size and the global batch_size become 1 / sp_size times smaller (c.f. get_total_train_batch_size()), hence max_steps needs to be sp_size times bigger to go through the same number of epochs. Think of decreasing per_device_train_batch_size in the pure dp scenario, which would result in more training steps, increasing sp_size alone (keeping per_device_train_batch_size fixed) has the same effect. To keep the training steps constant, the per_device_train_batch_size need to be manually adjusted to sp_times bigger (which restores the original global batch_size when multiplied by 1 / sp_size).


Another way to look at it is that [UlyssesSPDataLoaderAdapter](https://github.com/deepspeedai/DeepSpeed/blob/64c0052fa08438b4ecf4cae30af15091a92d2108/deepspeed/runtime/sequence_parallel/ulysses_sp.py#L442) yields sp_size times more data batches than the original DL, thus requires more steps to iterate through.

@zhangwj618 ok I added a fix in this PR: https://github.com/stas00/transformers/pull/4 if you want to double check on your end you can manually add the change

> @zhangwj618 ok I added a fix in this PR: [stas00#4](https://github.com/stas00/transformers/pull/4) if you want to double check on your end you can manually add the change

Yep, from my side, with this PR, the number of steps looks correct.

@stas00 There is a potential issue inside _deepspeed_sp_compute_loss(), that on some rank(s) good_tokens could be 0 and loss would be nan. Skipping these ranks when suming the total_loss can fix this problem. I did some tests and can confirm that, after fixing this, the loss convergence is aligned between sp+dp and dp-only setups (using the same global batch_size).

Indeed, in our reference implementation the nan's were already avoided earlier so weighted sum was safe:
https://github.com/snowflakedb/ArcticTraining/blob/d9328917d77fd888e80ff368f273447aa7c0ff09/arctic_training/trainer/sft_trainer.py#L129-L139

Looks like Kashif has already fixed it here:
https://github.com/huggingface/transformers/pull/41832/commits/6ae2bb0b2caf0e2ebe261605604db4daab75c535

Is there a mismatch between the docs and the code? `docs/source/en/deepspeed.md` says:
"By default, when you only configure `sp_size`, DP is automatically calculated as `dp_size = world_size / sp_size`."
However, when I run the code with sp_size != world_size, I get [this](https://github.com/stas00/accelerate/blob/a4005e7265575a31280e6db7ae909ded9ea83328/src/accelerate/parallelism_config.py#L353) error, unless I specify dp_replicate_size manully.


@zhangwj618 the doc is wrong... my bad let me fix it!

Thanks a lot @stas00 for your work 🤗 

super! Thanks a lot to Marc and Kashif for help with integration and Weijie Zhang for being the first early adopter!