context parallelism #7739

xrennvidia · 2023-10-17T18:11:50Z

What does this PR do ?

GPT training with long-context input (e.g., sequence length of 16, 32K, 64K) can easily overflow GPU memory with huge activations. Context parallelism splits long-context input along the dimension of sequence length, parallelizes partitioned sequence segments among multiple GPUs. In this way, each GPU only needs to store activations of a part of sequence length, so that we can avoid memory overflow.

Collection: [Note which collection this PR will affect]

Changelog

Add specific line by line info of high level changes in this PR.

Usage

You can potentially add a usage example below

# Add a code snippet demonstrating how to use this

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: xren <[email protected]>

Signed-off-by: Xiaowei Ren <[email protected]>

Signed-off-by: xren <[email protected]>

stu1130 · 2023-12-09T00:33:57Z

Hey @xrennvidia during the checkpoint stage, we ran into

 self.trainer.strategy.save_checkpoint(_checkpoint, filepath, storage_options=storage_options)
  File "/workspace/src/NeMo/nemo/collections/nlp/parts/nlp_overrides.py", line 305, in save_checkpoint
    dist_checkpointing.save(sharded_state_dict=checkpoint, checkpoint_dir=checkpoint_dir)
  File "/workspace/src/Megatron-LM/megatron/core/dist_checkpointing/serialization.py", line 221, in save
    validate_sharding_integrity(sharded_tensors)
  File "/workspace/src/Megatron-LM/megatron/core/dist_checkpointing/serialization.py", line 278, in validate_sharding_integrity
    _validate_sharding_for_key(shardings)
  File "/workspace/src/Megatron-LM/megatron/core/dist_checkpointing/serialization.py", line 316, in _validate_sharding_for_key
    raise CheckpointingException(f'Invalid access pattern for {rank_sharding[0][1]}')
megatron.core.dist_checkpointing.core.CheckpointingException: Invalid access pattern for ShardedTensor(key='optimizer.state.exp_avg.model.embedding.word_embeddings.weight')

You should be able to reproduce the issue on CP=2 and small GPT model. Let me know if you need more details on it.

Signed-off-by: Xiaowei Ren <[email protected]>

stu1130 · 2023-12-19T01:14:41Z

nemo/collections/nlp/models/language_modeling/megatron_base_model.py

