Enable MCore v0.14 checkpoint formats #14428

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Open

mikolajblaz wants to merge 57 commits into NVIDIA-NeMo:main from mikolajblaz:mblaz/complex-ckpt-deprecation

Collaborator

mikolajblaz commented Aug 7, 2025

Important

The Update branch button must only be pressed in very rare occassions.
An outdated branch is never blocking the merge of a PR.
Please reach out to the automation team before pressing that button.

What does this PR do ?

Enable checkpoint formats introduced in MCore v0.14. NeMo counterpart of the "enablement MR" !3743

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)

mikolajblaz and others added 30 commits

June 24, 2025 16:35


          Sketch dist-ckpt content versioning

b11cae8

Signed-off-by: Mikołaj Błaż <[email protected]>

# Conflicts:
#	nemo/lightning/pytorch/strategies/megatron_strategy.py


          Apply isort and black reformatting

2eb4469

Signed-off-by: mikolajblaz <[email protected]>
Signed-off-by: Mikołaj Błaż <[email protected]>


          Add docs

1f9392d

Signed-off-by: Mikołaj Błaż <[email protected]>


          Change dist_opt_sharding_type name

6c72e8a

Signed-off-by: Mikołaj Błaż <[email protected]>


          Remove MappingProxyType

c3fd4d9

Signed-off-by: Mikołaj Błaż <[email protected]>


          Apply isort and black reformatting

10bf743

Signed-off-by: mikolajblaz <[email protected]>
Signed-off-by: Mikołaj Błaż <[email protected]>


          Expand docs

b65d93e

Signed-off-by: Mikołaj Błaż <[email protected]>


          Fix .rst formatting

ed938d2

Signed-off-by: Mikołaj Błaż <[email protected]>


          Fix .rst formatting

b910a73

Signed-off-by: Mikołaj Błaż <[email protected]>


          Fix .rst formatting

62c95d8

Signed-off-by: Mikołaj Błaż <[email protected]>


          Fix .rst formatting

9e8b287

Signed-off-by: Mikołaj Błaż <[email protected]>


          Unindent code

f07e4b8

Signed-off-by: Mikołaj Błaż <[email protected]>


          Match MLM metadata creation logic

8e31c80

Signed-off-by: Mikołaj Błaż <[email protected]>


          Apply isort and black reformatting

05ceba1

Signed-off-by: mikolajblaz <[email protected]>
Signed-off-by: Mikołaj Błaż <[email protected]>


          Remove one Nemo1 TODO

b6d47d9

Signed-off-by: Mikołaj Błaż <[email protected]>


          Add doc

4a4fb84

Signed-off-by: Mikołaj Błaż <[email protected]>


          Remove collections/nlp TODOs (NeMo 1)

1cc6744

Signed-off-by: Mikołaj Błaż <[email protected]>


          Fix some TODOs

04408b2

Signed-off-by: Mikołaj Błaż <[email protected]>


          Merge branch 'main' into mblaz/dist-ckpt-content-versioning

02ab774


          Remove f string

92ec44d

Signed-off-by: Mikołaj Błaż <[email protected]>


          Fix MegatronStrategy typo

224dacf

Signed-off-by: Mikołaj Błaż <[email protected]>


          Handle TODOs

3edd343

Signed-off-by: Mikołaj Błaż <[email protected]>


          Add missing import

19d633d

Signed-off-by: Mikołaj Błaż <[email protected]>


          Merge branch 'main' into mblaz/dist-ckpt-content-versioning

9200cc3


          Add content metadata flag to async

bf28b78

Signed-off-by: Mikołaj Błaż <[email protected]>


          Add missing load_content_metadata

4ca9973

Signed-off-by: Mikołaj Błaż <[email protected]>


          Load content_metadata through unwrapped_ckpt_io

b07be93

Signed-off-by: Mikołaj Błaż <[email protected]>


          Pass content_metadata through storage_options

98c93c7

Signed-off-by: Mikołaj Błaż <[email protected]>


          Fix linting problems

bd8bf05

Signed-off-by: Mikołaj Błaż <[email protected]>


          Add sharded_state_dict_metadata to FabricMegatronStrategy

6f2e842

Signed-off-by: Mikołaj Błaż <[email protected]>

dimapihtar temporarily deployed to test

August 19, 2025 11:17

— with

GitHub Actions Inactive


          Merge branch 'main' into mblaz/complex-ckpt-deprecation

3ce099c

