base strategy simplification by dimapihtar · Pull Request #4001 · NVIDIA/Megatron-LM

dimapihtar · 2026-03-23T14:41:34Z

What does this PR do ?

Base strategy simplification. Getting rid of base's asbtract classes. Part of checkpointing simplification strategy.

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact the @mcore-oncall.

Contribution process

Pre-checks

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

Feel free to message or comment the @mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

When your PR is ready, click Ready for Review.
An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
- Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

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.

Signed-off-by: dimapihtar <dpihtar@gmail.com>

copy-pr-bot · 2026-03-23T14:41:39Z

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.

Signed-off-by: dimapihtar <dpihtar@gmail.com>

dimapihtar · 2026-03-30T18:10:11Z

/ok to test 67e70bf

Signed-off-by: dimapihtar <dpihtar@gmail.com>

dimapihtar · 2026-03-30T18:33:05Z

/ok to test 80283cc

Signed-off-by: dimapihtar <dpihtar@gmail.com>

Signed-off-by: dimapihtar <dpykhtar@nvidia.com>

dimapihtar · 2026-04-02T18:09:34Z

/ok to test 1d394b0

Signed-off-by: dimapihtar <dpykhtar@nvidia.com>

dimapihtar · 2026-04-03T12:22:08Z

/ok to test 778a158

Signed-off-by: dimapihtar <dpykhtar@nvidia.com>

dimapihtar · 2026-04-03T12:24:53Z

/ok to test 1e4e81e

Signed-off-by: dimapihtar <dpykhtar@nvidia.com>

dimapihtar · 2026-04-03T13:12:57Z

/ok to test 4aea010

svcnvidia-nemo-ci · 2026-04-03T16:02:25Z

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/23952697316

Signed-off-by: dimapihtar <dpihtar@gmail.com> Signed-off-by: dimapihtar <dpykhtar@nvidia.com> Signed-off-by: oliver könig <okoenig@nvidia.com>

…test config Upstream removed these flags from hybrid_mr_mcore_te_tp2_pp1_cp1_dgx_a100_1N8G/model_config.yaml in commit 2697b82 ("base strategy simplification NVIDIA#4001"), but they were accidentally re-introduced by a merge conflict resolution. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Signed-off-by: dimapihtar <dpihtar@gmail.com> Signed-off-by: dimapihtar <dpykhtar@nvidia.com>

Signed-off-by: dimapihtar <dpihtar@gmail.com> Signed-off-by: dimapihtar <dpykhtar@nvidia.com> Signed-off-by: yhgalaxy <yhgalaxy@outlook.com>

Signed-off-by: dimapihtar <dpihtar@gmail.com> Signed-off-by: dimapihtar <dpykhtar@nvidia.com> Signed-off-by: Jon Barker <jbarker@aws-cmh-slurm-1-vscode-02.cm.cluster>

Signed-off-by: dimapihtar <dpihtar@gmail.com> Signed-off-by: dimapihtar <dpykhtar@nvidia.com>

drop most of base.py strategy

027e064

Signed-off-by: dimapihtar <dpihtar@gmail.com>

dimapihtar added 14 commits March 24, 2026 06:57

code clenaup

17168af

Signed-off-by: dimapihtar <dpihtar@gmail.com>

fix docstring

8621dde

Signed-off-by: dimapihtar <dpihtar@gmail.com>

remove base import

8687661

Signed-off-by: dimapihtar <dpihtar@gmail.com>

fix imports

bb1ed10

Signed-off-by: dimapihtar <dpihtar@gmail.com>

fix imports

358d076

Signed-off-by: dimapihtar <dpihtar@gmail.com>

fix async

cf3f930

Signed-off-by: dimapihtar <dpihtar@gmail.com>

fix async

557a1d9

Signed-off-by: dimapihtar <dpihtar@gmail.com>

Merge branch 'main' into simplify_base_strategy

39fd4a2

fix unit test

36c1711

Signed-off-by: dimapihtar <dpihtar@gmail.com>

remove super init call

2436abc

Signed-off-by: dimapihtar <dpihtar@gmail.com>

fix unit tests

25b15a5

Signed-off-by: dimapihtar <dpihtar@gmail.com>

fix serialization unit test

b0990ed

Signed-off-by: dimapihtar <dpihtar@gmail.com>

remove super init call

a3f8d6d

Signed-off-by: dimapihtar <dpihtar@gmail.com>

fix import

f1d3f87

Signed-off-by: dimapihtar <dpihtar@gmail.com>

dimapihtar marked this pull request as ready for review March 30, 2026 18:09

dimapihtar requested review from a team as code owners March 30, 2026 18:09

Merge branch 'main' into simplify_base_strategy

67e70bf

dimapihtar added the complexity: high label Mar 30, 2026

svcnvidia-nemo-ci requested a review from a team March 30, 2026 18:10

svcnvidia-nemo-ci added this to the Core 0.16 milestone Mar 30, 2026

svcnvidia-nemo-ci added complexity: medium and removed complexity: high labels Mar 30, 2026