+        """
+        cp_stream = torch.cuda.Stream()
+
+        for module in self.get_gpt_module_list():


couldn't find method get_gpt_module_list. Is it get_model_module_list?

please pull the latest commit. this stale code.

Thanks. Did you mean latest NeMo context parallel commit or Megatron-LM/TransformerEngine?

latest NeMo context parallel commit

Signed-off-by: Xiaowei Ren <[email protected]>

xrennvidia · 2024-01-05T02:00:18Z

jenkins

xrennvidia · 2024-01-05T19:01:53Z

jenkins

stu1130 · 2024-01-09T06:50:19Z

Hey @xrennvidia I am using Megatron-LM (mcore r0.4.0). If I pull the latest change in the PR, would I also need to cherry-pick NVIDIA/Megatron-LM@5eaa937 to r0.4.0? Again, thanks for developing the feature, it really benefits our use case a lot!

xrennvidia · 2024-01-09T06:58:33Z

Hey @xrennvidia I am using Megatron-LM (mcore r0.4.0). If I pull the latest change in the PR, would I also need to cherry-pick NVIDIA/Megatron-LM@5eaa937 to r0.4.0? Again, thanks for developing the feature, it really benefits our use case a lot!

Hi @stu1130 , very happy to know this is helpful :) . If you want to run with PP > 1, you need to cherry-pick it.

Also FYI, I have a fix for your issue at here. I think the fix should be merged to MLM main branch soon, maybe tomorrow.

xrennvidia · 2024-01-09T23:21:55Z

jenkins

ericharper

LGTM. Thanks!

* make nemo recognize sequence_parallel_size Signed-off-by: xren <[email protected]> * add helper functions to set up SP running in TE Signed-off-by: xren <[email protected]> * slice seq length for a specific rank Signed-off-by: Xiaowei Ren <[email protected]> * fix data_parallel_size calculation Signed-off-by: Xiaowei Ren <[email protected]> * minor change Signed-off-by: Xiaowei Ren <[email protected]> * add missing argument of self Signed-off-by: Xiaowei Ren <[email protected]> * pass sp_global_ranks to TE transformer layer Signed-off-by: Xiaowei Ren <[email protected]> * fix nsys setting Signed-off-by: Xiaowei Ren <[email protected]> * fix seq_len calculation Signed-off-by: xren <[email protected]> * fix attn_mask split across seq-length dim Signed-off-by: xren <[email protected]> * code update of input split Signed-off-by: xren <[email protected]> * fix loss calculation Signed-off-by: xren <[email protected]> * fix loss_mask_sum calculation Signed-off-by: xren <[email protected]> * fix losss calculation Signed-off-by: xren <[email protected]> * rename sequence_parallelism to context_parallelism Signed-off-by: xren <[email protected]> * minor change Signed-off-by: xren <[email protected]> * fix loss_mask_sum calculation Signed-off-by: xren <[email protected]> * make sure do not call megatron-core parallel_state while cp_size is 1 Signed-off-by: xren <[email protected]> * slice position embedding for different CP rank Signed-off-by: xren <[email protected]> * fix mising property decorator Signed-off-by: xren <[email protected]> * typo fix Signed-off-by: xren <[email protected]> * fix rpe_bias CP slicing Signed-off-by: xren <[email protected]> * code style fix Signed-off-by: xren <[email protected]> * fix loss_mask_sum calculation Signed-off-by: xren <[email protected]> * do not load attention mask if it's not needed Signed-off-by: Xiaowei Ren <[email protected]> * bug fix Signed-off-by: xren <[email protected]> * fix ubuf size with CP > 1 Signed-off-by: Xiaowei Ren <[email protected]> * address naming confusion of mixed dp and cp Signed-off-by: xren <[email protected]> * rewrite cp code by assuming with_context_parallel=False Signed-off-by: xren <[email protected]> * pop context_parallel from dist opt kwargs Signed-off-by: xren <[email protected]> * make sure amax reduction group is aware of context parallelism Signed-off-by: xren <[email protected]> * remove use_fp8 from initialize_model_parallel Signed-off-by: xren <[email protected]> * make implementaitons of setup_transformer_engine_tp_groups and setup_transformer_engine_cp_running consistent Signed-off-by: xren <[email protected]> * cp function renaming Signed-off-by: xren <[email protected]> * make loss logging broadcast aware of cp Signed-off-by: xren <[email protected]> * fix a typo Signed-off-by: Xiaowei Ren <[email protected]> * var name fix Signed-off-by: Xiaowei Ren <[email protected]> * import transformer layer specs from MCore Signed-off-by: Xiaowei Ren <[email protected]> * upgrade MCore version Signed-off-by: Xiaowei Ren <[email protected]> * add add context_parallel into the kwargs of dist opt Signed-off-by: Xiaowei Ren <[email protected]> * remove redundant cp check Signed-off-by: Xiaowei Ren <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * code style fix Signed-off-by: Xiaowei Ren <[email protected]> * recover docker file Signed-off-by: Xiaowei Ren <[email protected]> * fix seq_length of CP Signed-off-by: Xiaowei Ren <[email protected]> * recover seq-length which has been fixed in mcore Signed-off-by: Xiaowei Ren <[email protected]> * function name fix Signed-off-by: Xiaowei Ren <[email protected]> --------- Signed-off-by: xren <[email protected]> Signed-off-by: Xiaowei Ren <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

This reverts commit 58d6bce.

This reverts commit a16f6c6.

This reverts commit 58d6bce.

This reverts commit a16f6c6.

This reverts commit 58d6bce. Signed-off-by: Jan Baczek <[email protected]>

* make nemo recognize sequence_parallel_size Signed-off-by: xren <[email protected]> * add helper functions to set up SP running in TE Signed-off-by: xren <[email protected]> * slice seq length for a specific rank Signed-off-by: Xiaowei Ren <[email protected]> * fix data_parallel_size calculation Signed-off-by: Xiaowei Ren <[email protected]> * minor change Signed-off-by: Xiaowei Ren <[email protected]> * add missing argument of self Signed-off-by: Xiaowei Ren <[email protected]> * pass sp_global_ranks to TE transformer layer Signed-off-by: Xiaowei Ren <[email protected]> * fix nsys setting Signed-off-by: Xiaowei Ren <[email protected]> * fix seq_len calculation Signed-off-by: xren <[email protected]> * fix attn_mask split across seq-length dim Signed-off-by: xren <[email protected]> * code update of input split Signed-off-by: xren <[email protected]> * fix loss calculation Signed-off-by: xren <[email protected]> * fix loss_mask_sum calculation Signed-off-by: xren <[email protected]> * fix losss calculation Signed-off-by: xren <[email protected]> * rename sequence_parallelism to context_parallelism Signed-off-by: xren <[email protected]> * minor change Signed-off-by: xren <[email protected]> * fix loss_mask_sum calculation Signed-off-by: xren <[email protected]> * make sure do not call megatron-core parallel_state while cp_size is 1 Signed-off-by: xren <[email protected]> * slice position embedding for different CP rank Signed-off-by: xren <[email protected]> * fix mising property decorator Signed-off-by: xren <[email protected]> * typo fix Signed-off-by: xren <[email protected]> * fix rpe_bias CP slicing Signed-off-by: xren <[email protected]> * code style fix Signed-off-by: xren <[email protected]> * fix loss_mask_sum calculation Signed-off-by: xren <[email protected]> * do not load attention mask if it's not needed Signed-off-by: Xiaowei Ren <[email protected]> * bug fix Signed-off-by: xren <[email protected]> * fix ubuf size with CP > 1 Signed-off-by: Xiaowei Ren <[email protected]> * address naming confusion of mixed dp and cp Signed-off-by: xren <[email protected]> * rewrite cp code by assuming with_context_parallel=False Signed-off-by: xren <[email protected]> * pop context_parallel from dist opt kwargs Signed-off-by: xren <[email protected]> * make sure amax reduction group is aware of context parallelism Signed-off-by: xren <[email protected]> * remove use_fp8 from initialize_model_parallel Signed-off-by: xren <[email protected]> * make implementaitons of setup_transformer_engine_tp_groups and setup_transformer_engine_cp_running consistent Signed-off-by: xren <[email protected]> * cp function renaming Signed-off-by: xren <[email protected]> * make loss logging broadcast aware of cp Signed-off-by: xren <[email protected]> * fix a typo Signed-off-by: Xiaowei Ren <[email protected]> * var name fix Signed-off-by: Xiaowei Ren <[email protected]> * import transformer layer specs from MCore Signed-off-by: Xiaowei Ren <[email protected]> * upgrade MCore version Signed-off-by: Xiaowei Ren <[email protected]> * add add context_parallel into the kwargs of dist opt Signed-off-by: Xiaowei Ren <[email protected]> * remove redundant cp check Signed-off-by: Xiaowei Ren <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * code style fix Signed-off-by: Xiaowei Ren <[email protected]> * recover docker file Signed-off-by: Xiaowei Ren <[email protected]> * fix seq_length of CP Signed-off-by: Xiaowei Ren <[email protected]> * recover seq-length which has been fixed in mcore Signed-off-by: Xiaowei Ren <[email protected]> * function name fix Signed-off-by: Xiaowei Ren <[email protected]> --------- Signed-off-by: xren <[email protected]> Signed-off-by: Xiaowei Ren <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Signed-off-by: Sasha Meister <[email protected]>

* make nemo recognize sequence_parallel_size Signed-off-by: xren <[email protected]> * add helper functions to set up SP running in TE Signed-off-by: xren <[email protected]> * slice seq length for a specific rank Signed-off-by: Xiaowei Ren <[email protected]> * fix data_parallel_size calculation Signed-off-by: Xiaowei Ren <[email protected]> * minor change Signed-off-by: Xiaowei Ren <[email protected]> * add missing argument of self Signed-off-by: Xiaowei Ren <[email protected]> * pass sp_global_ranks to TE transformer layer Signed-off-by: Xiaowei Ren <[email protected]> * fix nsys setting Signed-off-by: Xiaowei Ren <[email protected]> * fix seq_len calculation Signed-off-by: xren <[email protected]> * fix attn_mask split across seq-length dim Signed-off-by: xren <[email protected]> * code update of input split Signed-off-by: xren <[email protected]> * fix loss calculation Signed-off-by: xren <[email protected]> * fix loss_mask_sum calculation Signed-off-by: xren <[email protected]> * fix losss calculation Signed-off-by: xren <[email protected]> * rename sequence_parallelism to context_parallelism Signed-off-by: xren <[email protected]> * minor change Signed-off-by: xren <[email protected]> * fix loss_mask_sum calculation Signed-off-by: xren <[email protected]> * make sure do not call megatron-core parallel_state while cp_size is 1 Signed-off-by: xren <[email protected]> * slice position embedding for different CP rank Signed-off-by: xren <[email protected]> * fix mising property decorator Signed-off-by: xren <[email protected]> * typo fix Signed-off-by: xren <[email protected]> * fix rpe_bias CP slicing Signed-off-by: xren <[email protected]> * code style fix Signed-off-by: xren <[email protected]> * fix loss_mask_sum calculation Signed-off-by: xren <[email protected]> * do not load attention mask if it's not needed Signed-off-by: Xiaowei Ren <[email protected]> * bug fix Signed-off-by: xren <[email protected]> * fix ubuf size with CP > 1 Signed-off-by: Xiaowei Ren <[email protected]> * address naming confusion of mixed dp and cp Signed-off-by: xren <[email protected]> * rewrite cp code by assuming with_context_parallel=False Signed-off-by: xren <[email protected]> * pop context_parallel from dist opt kwargs Signed-off-by: xren <[email protected]> * make sure amax reduction group is aware of context parallelism Signed-off-by: xren <[email protected]> * remove use_fp8 from initialize_model_parallel Signed-off-by: xren <[email protected]> * make implementaitons of setup_transformer_engine_tp_groups and setup_transformer_engine_cp_running consistent Signed-off-by: xren <[email protected]> * cp function renaming Signed-off-by: xren <[email protected]> * make loss logging broadcast aware of cp Signed-off-by: xren <[email protected]> * fix a typo Signed-off-by: Xiaowei Ren <[email protected]> * var name fix Signed-off-by: Xiaowei Ren <[email protected]> * import transformer layer specs from MCore Signed-off-by: Xiaowei Ren <[email protected]> * upgrade MCore version Signed-off-by: Xiaowei Ren <[email protected]> * add add context_parallel into the kwargs of dist opt Signed-off-by: Xiaowei Ren <[email protected]> * remove redundant cp check Signed-off-by: Xiaowei Ren <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * code style fix Signed-off-by: Xiaowei Ren <[email protected]> * recover docker file Signed-off-by: Xiaowei Ren <[email protected]> * fix seq_length of CP Signed-off-by: Xiaowei Ren <[email protected]> * recover seq-length which has been fixed in mcore Signed-off-by: Xiaowei Ren <[email protected]> * function name fix Signed-off-by: Xiaowei Ren <[email protected]> --------- Signed-off-by: xren <[email protected]> Signed-off-by: Xiaowei Ren <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

xrennvidia added 30 commits June 5, 2023 18:40

make nemo recognize sequence_parallel_size

afce64e

Signed-off-by: xren <[email protected]>

merge with main

3f98473

Signed-off-by: xren <[email protected]>

add helper functions to set up SP running in TE

e313000

Signed-off-by: xren <[email protected]>

Merge branch 'main' into xren/extend_sp

52ff102

slice seq length for a specific rank

5580955

Signed-off-by: Xiaowei Ren <[email protected]>

Merge branch 'main' into xren/extend_sp

887c615

fix data_parallel_size calculation

ebd6323

Signed-off-by: Xiaowei Ren <[email protected]>

minor change

58cca3d

Signed-off-by: Xiaowei Ren <[email protected]>

add missing argument of self

87f027a

Signed-off-by: Xiaowei Ren <[email protected]>

pass sp_global_ranks to TE transformer layer

9ebfcf7

Signed-off-by: Xiaowei Ren <[email protected]>

fix nsys setting

728fd43

Signed-off-by: Xiaowei Ren <[email protected]>

fix seq_len calculation

66615e8

Signed-off-by: xren <[email protected]>

fix attn_mask split across seq-length dim

e1f5eb7

Signed-off-by: xren <[email protected]>

code update of input split

cf0c75c

Signed-off-by: xren <[email protected]>

fix loss calculation

b57e218

Signed-off-by: xren <[email protected]>

fix loss_mask_sum calculation

69f4ae8

Signed-off-by: xren <[email protected]>

fix losss calculation

a38dd9a

Signed-off-by: xren <[email protected]>

merge with main

b31e31f

Signed-off-by: xren <[email protected]>

rename sequence_parallelism to context_parallelism

8ac42f1

Signed-off-by: xren <[email protected]>

minor change

f7c9b5b

Signed-off-by: xren <[email protected]>

fix loss_mask_sum calculation

49b1052

Signed-off-by: xren <[email protected]>

merge with main

ae889fc

Signed-off-by: xren <[email protected]>

make sure do not call megatron-core parallel_state while cp_size is 1

2c43687

Signed-off-by: xren <[email protected]>

Merge branch 'main' into xren/context_parallelism

25bf369

slice position embedding for different CP rank

61af551

Signed-off-by: xren <[email protected]>

fix mising property decorator

dc8a540

Signed-off-by: xren <[email protected]>

typo fix

46479c6

Signed-off-by: xren <[email protected]>

fix rpe_bias CP slicing

b64b563

Signed-off-by: xren <[email protected]>

Merge branch 'main' into xren/context_parallelism

0362de6

code style fix

e1654fb

Signed-off-by: xren <[email protected]>

xrennvidia added 2 commits December 15, 2023 17:17

merge with main

b56ce02

Signed-off-by: Xiaowei Ren <[email protected]>

merge with main

3a29733

Signed-off-by: Xiaowei Ren <[email protected]>

stu1130 reviewed Dec 19, 2023

View reviewed changes

xrennvidia added 5 commits December 18, 2023 17:28

function name fix

5d25e67

Signed-off-by: Xiaowei Ren <[email protected]>

Merge branch 'main' into xren/context_parallelism

ead55a0

merge with main

3a36003

Signed-off-by: Xiaowei Ren <[email protected]>

Merge branch 'main' into xren/context_parallelism

2d42b1c

merge with main

f66a5aa

Signed-off-by: Xiaowei Ren <[email protected]>

Merge branch 'main' into xren/context_parallelism

5d464c9

Merge branch 'main' into xren/context_parallelism

2c9c95e

ericharper approved these changes Jan 10, 2024

View reviewed changes

ericharper merged commit 58d6bce into main Jan 10, 2024
15 checks passed

ericharper deleted the xren/context_parallelism branch January 10, 2024 06:32

minitu pushed a commit to minitu/NeMo that referenced this pull request Jan 19, 2024

Revert "context parallelism (NVIDIA#7739)"

3867217

This reverts commit 58d6bce.

minitu pushed a commit to minitu/NeMo that referenced this pull request Jan 22, 2024

Revert "context parallelism (NVIDIA#7739)"

a16f6c6

This reverts commit 58d6bce.

minitu pushed a commit to minitu/NeMo that referenced this pull request Jan 22, 2024

Revert "Revert "context parallelism (NVIDIA#7739)""

adc7412

This reverts commit a16f6c6.

minitu pushed a commit to minitu/NeMo that referenced this pull request Jan 24, 2024

Revert "context parallelism (NVIDIA#7739)"

11df442

This reverts commit 58d6bce.

minitu pushed a commit to minitu/NeMo that referenced this pull request Jan 29, 2024

Revert "context parallelism (NVIDIA#7739)"

a2d409a

This reverts commit 58d6bce.

layalir pushed a commit to layalir/NeMo that referenced this pull request Jan 31, 2024

Revert "context parallelism (NVIDIA#7739)"

cb97123

This reverts commit 58d6bce.

minitu pushed a commit to minitu/NeMo that referenced this pull request Feb 1, 2024

Reinstate "context parallelism (NVIDIA#7739)"

f86f402

This reverts commit a16f6c6.

jbaczek pushed a commit to jbaczek/NeMo that referenced this pull request Feb 2, 2024

Revert "context parallelism (NVIDIA#7739)"

8f74c8c

This reverts commit 58d6bce. Signed-off-by: Jan Baczek <[email protected]>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

context parallelism #7739

context parallelism #7739

xrennvidia commented Oct 17, 2023

stu1130 commented Dec 9, 2023 •

edited

Loading

stu1130 Dec 19, 2023 •

edited

Loading

xrennvidia Dec 19, 2023

stu1130 Dec 19, 2023

xrennvidia Dec 19, 2023

stu1130 Dec 19, 2023

xrennvidia commented Jan 5, 2024

xrennvidia commented Jan 5, 2024

stu1130 commented Jan 9, 2024

xrennvidia commented Jan 9, 2024 •

edited

Loading

xrennvidia commented Jan 9, 2024

ericharper left a comment

context parallelism #7739

context parallelism #7739

Conversation

xrennvidia commented Oct 17, 2023

What does this PR do ?

Changelog

Usage

Before your PR is "Ready for review"

Who can review?

Additional Information

stu1130 commented Dec 9, 2023 • edited Loading

stu1130 Dec 19, 2023 • edited Loading

Choose a reason for hiding this comment

xrennvidia Dec 19, 2023

Choose a reason for hiding this comment

stu1130 Dec 19, 2023

Choose a reason for hiding this comment

xrennvidia Dec 19, 2023

Choose a reason for hiding this comment

stu1130 Dec 19, 2023

Choose a reason for hiding this comment

xrennvidia commented Jan 5, 2024

xrennvidia commented Jan 5, 2024

stu1130 commented Jan 9, 2024

xrennvidia commented Jan 9, 2024 • edited Loading

xrennvidia commented Jan 9, 2024

ericharper left a comment

Choose a reason for hiding this comment

stu1130 commented Dec 9, 2023 •

edited

Loading

stu1130 Dec 19, 2023 •

edited

Loading

xrennvidia commented Jan 9, 2024 •

edited

Loading