Extend sequence length padding for GPT SFT to account for context parallel
State: closed | Pull request: yes
Labels: name: NLP

# What does this PR do ?

Adds a check for MegatronGPTSFTModel to ensure sequence dimension will be padded to a multiple of 16 after being split over context parallel ranks, which TE requires during fp8 training; this fixes issues like `AssertionError: FP8 execution requires 2D input matrices with height divisible by 8 and width divisible by 16, but got tensor with dims=[5504, 184]` when using context parallel with datasets that have a variable sequence length.

**Collection**: nlp

# Changelog 
- Alters `pad_seq_length_to_mult` input to GPT SFT datasets to account for context parallel

# Usage
* You can potentially add a usage example below

```python
# Add a code snippet demonstrating how to use this 
```

# Jenkins CI
To run Jenkins, a NeMo User with write access must comment `jenkins` on the PR.

# Before your PR is "Ready for review"
**Pre checks**:
- [x] Make sure you read and followed [Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md)
- [ ] Did you write any new necessary tests?
- [ ] Did you add or update any necessary documentation?
- [ ] Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  - [ ] Reviewer: Does the PR have correct import guards for all optional libraries?
  
**PR Type**:
- [ ] New Feature
- [x] Bugfix
- [ ] Documentation

If you haven't finished some of the above items you can still open "Draft" PR.


## Who can review?

Anyone in the community is free to review the PR once the checks have passed. 
[Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md) contains specific people who can review PRs to various areas.

# Additional Information
* Related to # (issue)


jenkins

Hi @vysarge , I do not totally understand what the first input dim and the second dim mean, could you please describe it? And how does it work if `sequence_parallel=True` and `context_parallel_size > 1`? Thanks.

> Hi @vysarge , I do not totally understand what the first input dim and the second dim mean, could you please describe it? And how does it work if `sequence_parallel=True` and `context_parallel_size > 1`? Thanks.

Hm, I thought sequence parallel and context parallel were mutually exclusive.  If they're used together, do CP and TP counts split on the same dimension?  Making the change you suggested should fix this case too if so, let me test it.

> > Hi @vysarge , I do not totally understand what the first input dim and the second dim mean, could you please describe it? And how does it work if `sequence_parallel=True` and `context_parallel_size > 1`? Thanks.
> 
> Hm, I thought sequence parallel and context parallel were mutually exclusive. If they're used together, do CP and TP counts split on the same dimension? Making the change you suggested should fix this case too if so, let me test it.

Sequence parallel and context parallel can work together. Sequence parallel only split sequence of LayerNorm. In addition, context parallel split sequence of the whole transformer layer. Yeah, you can test my suggested fix, I think it should work for all cases. Thanks.

> > > Hi @vysarge , I do not totally understand what the first input dim and the second dim mean, could you please describe it? And how does it work if `sequence_parallel=True` and `context_parallel_size > 1`? Thanks.
> > 
> > 
> > Hm, I thought sequence parallel and context parallel were mutually exclusive. If they're used together, do CP and TP counts split on the same dimension? Making the change you suggested should fix this case too if so, let me test it.
> 
> Sequence parallel and context parallel can work together. Sequence parallel only split sequence of LayerNorm. In addition, context parallel split sequence of the whole transformer layer. Yeah, you can test my suggested fix, I think it should work for all cases. Thanks.

Tested and this fix is working well; PR has been updated, thanks.

LGTM, thanks.

jenkins

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.