Mcore dist opt ckpt fix #9156

akoumpa · 2024-05-09T21:00:05Z

What does this PR do ?

Enable saving/restoring checkpoint when using mcore dist opt in NeMo.

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

GitHub Actions CI

The Jenkins CI system has been replaced by GitHub Actions self-hosted runners.

The GitHub Actions CI will run automatically when the "Run CICD" label is added to the PR.
To re-run CI remove and add the label again.
To run CI on an untrusted fork, a NeMo user with write access must first click "Approve and run".

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: Alexandros Koumparoulis <akoumparouli@nvidia.com>

Signed-off-by: akoumpa <akoumpa@users.noreply.github.com>

Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>

Signed-off-by: akoumpa <akoumpa@users.noreply.github.com>

examples/nlp/language_modeling/conf/megatron_gpt_config.yaml

nemo/core/optim/mcore_optim.py

… ckpt Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>

Signed-off-by: akoumpa <akoumpa@users.noreply.github.com>

examples/nlp/language_modeling/conf/megatron_gpt_config.yaml

nemo/core/optim/mcore_optim.py

Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>

Signed-off-by: akoumpa <akoumpa@users.noreply.github.com>

dimapihtar

LGTM. Thank you!

Signed-off-by: Alexandros Koumparoulis <153118171+akoumpa@users.noreply.github.com>

nemo/core/optim/mcore_optim.py

Co-authored-by: mikolajblaz <mikolajblaz@users.noreply.github.com> Signed-off-by: Alexandros Koumparoulis <153118171+akoumpa@users.noreply.github.com>

* Mcore dist opt ckpt fix Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com> * pass dp_zero_gather_scatter to starded-state-dict Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com> * Apply isort and black reformatting Signed-off-by: akoumpa <akoumpa@users.noreply.github.com> * introduce dist_ckpt_parallel_save option Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com> * determine sharding type from dist_ckpt_parallel_save Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com> * Apply isort and black reformatting Signed-off-by: akoumpa <akoumpa@users.noreply.github.com> * read model.disk_ckpt_parallel_save from cfg and pass it to mcore dist ckpt Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com> * Apply isort and black reformatting Signed-off-by: akoumpa <akoumpa@users.noreply.github.com> * Pass is_loading to mcore_optim.py's sharded_state_dict Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com> * Apply isort and black reformatting Signed-off-by: akoumpa <akoumpa@users.noreply.github.com> * Update nemo/core/optim/mcore_optim.py Co-authored-by: mikolajblaz <mikolajblaz@users.noreply.github.com> Signed-off-by: Alexandros Koumparoulis <153118171+akoumpa@users.noreply.github.com> --------- Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com> Signed-off-by: akoumpa <akoumpa@users.noreply.github.com> Signed-off-by: Alexandros Koumparoulis <153118171+akoumpa@users.noreply.github.com> Co-authored-by: akoumpa <akoumpa@users.noreply.github.com> Co-authored-by: mikolajblaz <mikolajblaz@users.noreply.github.com> Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>

* Mcore dist opt ckpt fix Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com> * pass dp_zero_gather_scatter to starded-state-dict Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com> * Apply isort and black reformatting Signed-off-by: akoumpa <akoumpa@users.noreply.github.com> * introduce dist_ckpt_parallel_save option Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com> * determine sharding type from dist_ckpt_parallel_save Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com> * Apply isort and black reformatting Signed-off-by: akoumpa <akoumpa@users.noreply.github.com> * read model.disk_ckpt_parallel_save from cfg and pass it to mcore dist ckpt Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com> * Apply isort and black reformatting Signed-off-by: akoumpa <akoumpa@users.noreply.github.com> * Pass is_loading to mcore_optim.py's sharded_state_dict Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com> * Apply isort and black reformatting Signed-off-by: akoumpa <akoumpa@users.noreply.github.com> * Update nemo/core/optim/mcore_optim.py Co-authored-by: mikolajblaz <mikolajblaz@users.noreply.github.com> Signed-off-by: Alexandros Koumparoulis <153118171+akoumpa@users.noreply.github.com> --------- Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com> Signed-off-by: akoumpa <akoumpa@users.noreply.github.com> Signed-off-by: Alexandros Koumparoulis <153118171+akoumpa@users.noreply.github.com> Co-authored-by: akoumpa <akoumpa@users.noreply.github.com> Co-authored-by: mikolajblaz <mikolajblaz@users.noreply.github.com> Signed-off-by: Jan Lasek <janek.lasek@gmail.com>

* Mcore dist opt ckpt fix Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com> * pass dp_zero_gather_scatter to starded-state-dict Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com> * Apply isort and black reformatting Signed-off-by: akoumpa <akoumpa@users.noreply.github.com> * introduce dist_ckpt_parallel_save option Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com> * determine sharding type from dist_ckpt_parallel_save Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com> * Apply isort and black reformatting Signed-off-by: akoumpa <akoumpa@users.noreply.github.com> * read model.disk_ckpt_parallel_save from cfg and pass it to mcore dist ckpt Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com> * Apply isort and black reformatting Signed-off-by: akoumpa <akoumpa@users.noreply.github.com> * Pass is_loading to mcore_optim.py's sharded_state_dict Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com> * Apply isort and black reformatting Signed-off-by: akoumpa <akoumpa@users.noreply.github.com> * Update nemo/core/optim/mcore_optim.py Co-authored-by: mikolajblaz <mikolajblaz@users.noreply.github.com> Signed-off-by: Alexandros Koumparoulis <153118171+akoumpa@users.noreply.github.com> --------- Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com> Signed-off-by: akoumpa <akoumpa@users.noreply.github.com> Signed-off-by: Alexandros Koumparoulis <153118171+akoumpa@users.noreply.github.com> Co-authored-by: akoumpa <akoumpa@users.noreply.github.com> Co-authored-by: mikolajblaz <mikolajblaz@users.noreply.github.com>

