Packed Sequence #7945

cuichenx · 2023-11-27T22:01:12Z

What does this PR do ?

Support training with packed sequences for SFT and PEFT.
In this scenario, sequences in a batch are "concatenated" in the sequence length dimension, instead of "stacked" in a batch dimension.
This is shown to significantly improve the training throughput.
Packed sequences will need to be prepared in a particular way (script in #8682 )

Collection: NLP

Changelog

new GPTSFTPackedDataset class to handle the new dataset format
pass cu_seqlens and qkv_format to the forward() function of mcore model

Usage

Note: only train_ds is packed. Validation and test datasets are unchanged.

python /opt/NeMo/examples/nlp/language_modeling/tuning/megatron_gpt_peft_tuning.py \
...
+model.data.train_ds.packed_sequence=True

Before your PR is "Ready for review"

Pre checks:

Make sure you read and followed Contributor guidelines
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
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 contains specific people who can review PRs to various areas.

Additional Information

Related to # (issue)

Signed-off-by: Chen Cui <[email protected]>

…ked_sft

for more information, see https://pre-commit.ci

Signed-off-by: Chen Cui <[email protected]>

for more information, see https://pre-commit.ci

Signed-off-by: Chen Cui <[email protected]>

for more information, see https://pre-commit.ci

Signed-off-by: Chen Cui <[email protected]>

…ked_sft

Signed-off-by: Chen Cui <[email protected]>

for more information, see https://pre-commit.ci

Signed-off-by: Chen Cui <[email protected]>

Modify GPTSFTPackedDataset to respond to pad_to_max_length setting Signed-off-by: Chen Cui <[email protected]>

ericharper

LGTM. Thanks!

ericharper

Sorry ignore my previous approval.

Could you please add unit tests?

See tests/collections/nlp/test_chat_sft_dataset.py as an example

ericharper · 2023-12-05T04:07:11Z

jenkins

Signed-off-by: Chen Cui <[email protected]>

for more information, see https://pre-commit.ci

ericharper · 2023-12-15T20:52:54Z

jenkins

ericharper

LGTM. Thanks for adding the unit test.

@vysarge

* support packed dataset Signed-off-by: Chen Cui <[email protected]> * support packed dataset Signed-off-by: Chen Cui <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix collate_fn bug for TP > 1 Signed-off-by: Chen Cui <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * make packed dataset work Signed-off-by: Chen Cui <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix nan bug Signed-off-by: Chen Cui <[email protected]> * support answer only loss Signed-off-by: Chen Cui <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * account for padding in cu_seqlens during dataloading for attn kernel Signed-off-by: Chen Cui <[email protected]> * fix path for answer_only_loss = false Signed-off-by: Chen Cui <[email protected]> * @vysarge Modify GPTSFTPackedDataset to respond to pad_to_max_length setting Signed-off-by: Chen Cui <[email protected]> * add unit test Signed-off-by: Chen Cui <[email protected]> * fix loss mask bug for answer_only_loss=False Signed-off-by: Chen Cui <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Chen Cui <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Eric Harper <[email protected]> Signed-off-by: Piotr Żelasko <[email protected]>

@vysarge

* support packed dataset Signed-off-by: Chen Cui <[email protected]> * support packed dataset Signed-off-by: Chen Cui <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix collate_fn bug for TP > 1 Signed-off-by: Chen Cui <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * make packed dataset work Signed-off-by: Chen Cui <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix nan bug Signed-off-by: Chen Cui <[email protected]> * support answer only loss Signed-off-by: Chen Cui <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * account for padding in cu_seqlens during dataloading for attn kernel Signed-off-by: Chen Cui <[email protected]> * fix path for answer_only_loss = false Signed-off-by: Chen Cui <[email protected]> * @vysarge Modify GPTSFTPackedDataset to respond to pad_to_max_length setting Signed-off-by: Chen Cui <[email protected]> * add unit test Signed-off-by: Chen Cui <[email protected]> * fix loss mask bug for answer_only_loss=False Signed-off-by: Chen Cui <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Chen Cui <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Eric Harper <[email protected]> Signed-off-by: Sasha Meister <[email protected]>

@vysarge

* support packed dataset Signed-off-by: Chen Cui <[email protected]> * support packed dataset Signed-off-by: Chen Cui <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix collate_fn bug for TP > 1 Signed-off-by: Chen Cui <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * make packed dataset work Signed-off-by: Chen Cui <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix nan bug Signed-off-by: Chen Cui <[email protected]> * support answer only loss Signed-off-by: Chen Cui <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * account for padding in cu_seqlens during dataloading for attn kernel Signed-off-by: Chen Cui <[email protected]> * fix path for answer_only_loss = false Signed-off-by: Chen Cui <[email protected]> * @vysarge Modify GPTSFTPackedDataset to respond to pad_to_max_length setting Signed-off-by: Chen Cui <[email protected]> * add unit test Signed-off-by: Chen Cui <[email protected]> * fix loss mask bug for answer_only_loss=False Signed-off-by: Chen Cui <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Chen Cui <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Eric Harper <[email protected]>

cuichenx added 4 commits November 22, 2023 13:43

support packed dataset

9a107f2

Signed-off-by: Chen Cui <[email protected]>

support packed dataset

74ac00e

Signed-off-by: Chen Cui <[email protected]>

Merge remote-tracking branch 'origin/chcui/packed_sft' into chcui/pac…

518130f

…ked_sft

Merge branch 'main' into chcui/packed_sft

be77d93

github-actions bot added the NLP label Nov 27, 2023

pre-commit-ci bot and others added 13 commits November 27, 2023 22:02

[pre-commit.ci] auto fixes from pre-commit.com hooks

ef75189

for more information, see https://pre-commit.ci

fix collate_fn bug for TP > 1

e11c696

Signed-off-by: Chen Cui <[email protected]>

[pre-commit.ci] auto fixes from pre-commit.com hooks

a07c5d6

for more information, see https://pre-commit.ci

make packed dataset work

6dc84be

Signed-off-by: Chen Cui <[email protected]>

[pre-commit.ci] auto fixes from pre-commit.com hooks

2d5238d

for more information, see https://pre-commit.ci

fix nan bug

801fe7a

Signed-off-by: Chen Cui <[email protected]>

Merge remote-tracking branch 'origin/chcui/packed_sft' into chcui/pac…

f88a2bb

…ked_sft

support answer only loss

6b23131

Signed-off-by: Chen Cui <[email protected]>

[pre-commit.ci] auto fixes from pre-commit.com hooks

cdcaf9a

for more information, see https://pre-commit.ci

account for padding in cu_seqlens during dataloading for attn kernel

a3d88f7

Signed-off-by: Chen Cui <[email protected]>

fix path for answer_only_loss = false

0faa0f4

Signed-off-by: Chen Cui <[email protected]>

@vysarge

edc2421

Modify GPTSFTPackedDataset to respond to pad_to_max_length setting Signed-off-by: Chen Cui <[email protected]>

Merge branch 'main' into chcui/packed_sft

d58d723

cuichenx marked this pull request as ready for review December 5, 2023 01:10

ericharper approved these changes Dec 5, 2023

View reviewed changes

Merge branch 'main' into chcui/packed_sft

09107a8

ericharper requested changes Dec 5, 2023

View reviewed changes

vysarge mentioned this pull request Dec 5, 2023

Token count and sequence length logging for MegatronGPTSFTModel #7973

Closed

8 tasks

cuichenx and others added 4 commits December 12, 2023 08:42

add unit test

cdb0961

Signed-off-by: Chen Cui <[email protected]>

fix loss mask bug for answer_only_loss=False

7303f2d

Signed-off-by: Chen Cui <[email protected]>

Merge branch 'main' into chcui/packed_sft

54cd352

[pre-commit.ci] auto fixes from pre-commit.com hooks

b7e722e

for more information, see https://pre-commit.ci

cuichenx requested a review from ericharper December 12, 2023 22:59

Merge branch 'main' into chcui/packed_sft

38fbc59

ericharper approved these changes Dec 15, 2023

View reviewed changes

ericharper merged commit ed0f681 into main Dec 16, 2023
15 checks passed

ericharper deleted the chcui/packed_sft branch December 16, 2023 01:26

cuichenx mentioned this pull request Mar 16, 2024

Add script for packed sequence data preparation #8682

Merged

8 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Packed Sequence #7945

Packed Sequence #7945

cuichenx commented Nov 27, 2023 •

edited

Loading

ericharper left a comment

ericharper left a comment

ericharper commented Dec 5, 2023

ericharper commented Dec 15, 2023

ericharper left a comment

Packed Sequence #7945

Packed Sequence #7945

Conversation

cuichenx commented Nov 27, 2023 • edited Loading

What does this PR do ?

Changelog

Usage

Before your PR is "Ready for review"

Who can review?

Additional Information

ericharper left a comment

Choose a reason for hiding this comment

ericharper left a comment

Choose a reason for hiding this comment

ericharper commented Dec 5, 2023

ericharper commented Dec 15, 2023

ericharper left a comment

Choose a reason for hiding this comment

cuichenx commented Nov 27, 2023 •

edited

Loading