[Dev] Add more tests for LayerwiseDistOpt with dist_ckpt by BoxiangW · Pull Request #2132 · NVIDIA/Megatron-LM

BoxiangW · 2025-11-04T21:02:52Z

What does this PR do ?

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share discuss a design-doc with the team.

Contribution process

flowchart LR
    A[Pre-checks] --> B[PR Tests]
    subgraph Code Review/Approval
        C1[Expert Review] --> C2[Final Review]
    end
    B --> C1
    C2 --> D[Merge]

Pre-checks

I want this PR in a versioned release and have added the appropriate Milestone (e.g., Core 0.8)
I have added relevant unit tests
I have added relevant functional tests
I have added proper typing to my code Typing guidelines
I have added relevant documentation
I have run the autoformatter.sh on my PR

Code review

The following process is enforced via the CODEOWNERS file for changes into megatron/core. For changes outside of megatron/core, it is up to the PR author whether or not to tag the Final Reviewer team.

For MRs into `main` branch

(Step 1): Add PR label `Expert Review`

(Step 2): Collect the expert reviewers reviews

Attach the Expert Review label when your PR is ready for review.
GitHub auto-assigns expert reviewers based on your changes. They will get notified and pick up your PR soon.

⚠️ Only proceed to the next step once all reviewers have approved, merge-conflict are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

(Step 3): Final Review

Add Final Review label
GitHub auto-assigns final reviewers based on your changes. They will get notified and pick up your PR soon.

(Optional Step 4): Cherry-pick into release branch

If this PR also needs to be merged into core_r* release branches, after this PR has been merged, select Cherry-pick to open a new PR into the release branch.

For MRs into `dev` branch

The proposed review process for `dev` branch is under active discussion.

MRs are mergable after one approval by either eharper@nvidia.com or zijiey@nvidia.com.

Merging your PR

Any member of core-adlr and core-nemo will be able to merge your PR.

