Allows non-strict load with distributed checkpoints #9613

mikolajblaz · 2024-07-04T11:10:16Z

What does this PR do ?

With distributed checkpoints, the mismatches between the runtime model and checkpoint model manifest during dist_checkpoint.load (not during model.load_state_dict as with regular checkpoints).
This PR adds a flag that allows to adjust load strictness (e.g. ignore unexpected keys).

This PR relies on MCore feature that is not merged yet (merge ETA 5th July): https://gitlab-master.nvidia.com/ADLR/megatron-lm/-/merge_requests/1628

Collection: NLP

Changelog

Add model.dist_ckpt_load_strictness flag to control dist ckpt load strictness. The most useful value is log_all which warns about all mismatches but performs the checkpoint load for a matching state dict subset.

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: Mikołaj Błaż <mblaz@nvidia.com>

Dockerfile.ci

Signed-off-by: Mikołaj Błaż <mblaz@nvidia.com>

dimapihtar

LGTM. Thank you!

* Allow non-strict load Signed-off-by: Mikołaj Błaż <mblaz@nvidia.com> * Point to non-stric load MCore branch Signed-off-by: Mikołaj Błaż <mblaz@nvidia.com> * Avoid module level StrictHandling Signed-off-by: Mikołaj Błaż <mblaz@nvidia.com> * Use MCore fork Signed-off-by: Mikołaj Błaż <mblaz@nvidia.com> * Update to MCore fix Signed-off-by: Mikołaj Błaż <mblaz@nvidia.com> * Restore ackward compatibility Signed-off-by: Mikołaj Błaż <mblaz@nvidia.com> * Update flag defaults Signed-off-by: Mikołaj Błaż <mblaz@nvidia.com> * Update MCore tag Signed-off-by: Mikołaj Błaż <mblaz@nvidia.com> * Update PyT Dist interface Signed-off-by: Mikołaj Błaż <mblaz@nvidia.com> * Update to latest core_r0.8.0 Signed-off-by: Mikołaj Błaż <mblaz@nvidia.com> --------- Signed-off-by: Mikołaj Błaż <mblaz@nvidia.com>

* Allow non-strict load * Point to non-stric load MCore branch * Avoid module level StrictHandling * Use MCore fork * Update to MCore fix * Restore ackward compatibility * Update flag defaults * Update MCore tag * Update PyT Dist interface * Update to latest core_r0.8.0 --------- Signed-off-by: Mikołaj Błaż <mblaz@nvidia.com> Co-authored-by: mikolajblaz <mikolajblaz@users.noreply.github.com>

…IDIA#9715) * Allow non-strict load * Point to non-stric load MCore branch * Avoid module level StrictHandling * Use MCore fork * Update to MCore fix * Restore ackward compatibility * Update flag defaults * Update MCore tag * Update PyT Dist interface * Update to latest core_r0.8.0 --------- Signed-off-by: Mikołaj Błaż <mblaz@nvidia.com> Co-authored-by: mikolajblaz <mikolajblaz@users.noreply.github.com>

* Allow non-strict load * Point to non-stric load MCore branch * Avoid module level StrictHandling * Use MCore fork * Update to MCore fix * Restore ackward compatibility * Update flag defaults * Update MCore tag * Update PyT Dist interface * Update to latest core_r0.8.0 --------- Signed-off-by: Mikołaj Błaż <mblaz@nvidia.com> Co-authored-by: mikolajblaz <mikolajblaz@users.noreply.github.com> Signed-off-by: Tugrul Konuk <ertkonuk@gmail.com>

…IDIA#9715) * Allow non-strict load * Point to non-stric load MCore branch * Avoid module level StrictHandling * Use MCore fork * Update to MCore fix * Restore ackward compatibility * Update flag defaults * Update MCore tag * Update PyT Dist interface * Update to latest core_r0.8.0 --------- Signed-off-by: Mikołaj Błaż <mblaz@nvidia.com> Co-authored-by: mikolajblaz <mikolajblaz@users.noreply.github.com> Signed-off-by: Malay Nagda <malayn@malayn-mlt.client.nvidia.com>

…IDIA#9715) * Allow non-strict load * Point to non-stric load MCore branch * Avoid module level StrictHandling * Use MCore fork * Update to MCore fix * Restore ackward compatibility * Update flag defaults * Update MCore tag * Update PyT Dist interface * Update to latest core_r0.8.0 --------- Signed-off-by: Mikołaj Błaż <mblaz@nvidia.com> Co-authored-by: mikolajblaz <mikolajblaz@users.noreply.github.com> Signed-off-by: tonyjie <jl4257@cornell.edu>

* Allow non-strict load Signed-off-by: Mikołaj Błaż <mblaz@nvidia.com> * Point to non-stric load MCore branch Signed-off-by: Mikołaj Błaż <mblaz@nvidia.com> * Avoid module level StrictHandling Signed-off-by: Mikołaj Błaż <mblaz@nvidia.com> * Use MCore fork Signed-off-by: Mikołaj Błaż <mblaz@nvidia.com> * Update to MCore fix Signed-off-by: Mikołaj Błaż <mblaz@nvidia.com> * Restore ackward compatibility Signed-off-by: Mikołaj Błaż <mblaz@nvidia.com> * Update flag defaults Signed-off-by: Mikołaj Błaż <mblaz@nvidia.com> * Update MCore tag Signed-off-by: Mikołaj Błaż <mblaz@nvidia.com> * Update PyT Dist interface Signed-off-by: Mikołaj Błaż <mblaz@nvidia.com> * Update to latest core_r0.8.0 Signed-off-by: Mikołaj Błaż <mblaz@nvidia.com> --------- Signed-off-by: Mikołaj Błaż <mblaz@nvidia.com>