github-actions bot added the NLP label May 9, 2024

akoumpa force-pushed the akoumparouli/mcore_dist_opt_ckpt branch from 28b6376 to be23a09 Compare May 15, 2024 07:14

github-actions bot added the core Changes to NeMo Core label May 15, 2024

akoumpa force-pushed the akoumparouli/mcore_dist_opt_ckpt branch from be23a09 to b1df810 Compare May 15, 2024 16:24

akoumpa and others added 3 commits May 16, 2024 09:43

Mcore dist opt ckpt fix

9d13f15

Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>

pass dp_zero_gather_scatter to starded-state-dict

1a64904

Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>

Apply isort and black reformatting

cb78ce9

Signed-off-by: akoumpa <akoumpa@users.noreply.github.com>

akoumpa force-pushed the akoumparouli/mcore_dist_opt_ckpt branch from dcd692e to cb78ce9 Compare May 16, 2024 16:43

akoumpa added 2 commits May 16, 2024 09:45

introduce dist_ckpt_parallel_save option

006b6d8

Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>

determine sharding type from dist_ckpt_parallel_save

545fc51

Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>

akoumpa force-pushed the akoumparouli/mcore_dist_opt_ckpt branch from 00f478e to 545fc51 Compare May 16, 2024 16:59

Apply isort and black reformatting

4643470

Signed-off-by: akoumpa <akoumpa@users.noreply.github.com>

akoumpa force-pushed the akoumparouli/mcore_dist_opt_ckpt branch from 4643470 to a2d516c Compare May 17, 2024 09:11

akoumpa added the Run CICD label May 17, 2024

akoumpa marked this pull request as ready for review May 17, 2024 09:11

akoumpa requested a review from mikolajblaz May 17, 2024 09:11

akoumpa marked this pull request as draft May 17, 2024 09:12

mikolajblaz reviewed May 17, 2024

View reviewed changes

examples/nlp/language_modeling/conf/megatron_gpt_config.yaml Show resolved Hide resolved

nemo/core/optim/mcore_optim.py Outdated Show resolved Hide resolved

read model.disk_ckpt_parallel_save from cfg and pass it to mcore dist…

8fa988d

… ckpt Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>

akoumpa force-pushed the akoumparouli/mcore_dist_opt_ckpt branch from a2d516c to 8fa988d Compare May 17, 2024 22:15

Apply isort and black reformatting

4b297e0

Signed-off-by: akoumpa <akoumpa@users.noreply.github.com>

akoumpa marked this pull request as ready for review May 17, 2024 22:19

mikolajblaz reviewed May 21, 2024

View reviewed changes

examples/nlp/language_modeling/conf/megatron_gpt_config.yaml Show resolved Hide resolved

nemo/core/optim/mcore_optim.py Outdated Show resolved Hide resolved

akoumpa and others added 2 commits May 21, 2024 02:23

Pass is_loading to mcore_optim.py's sharded_state_dict

82b07c9

Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>

Apply isort and black reformatting

27eb553

Signed-off-by: akoumpa <akoumpa@users.noreply.github.com>

dimapihtar previously approved these changes May 21, 2024

View reviewed changes

Merge branch 'main' into akoumparouli/mcore_dist_opt_ckpt

d7cf7f0

Signed-off-by: Alexandros Koumparoulis <153118171+akoumpa@users.noreply.github.com>

akoumpa dismissed dimapihtar’s stale review via d7cf7f0 May 21, 2024 20:43

akoumpa requested a review from dimapihtar May 21, 2024 22:15

mikolajblaz reviewed May 22, 2024

View reviewed changes

nemo/core/optim/mcore_optim.py Outdated Show resolved Hide resolved

dimapihtar previously approved these changes May 22, 2024

View reviewed changes

Update nemo/core/optim/mcore_optim.py

0a9cd71

Co-authored-by: mikolajblaz <mikolajblaz@users.noreply.github.com> Signed-off-by: Alexandros Koumparoulis <153118171+akoumpa@users.noreply.github.com>

akoumpa dismissed dimapihtar’s stale review via 0a9cd71 May 22, 2024 17:00

akoumpa added Run CICD and removed Run CICD labels May 22, 2024

mikolajblaz approved these changes May 22, 2024

View reviewed changes

akoumpa merged commit 286d387 into main May 22, 2024
17 of 133 checks passed

akoumpa deleted the akoumparouli/mcore_dist_opt_ckpt branch May 22, 2024 18:45

ko3n1g mentioned this pull request Jul 18, 2024

Release 2.0.0rc1 #9786

Closed

2 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Mcore dist opt ckpt fix #9156

Mcore dist opt ckpt fix #9156

akoumpa commented May 9, 2024 •

edited

Loading

dimapihtar left a comment

Mcore dist opt ckpt fix #9156

Mcore dist opt ckpt fix #9156

Conversation

akoumpa commented May 9, 2024 • edited Loading

What does this PR do ?

Changelog

Usage

GitHub Actions CI

Before your PR is "Ready for review"

Who can review?

Additional Information

dimapihtar left a comment

Choose a reason for hiding this comment

akoumpa commented May 9, 2024 •

edited

Loading