…factor` Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>

copy-pr-bot · 2025-11-04T21:02:57Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

BoxiangW · 2025-11-04T21:05:10Z

/ok to test 669a16c

BoxiangW · 2025-11-07T00:00:33Z

/ok to test 04154c8

Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>

BoxiangW · 2025-11-07T04:07:39Z

/ok to test 6e4c888

Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>

BoxiangW · 2025-11-07T04:43:12Z

/ok to test 267b06b

Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>

BoxiangW · 2025-11-07T05:39:12Z

/ok to test 7057451

Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>

BoxiangW · 2025-11-07T07:52:39Z

/ok to test aaf7235

dimapihtar · 2025-11-07T12:52:53Z

@BoxiangW looks like we need to add "functional_tests" label in order to trigger that functional tests you added in this PR.

dimapihtar

LGTM. Thank you!

BoxiangW · 2025-11-07T23:24:46Z

/ok to test 669a16c

copy-pr-bot · 2025-11-07T23:24:50Z

/ok to test 669a16c

@BoxiangW, there was an error processing your request: E2

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/

skyw

This PR should not be merged.

The title says it is only adding tests, but there a lot of code change including public API change. They should be spereate.
Test is not thorough, i.e. it doesn't catch a know issue that load fail on empty state_dict.
Plus some other improvement suggested.

skyw · 2025-11-10T16:41:23Z


-    muon_extra_scale_factor: float = 1.0
-    """Additional scale factor for the muon update."""
+    muon_extra_scale_factor: float = 0.2


FIX: Revert this change. Don't change this default scale factor.
ESPECIALLY not in a unrelated PR.

Reverted, this was recommended by Mkail and we can change it in a separate PR then

skyw · 2025-11-10T16:42:07Z

 def inspect_types(x: Any, prefix: Tuple = (), indent: int = 4):
    """Helper to print types of (nested) dict values."""
-    print_indent = lambda: print(" " * indent * len(prefix), end="")
+    print_indent = lambda: print(" " * indent * len(prefix), end="")  # pylint: disable=bad-builtin


too many inline disable. just disable it for the section.

# pylint: disable=bad-builtin ... # pylint: enable=bad-builtin

I'd suggest we just remove changes from this file(the only change is a comment). and thus we don't need to fix lint for it.
For that comment, I think it is ok for it to catch our empty list issue and we should fix it from layerwise side

skyw · 2025-11-10T16:42:47Z

                       choices=['blockwise', 'duplicated', 'distributed'],
                       help='How to perform NS calculation for tensor model parallel weights')
-    group.add_argument('--muon-extra-scale-factor', type=float, default=1.0,
+    group.add_argument('--muon-extra-scale-factor', type=float, default=0.2,


FIX: Revert this change.

skyw · 2025-11-10T16:44:05Z


        if succeeded:
-            logger.info(f"Job succeeded with status: {job_dict["status"]}")
+            logger.info(f"Job succeeded with status: {job_dict['status']}")


revert meaningless format change, unless their is a new style guide?

Yes, it was a pylint guide

but lines above and below all use double quote?

skyw · 2025-11-10T16:49:38Z

+            )
+            assert total_params > 0, "No parameters found in optimizer"
+
+    @pytest.mark.parametrize('tp_pp', [(1, 2), (2, 1), (2, 2)])


Fix: Bad variable grouping. should use TP and PP separately.
Same for all other tests.

This is not applicable for all the cases, since we have 8GPU max, if we do

@pytest.mark.parametrize('tp', [2, 4]) @pytest.mark.parametrize('pp', [2, 4])

It will reach 16 GPUs in total

you can skip test (with a message) if total number of GPU exceeded.

skyw · 2025-11-10T16:57:07Z

+                try:
+                    state_dict = load(load_sharded_sd, ckpt_dir)
+                    optimizer_B.load_state_dict(state_dict)
+                except Exception:


Don't silently pass tests. Can skip the test if there is a known reason.

Removed this special treatment

skyw · 2025-11-10T17:01:46Z

+                optim_param_state_B = optimizer_B.state_dict()
+
+                # Test both param state dicts are equal
+                diffs = diff(optim_param_state_A, optim_param_state_B)


Try not to use self made diff function as testing metric, it mixes things together unnecessarily. e.g. a bug in diff could have very large blast radius because all tests depend on it become faulty.
Better to explicitly compare everything here.

Added a similar but simpler explicitly checking in this file as a util function

FDecaYed · 2025-11-11T11:24:30Z

 def inspect_types(x: Any, prefix: Tuple = (), indent: int = 4):
    """Helper to print types of (nested) dict values."""
-    print_indent = lambda: print(" " * indent * len(prefix), end="")
+    print_indent = lambda: print(" " * indent * len(prefix), end="")  # pylint: disable=bad-builtin


I'd suggest we just remove changes from this file(the only change is a comment). and thus we don't need to fix lint for it.
For that comment, I think it is ok for it to catch our empty list issue and we should fix it from layerwise side

Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>

BoxiangW · 2025-11-12T21:30:28Z

/ok to test a6693fc

Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>

BoxiangW · 2025-11-12T23:05:37Z

/ok to test 73e424b

skyw

I mostly reviewed unit test part and suggested more improvement. LGTM overall.

I didn't review "functional" test as I don't exactly know what to expect there.

skyw · 2025-11-12T23:17:34Z

+            )
+            assert total_params > 0, "No parameters found in optimizer"
+
+    @pytest.mark.parametrize('tp_pp', [(1, 2), (2, 1), (2, 2)])


you can skip test (with a message) if total number of GPU exceeded.

skyw · 2025-11-12T23:20:32Z

+        pp = dp_size
+        tp = world_size // pp
+
+        if tp == 0:


Should check broadly whether world_size is divisible by pp.

Fixed this part

BoxiangW · 2025-11-12T23:21:58Z

/ok to test 896b2d7

Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>

BoxiangW · 2025-11-13T00:48:25Z

/ok to test b478987

FDecaYed

LG after changes

Victarry · 2025-11-13T09:31:31Z

Hi @BoxiangW , the golden values of this MR is not added, which cause the functional tests fail https://gitlab-master.nvidia.com/ADLR/megatron-lm/-/pipelines/38398052.

Add functional test yaml and change degfaults on `--muon-extra-scale-…