* Allow non-strict load * Point to non-stric load MCore branch * Avoid module level StrictHandling * Use MCore fork * Update to MCore fix * Restore ackward compatibility * Update flag defaults * Update MCore tag * Update PyT Dist interface * Update to latest core_r0.8.0 --------- Signed-off-by: Mikołaj Błaż <mblaz@nvidia.com> Co-authored-by: mikolajblaz <mikolajblaz@users.noreply.github.com>

…IDIA#9715) * Allow non-strict load * Point to non-stric load MCore branch * Avoid module level StrictHandling * Use MCore fork * Update to MCore fix * Restore ackward compatibility * Update flag defaults * Update MCore tag * Update PyT Dist interface * Update to latest core_r0.8.0 --------- Signed-off-by: Mikołaj Błaż <mblaz@nvidia.com> Co-authored-by: mikolajblaz <mikolajblaz@users.noreply.github.com> Signed-off-by: Hainan Xu <hainanx@nvidia.com>

mikolajblaz self-assigned this Jul 4, 2024

github-actions bot added the NLP label Jul 4, 2024

mikolajblaz added 3 commits July 4, 2024 13:13

Allow non-strict load

812ea87

Signed-off-by: Mikołaj Błaż <mblaz@nvidia.com>

Point to non-stric load MCore branch

be985fb

Signed-off-by: Mikołaj Błaż <mblaz@nvidia.com>

Avoid module level StrictHandling

cc8e4a1

Signed-off-by: Mikołaj Błaż <mblaz@nvidia.com>

mikolajblaz force-pushed the mblaz/non-strict-load branch from ddb6be3 to cc8e4a1 Compare July 4, 2024 11:13

mikolajblaz added the Run CICD label Jul 4, 2024

Use MCore fork

bf95042

Signed-off-by: Mikołaj Błaż <mblaz@nvidia.com>

mikolajblaz added Run CICD and removed Run CICD labels Jul 4, 2024

Update to MCore fix

a9a53f6

Signed-off-by: Mikołaj Błaż <mblaz@nvidia.com>

mikolajblaz added Run CICD and removed Run CICD labels Jul 4, 2024

mikolajblaz added 3 commits July 4, 2024 17:31

Restore ackward compatibility

a08d873

Signed-off-by: Mikołaj Błaż <mblaz@nvidia.com>

Update flag defaults

096b10a

Signed-off-by: Mikołaj Błaż <mblaz@nvidia.com>

Update MCore tag

e95d5cc

Signed-off-by: Mikołaj Błaż <mblaz@nvidia.com>

mikolajblaz added Run CICD and removed Run CICD labels Jul 4, 2024

mikolajblaz requested a review from dimapihtar July 4, 2024 15:36

mikolajblaz commented Jul 4, 2024

View reviewed changes

Dockerfile.ci Outdated Show resolved Hide resolved

dimapihtar reviewed Jul 4, 2024

View reviewed changes

Dockerfile.ci Outdated Show resolved Hide resolved

Update PyT Dist interface

ef61804

Signed-off-by: Mikołaj Błaż <mblaz@nvidia.com>

mikolajblaz added Run CICD and removed Run CICD labels Jul 4, 2024

mikolajblaz mentioned this pull request Jul 5, 2024

Wrap all fp8 extra states in LocalNonpersistentObject #9422

Closed

8 tasks

mikolajblaz changed the base branch from main to r2.0.0rc1 July 8, 2024 12:00

terrykong force-pushed the r2.0.0rc1 branch from 6697bbd to 66c960e Compare July 8, 2024 20:53

mikolajblaz and others added 3 commits July 10, 2024 11:10

Merge branch 'r2.0.0rc1' into mblaz/non-strict-load

96b7bd6

Update to latest core_r0.8.0

8406408

Signed-off-by: Mikołaj Błaż <mblaz@nvidia.com>

Merge branch 'r2.0.0rc1' into mblaz/non-strict-load

3e6df2c

mikolajblaz added Run CICD and removed Run CICD labels Jul 11, 2024

Merge branch 'r2.0.0rc1' into mblaz/non-strict-load

42139b2

mikolajblaz removed the Run CICD label Jul 11, 2024

Merge branch 'r2.0.0rc1' into mblaz/non-strict-load

7d01941

mikolajblaz added the Run CICD label Jul 11, 2024

Merge branch 'r2.0.0rc1' into mblaz/non-strict-load

a91c9c7

mikolajblaz added Run CICD and removed Run CICD labels Jul 11, 2024

Merge branch 'r2.0.0rc1' into mblaz/non-strict-load

5614d6b

mikolajblaz added Run CICD and removed Run CICD labels Jul 12, 2024

dimapihtar self-requested a review July 12, 2024 10:19

dimapihtar approved these changes Jul 12, 2024

View reviewed changes

mikolajblaz added Run CICD and removed Run CICD labels Jul 12, 2024

mikolajblaz merged commit 4239511 into r2.0.0rc1 Jul 12, 2024
217 of 219 checks passed

mikolajblaz deleted the mblaz/non-strict-load branch July 12, 2024 13:12

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

Allows non-strict load with distributed checkpoints #9613

Allows non-strict load with distributed checkpoints #9613

mikolajblaz commented Jul 4, 2024 •

edited

Loading

dimapihtar left a comment

Allows non-strict load with distributed checkpoints #9613

Allows non-strict load with distributed checkpoints #9613

Conversation

mikolajblaz commented Jul 4, 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

mikolajblaz commented Jul 4, 2024 •

edited

Loading