Context Parallelism

[rrutmann]

What does this PR do?

This PR adds end-to-end context parallel support for the GPT2 path. The current implementation requires the usage of pytorch's flash attention when enabling context parallelism

General Changes

• Applies CP-aware SDPA dispatch in attention.
• Adds trainer-side sequence sharding for inputs and targets on the CP mesh.
• Propagates CP load-balancer configuration from model setup into trainer sharding.
• Registers and validates the new GPT2 CP model variant.
• Adds focused unit coverage and an e2e CP-vs-non-CP parity test.

Breaking Changes

• ..

Checklist before submitting final PR

• My PR is minimal and addresses one issue in isolation
• I have merged the latest version of the target branch into this feature branch
• I have reviewed my own code w.r.t. correct implementation, missing type hints, proper documentation, etc.
• I have run a sample config for model training
• I have checked that all tests run through (python tests/tests.py)
• I have updated the internal changelog (CHANGELOG_DEV.md)