669a16c

…factor` Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>

BoxiangW self-assigned this Nov 4, 2025

BoxiangW requested review from a team as code owners November 4, 2025 21:02

BoxiangW added this to the Core 0.15 milestone Nov 4, 2025

BoxiangW added Run tests Expert Review [deprecated] Apply this label to indicate that your PR is ready for expert review. labels Nov 4, 2025

copy-pr-bot Bot temporarily deployed to nemo-ci November 4, 2025 21:05 Inactive

copy-pr-bot Bot temporarily deployed to test November 4, 2025 21:06 Inactive

copy-pr-bot Bot temporarily deployed to nemo-ci November 4, 2025 21:10 Inactive

copy-pr-bot Bot temporarily deployed to nemo-ci November 4, 2025 21:12 Inactive

copy-pr-bot Bot temporarily deployed to nemo-ci November 4, 2025 21:29 Inactive

Deal with Layerwise special case

6e4c888

Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>

Fix lint

267b06b

Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>

Fix functional test

7057451

Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>

BoxiangW added 3 commits November 6, 2025 22:07

Fix

35a18e1

Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>

Merge branch 'dev' into boxiangw/dist_ckpt_tests

ce77509

Fix config issue

aaf7235

Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>

dimapihtar approved these changes Nov 7, 2025

View reviewed changes

sbak5 reviewed Nov 7, 2025

View reviewed changes

Comment thread tests/unit_tests/dist_checkpointing/test_layer_wise_optimizer.py

skyw requested changes Nov 10, 2025

View reviewed changes

FDecaYed requested changes Nov 11, 2025

View reviewed changes

Address comments

a6693fc

Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>

Address comments

73e424b

Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>

Merge branch 'dev' into boxiangw/dist_ckpt_tests

896b2d7

skyw approved these changes Nov 12, 2025

View reviewed changes

BoxiangW added 2 commits November 12, 2025 16:48

Fix tests and split tp_pp

5882faf

Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>

Lint

b478987

Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>

FDecaYed approved these changes Nov 13, 2025

View reviewed changes

Uh oh!

Conversation

BoxiangW commented Nov 4, 2025 • edited by FDecaYed Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do ?

Contribution process

Pre-checks

Code review

(Step 1): Add PR label Expert Review

(Step 2): Collect the expert reviewers reviews

(Step 3): Final Review

(Optional Step 4): Cherry-pick into release branch

Merging your PR

Uh oh!

copy-pr-bot Bot commented Nov 4, 2025

Uh oh!

BoxiangW commented Nov 4, 2025

Uh oh!

BoxiangW commented Nov 7, 2025

Uh oh!

BoxiangW commented Nov 7, 2025

Uh oh!

BoxiangW commented Nov 7, 2025

Uh oh!

BoxiangW commented Nov 7, 2025

Uh oh!

BoxiangW commented Nov 7, 2025

Uh oh!

dimapihtar commented Nov 7, 2025

Uh oh!

dimapihtar left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

BoxiangW commented Nov 7, 2025

Uh oh!

copy-pr-bot Bot commented Nov 7, 2025

Uh oh!

skyw left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

BoxiangW commented Nov 12, 2025

Uh oh!

BoxiangW commented Nov 12, 2025

Uh oh!

skyw left a comment

Choose a reason for hiding this comment

BoxiangW commented Nov 4, 2025 •

edited by FDecaYed

Loading

(Step 1): Add PR label `Expert Review`