fix code style

80283cc

Signed-off-by: dimapihtar <dpihtar@gmail.com>

remove func

2132326

Signed-off-by: dimapihtar <dpihtar@gmail.com>

jaredcasper approved these changes Apr 2, 2026

View reviewed changes

fix unit test

1d394b0

Signed-off-by: dimapihtar <dpykhtar@nvidia.com>

copy-pr-bot Bot temporarily deployed to test April 2, 2026 18:10 Inactive

dimapihtar requested a review from yaoyu-33 April 2, 2026 18:14

yaoyu-33 approved these changes Apr 2, 2026

View reviewed changes

svcnvidia-nemo-ci added Approved All necessary approvals have been made and removed Final Review PR is in the "final review" stage labels Apr 2, 2026

dimapihtar and others added 3 commits April 3, 2026 05:20

fix depreaction warnings

c03e751

Signed-off-by: dimapihtar <dpykhtar@nvidia.com>

fix deprecation warnings

d8f2bcb

Signed-off-by: dimapihtar <dpykhtar@nvidia.com>

Merge branch 'main' into simplify_base_strategy

778a158

fix code style

1e4e81e

Signed-off-by: dimapihtar <dpykhtar@nvidia.com>

copy-pr-bot Bot temporarily deployed to test April 3, 2026 12:25 Inactive

fix unit tests

4aea010

Signed-off-by: dimapihtar <dpykhtar@nvidia.com>

copy-pr-bot Bot temporarily deployed to test April 3, 2026 13:14 Inactive

dimapihtar added this pull request to the merge queue Apr 3, 2026

Merged via the queue into NVIDIA:main with commit 2697b82 Apr 3, 2026
129 checks passed

dimapihtar deleted the simplify_base_strategy branch April 3, 2026 16:39

ko3n1g pushed a commit that referenced this pull request Apr 4, 2026

base strategy simplification (#4001)

d656753

Signed-off-by: dimapihtar <dpihtar@gmail.com> Signed-off-by: dimapihtar <dpykhtar@nvidia.com> Signed-off-by: oliver könig <okoenig@nvidia.com>

ko3n1g pushed a commit to ko3n1g/Megatron-LM that referenced this pull request Apr 4, 2026

base strategy simplification (NVIDIA#4001)

22dfd24

Signed-off-by: dimapihtar <dpihtar@gmail.com> Signed-off-by: dimapihtar <dpykhtar@nvidia.com> Signed-off-by: oliver könig <okoenig@nvidia.com>

yangbofun pushed a commit to xlm-research/Megatron-LM that referenced this pull request May 22, 2026

base strategy simplification (NVIDIA#4001)

0a66ee8

Signed-off-by: dimapihtar <dpihtar@gmail.com> Signed-off-by: dimapihtar <dpykhtar@nvidia.com>

yhgalaxy pushed a commit to yhgalaxy/Megatron-LM that referenced this pull request Jun 17, 2026

base strategy simplification (NVIDIA#4001)

497f6c0

Signed-off-by: dimapihtar <dpihtar@gmail.com> Signed-off-by: dimapihtar <dpykhtar@nvidia.com> Signed-off-by: yhgalaxy <yhgalaxy@outlook.com>

terminator123 pushed a commit to 021ai/Megatron-LM that referenced this pull request Aug 3, 2026

base strategy simplification (NVIDIA#4001)

5078e16

Signed-off-by: dimapihtar <dpihtar@gmail.com> Signed-off-by: dimapihtar <dpykhtar@nvidia.com>

Navigation Menu

Uh oh!

base strategy simplification - #4001

base strategy simplification#4001
dimapihtar merged 36 commits into
NVIDIA:mainfrom
dimapihtar:simplify_base_strategy

dimapihtar commented Mar 23, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Mar 23, 2026

Uh oh!

dimapihtar commented Mar 30, 2026

Uh oh!

dimapihtar commented Mar 30, 2026

Uh oh!

dimapihtar commented Apr 2, 2026

Uh oh!

dimapihtar commented Apr 3, 2026

Uh oh!

dimapihtar commented Apr 3, 2026

Uh oh!

dimapihtar commented Apr 3, 2026

Uh oh!

svcnvidia-nemo-ci commented Apr 3, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

Uh oh!

Conversation

dimapihtar commented Mar 23, 2026 • edited 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: Mark PR as "Ready for Review"

Step 2: Final Review

Step 3: Approved

Merge

Uh oh!

copy-pr-bot Bot commented Mar 23, 2026

Uh oh!

dimapihtar commented Mar 30, 2026

Uh oh!

dimapihtar commented Mar 30, 2026

Uh oh!

dimapihtar commented Apr 2, 2026

Uh oh!

dimapihtar commented Apr 3, 2026

Uh oh!

dimapihtar commented Apr 3, 2026

Uh oh!

dimapihtar commented Apr 3, 2026

Uh oh!

svcnvidia-nemo-ci commented Apr 3, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

dimapihtar commented Mar 23, 2026 •

edited

Loading