ko3n1g added Run CICD and removed Run CICD labels

dimapihtar added Run CICD and removed Run CICD labels

dimapihtar temporarily deployed to test

August 26, 2025 16:37

— with

GitHub Actions Inactive


          add unit tests

08acb91

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

dimapihtar dismissed their stale review via

08acb91

September 2, 2025 14:44

dimapihtar added 2 commits

September 2, 2025 07:46


          remove print

13274cd

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


          fix style

5a56cce

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

dimapihtar removed the Run CICD label

dimapihtar self-requested a review

September 2, 2025 14:50

dimapihtar previously approved these changes

View reviewed changes

dimapihtar added the Run CICD label

dimapihtar temporarily deployed to test

September 2, 2025 14:51

— with

GitHub Actions Inactive

github-advanced-security bot found potential problems

View reviewed changes

tests/lightning/pytorch/strategies/test_megatron_strategy.py

@@ @@ -67,3 +91,24 @@ @@
                       checkpoint = {"other": 123}
                       strategy.load_model_state_dict(checkpoint)
                       strategy.optimizers[0].reload_model_params.assert_called_once_with(checkpoint)
+                  def test_sharded_state_dict_metadata(self):
+                      strategy = MegatronStrategy(ckpt_save_pre_mcore_014=False, ckpt_parallel_save_optim=True)

Check warning

Code scanning / CodeQL

Variable defined multiple times Warning test

This assignment to 'strategy' is unnecessary as it is

redefined

before this value is used.


          add more unit tests

0bc4a75

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

dimapihtar dismissed their stale review via

0bc4a75

September 4, 2025 13:10

chtruong814 added Run CICD and removed Run CICD labels


          Apply isort and black reformatting

163bbb2

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

chtruong814 added Run CICD and removed Run CICD labels

chtruong814 temporarily deployed to test

September 4, 2025 13:13

— with

GitHub Actions Inactive

github-advanced-security bot found potential problems

View reviewed changes

tests/lightning/pytorch/strategies/test_megatron_strategy.py

+                  def test_init_errors(self):
+                      with pytest.raises(ValueError):
+                          strategy = MegatronStrategy(ddp="pytorch", fsdp="fsdp")

Check warning

Code scanning / CodeQL

Variable defined multiple times Warning test

This assignment to 'strategy' is unnecessary as it is

redefined

before this value is used.

tests/lightning/pytorch/strategies/test_megatron_strategy.py

+                          strategy = MegatronStrategy(ddp="pytorch", fsdp="fsdp")
+                      with pytest.raises(ValueError):
+                          strategy = MegatronStrategy(ddp="test")

Check warning

Code scanning / CodeQL

Variable defined multiple times Warning test

This assignment to 'strategy' is unnecessary as it is

redefined

before this value is used.

tests/lightning/pytorch/strategies/test_megatron_strategy.py

+                          strategy = MegatronStrategy(ddp="test")
+                      with pytest.raises(NotImplementedError):
+                          strategy = MegatronStrategy(fsdp="pytorch")

Check warning

Code scanning / CodeQL

Variable defined multiple times Warning test

This assignment to 'strategy' is unnecessary as it is

redefined

before this value is used.

tests/lightning/pytorch/strategies/test_megatron_strategy.py


		ddp = DistributedDataParallelConfig(use_custom_fsdp=True)
		strategy = MegatronStrategy(ddp=ddp, fsdp=None)

Check warning

Code scanning / CodeQL

Variable defined multiple times Warning test

This assignment to 'strategy' is unnecessary as it is

redefined

before this value is used.

tests/lightning/pytorch/strategies/test_megatron_strategy.py


		ddp = DistributedDataParallelConfig(use_custom_fsdp=False)
		strategy = MegatronStrategy(ddp=ddp, fsdp="megatron")

Check warning

Code scanning / CodeQL

Variable defined multiple times Warning test

This assignment to 'strategy' is unnecessary as it is

redefined

before this value is used.

tests/lightning/pytorch/strategies/test_megatron_strategy.py


		with pytest.raises(ValueError):
		strategy = MegatronStrategy(fsdp="test")

Check notice

Code scanning / CodeQL

Unused local variable Note test

Variable strategy is not used.


          Merge branch 'main' into mblaz/complex-ckpt-deprecation

47121b4

chtruong814 added Run CICD and removed Run CICD labels

chtruong814 temporarily deployed to test

September 4, 2025 22:28

— with

GitHub Actions Inactive

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels