[PyTorch] Add context parallel support for packed dataset in THD format
You code indentation is not consistent, some places have 4 spaces, and other places have 2 spaces.
NeMo code has 4 spaces always, so please make sure all your code have 4-space indentation. Thanks for fixing the comments, it looks much better now.
The total sequence length (size of t in THD format) is a constant, right? If so, we should have some padded tokens at the end? how those padded tokens are split across different CP ranks? This PR is stale because it has been open for 14 days with no activity. Remove stale label or comment or update or this will be closed in 7 days. This PR is stale because it has been open for 14 days with no activity. Remove stale label or comment or update or this will be closed in 7 days. This PR was closed because it has been inactive for 7 days since being marked as stale. > # What does this PR do ?
> This PR adds context parallel support for packed dataset in THD format in NeMo in response to this TE PR: [NVIDIA/TransformerEngine#641](https://github.com/NVIDIA/TransformerEngine/pull/641). Currently, the TE PR requires each individual sequence length is divisible by (2*context_parallel_size).
> 
> # Changes
> * Add support to split packed dataset across different CP ranks in a load balanced way
> * Add necessary paddings to dataset during packing stage to make sure the individual sequence length is a multiple of 2*cp_size
> 
> **PR Type**:
> 
> * [x]  New Feature
> * [ ]  Bugfix
> * [ ]  Documentation

Hello, I have some problems with cp need seqlen divisibility world_size * 2. I see that you are padding data in the code, but this will cause the pad token id to enter the flash attention calculation. I don't know if this is correct.

# What does this PR do ?

This PR adds context parallel support for packed dataset in THD format in NeMo in response to this TE PR: https://github.com/NVIDIA/TransformerEngine/pull/641. Currently, the TE PR requires each individual sequence length is divisible by (2*context_parallel_size).

# Changes

- Add support to split packed dataset across different CP ranks in a load balanced way
- Add necessary paddings to dataset during packing stage to make sure the individual sequence length is a multiple of 2*cp_size

**PR Type**:
- [x] New Feature
- [ ] Bugfix
- [ ] Documentation

