Add Context Parallelism support to cudnn Flash Attention

[kocchop]

Description

This PR adds Context Parallelism support to GPU Flash Attention. It is necessary to support large sequence lengths in MaxText. Right now, the support is offered through Transformer-Engine and uses an All-Gather type implementation. Note that, it requires mask type to be causal and does not work with sliding window attention yet. Also it requires transformer-engine==1.13 or above.

NEW

1. Implemented input sequence re-ordering at the data loading stage
2. Added context_parallel_load_balance flag in base.yml to turn on/off load balancing.
3. Added sequence packing flag enable_packing in base.yml and also modified the associated data processing files.
4. Added/modified unit tests for context parallelism and GPU flash attention cudnn_flash_te
5. Added sharding of parameters across CP rank

Tests

Unit test is included with the PR with the base model for 4 x a100 gpus,

Checklist

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed.

[gobbleturk]

Awesome, this looks great! Just one change to the sharding rules (with current state of this PR it looks like the old sequence parallelism is broken)

I'd appreciate it if you could also run bash code_style.sh to run our linter (I recommend saving the branch before running this just in case...)

[A9isha]

This is great! thank you so much

Curious, what kind of improvement are you observing with the loadbalancing enabled?

[A9isha]

Oh also please merge all commits into one before pushing

[kocchop]

This is great! thank you so much

Curious, what kind of improvement are you observing with the loadbalancing enabled?

for llama3-8b, with load balancing and cp=2, we're able to see around ~10% perf improvement

[kocchop]

@gobbleturk @A9isha once I get a LGTM from you, I can work on resolving the merge conflicts and also merge all the commits into one.

[gobbleturk]

LGTM! I would like either @khatwanimohit or @aireenmei to review for the data loading change part enable_packing=True

[A9isha]

LGTM - thank you!

[kocchop]

Hi @gobbleturk could you please approve so that the unit tests can run?

[A9isha]

Thank you!

[aireenmei]

LGTM for data input changes!

[gobbleturk]

All tests passed, manually adding pull ready

[gobbleturk]

pull ready might have added successfully automatically after an additional allowing step by me to run next steps, not sure