[tests] Add Context-parallel CI tests

[kashif]

What does this PR do?

Adds two context parallel tests for the CI

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,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? 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, and
  here are tips on formatting docstrings.
• 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.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[stas00]

Looks fantastic, Kashif! Thank you for adding the missing tests. Now I don't need to worry about breaking your CP integration while integrating ALST/UlyssesSP.

[stas00]

oh, one request, since we will now have more than one context parallel, you might want to rename the file+class to include torch in it - via #41832 and huggingface/accelerate#3817 we will now have pc.cp_backend = {torch|deepspeed} torch being the default.

[stas00]

It's interesting that "hf-internal-testing/tiny-random-LlamaForCausalLM" fails here:

stderr: [rank0]:   File "/code/users/stas/github/transformers-alst-integration/src/transformers/models/llama/modeling_llama.py", line 292, in forward
stderr: [rank0]:     attn_output = self.o_proj(attn_output)
stderr: [rank0]:                   ^^^^^^^^^^^^^^^^^^^^^^^^
stderr: [rank0]:   File "/home/yak/miniconda3/envs/dev/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
stderr: [rank0]:     return self._call_impl(*args, **kwargs)
stderr: [rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
stderr: [rank0]:   File "/home/yak/miniconda3/envs/dev/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
stderr: [rank0]:     return forward_call(*args, **kwargs)
stderr: [rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
stderr: [rank0]:   File "/home/yak/miniconda3/envs/dev/lib/python3.12/site-packages/torch/nn/modules/linear.py", line 125, in forward
stderr: [rank0]:     return F.linear(input, self.weight, self.bias)
stderr: [rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
stderr: [rank0]: RuntimeError: mat1 and mat2 shapes cannot be multiplied (64x32 and 16x16)

[kashif]

yes... i got the same so i switched to "HuggingFaceTB/SmolLM-135M"

[ydshieh]

Thanks a lot. I will check on our multi-gpu CI runners too.

[stas00]

would it be possible to exercise the rest of the torch/fsdp PC config here:

parallelism_config:
  parallelism_config_dp_replicate_size: 1
  parallelism_config_dp_shard_size: 1
  parallelism_config_tp_size: 1
  parallelism_config_cp_size: 2

  --------------
  parallelism_config_cp_comm_strategy: alltoall

that is to add all config values for the Torch PC plugin, like cp_comm_strategy? it also helps with tests-as-documentation/example

[stas00]

Kashif, I'd suggest dropping the basic test, since it gets repeated almost identically in the 2nd test - will save $$ and time.

[stas00]

@kashif, I took your test as a foundation for mine and further worked on it, while also making it self-contained w/o needing to include in the already bloated repo more files.

https://github.com/huggingface/transformers/pull/41832/files#diff-bfd2f0924ca5096a9cf7bff5929081cf6552b70df8bb40632d3b11273c9554af

Please feel free to re-use, with the note that I made some tweaks to testing_utils.py that it relies on - your version of my lib is quite outdated, I have made quite a few changes over the recent years. If you want to sync back the latest version is here https://github.com/stas00/ml-engineering/blob/master/testing/testing_utils.py

[kashif]

@stas00 I've made the script in the self-contained style and incorporated your suggestions

[kashif]

@ydshieh we should ideally run this only on the multi-gpu setup, is that possible?

[stas00]

Looks great, Kashif.

[SunMarc]

Thanks for adding this ! Just a minor nit

[kashif]

@SunMarc moved to fsdp folder

[ydshieh]

@ydshieh we should ideally run this only on the multi-gpu setup, is that possible?

Yes, we have multi-gpu runners (A10). I will check how it goes with that.

[ydshieh]

@kashif The test is run and pass in our multi-gpu A10 runner, great!

[kashif]

amazing thank you @ydshieh