Skip to content

Extract args init to launch scripts - #4225

Merged
maanug-nv merged 10 commits into
NVIDIA:mainfrom
maanug-nv:refactor-initialize
Apr 15, 2026
Merged

Extract args init to launch scripts#4225
maanug-nv merged 10 commits into
NVIDIA:mainfrom
maanug-nv:refactor-initialize

Conversation

@maanug-nv

@maanug-nv maanug-nv commented Apr 9, 2026

Copy link
Copy Markdown
Contributor

What does this PR do ?

Extract args parsing and validation into a dedicated function out of initialize_megatron(), and call it in pretrain scripts. This is to facilitate separating args logic from training code.

We will build the config classes from the args in the pretrain scripts and pass that to pretrain(), in #4227 .

⚠️ 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"

  1. When your PR is ready, click Ready for Review.
  2. 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.

@copy-pr-bot

copy-pr-bot Bot commented Apr 9, 2026

Copy link
Copy Markdown

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

@maanug-nv
maanug-nv force-pushed the refactor-initialize branch 2 times, most recently from 676fbcb to b94beb4 Compare April 9, 2026 20:08
@maanug-nv
maanug-nv marked this pull request as ready for review April 9, 2026 20:09
@maanug-nv
maanug-nv requested review from a team as code owners April 9, 2026 20:09
@maanug-nv

Copy link
Copy Markdown
Contributor Author

/claude review

@svcnvidia-nemo-ci
svcnvidia-nemo-ci requested a review from a team April 9, 2026 20:09
@svcnvidia-nemo-ci svcnvidia-nemo-ci added this to the Core 0.16 milestone Apr 9, 2026
Comment thread megatron/training/arguments.py

@claude claude Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Two issues found:

  1. Bug: parse_and_validate_args missing return args — commented inline on megatron/training/arguments.py.

  2. Dead parameter: parsed_args in initialize_megatron — the parsed_args=None parameter on initialize_megatron() is now unused. The function body no longer references it (it just calls get_args()). It should be removed to avoid confusion, along with any callers that pass it.

@maanug-nv

Copy link
Copy Markdown
Contributor Author

Code review

Found 2 issues:

  1. pretrain() no longer accepts extra_args_provider or args_defaults, but pretrain_bert.py, pretrain_t5.py, pretrain_vlm.py, and train_rl.py still pass these — they will raise TypeError at runtime.

def pretrain(
train_valid_test_dataset_provider,
model_provider,
model_type,
forward_step_func,
process_non_loss_data_func=None,
get_embedding_ranks=None,
get_position_embedding_ranks=None,
non_loss_data_func=None,
store=None,
inprocess_call_wrapper: Optional[CallWrapper] = None,
):
"""Main training program.

https://github.com/NVIDIA/Megatron-LM/blob/14fedaebc2dfc58d87021427ba52b82fdd24d4aa/pretrain_bert.py#L193-L197

  1. In pretrain_mamba.py, from megatron.training.arguments import parse_and_validate_args was inserted before _PROGRAM_START_TIME = time.time(), violating the comment's stated intent to capture program start time before heavy imports. In pretrain_gpt.py the same import is correctly placed after the timestamp.

import time
from megatron.training.arguments import parse_and_validate_args
_PROGRAM_START_TIME = time.time()

🤖 Generated with Claude Code

- If this code review was useful, please react with 👍. Otherwise, react with 👎.

@maanug-nv
maanug-nv requested a review from a team as a code owner April 9, 2026 20:42
Comment thread megatron/training/arguments.py Outdated
maanug-nv and others added 8 commits April 10, 2026 10:42
Signed-off-by: Maanu Grover <maanug@nvidia.com>
Signed-off-by: Maanu Grover <maanug@nvidia.com>
Signed-off-by: Maanu Grover <maanug@nvidia.com>
Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com>
Signed-off-by: Maanu Grover <maanug@nvidia.com>
Signed-off-by: Maanu Grover <maanug@nvidia.com>
Signed-off-by: Maanu Grover <maanug@nvidia.com>
Signed-off-by: Maanu Grover <maanug@nvidia.com>
@svcnvidia-nemo-ci svcnvidia-nemo-ci added the Approved All necessary approvals have been made label Apr 14, 2026
@maanug-nv
maanug-nv enabled auto-merge April 14, 2026 20:50
@maanug-nv
maanug-nv added this pull request to the merge queue Apr 14, 2026
@svcnvidia-nemo-ci

Copy link
Copy Markdown
Contributor

🔄 Merge queue validation started!

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

@github-merge-queue
github-merge-queue Bot removed this pull request from the merge queue due to failed status checks Apr 14, 2026
Signed-off-by: Maanu Grover <maanug@nvidia.com>
Signed-off-by: Maanu Grover <maanug@nvidia.com>
@maanug-nv
maanug-nv enabled auto-merge April 14, 2026 22:12
@maanug-nv
maanug-nv added this pull request to the merge queue Apr 14, 2026
@svcnvidia-nemo-ci

Copy link
Copy Markdown
Contributor

🔄 Merge queue validation started!

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

Merged via the queue into NVIDIA:main with commit 97aca2f Apr 15, 2026
128 of 129 checks passed
@JordiBayarri

Copy link
Copy Markdown

Missing a update in "examples/multimodal/train.py"

lmcafee-nvidia pushed a commit to lmcafee-nvidia/Megatron-LM that referenced this pull request Apr 17, 2026
…init split

PR NVIDIA#4225 extracted argument parsing out of initialize_megatron(); call
parse_and_validate_args() separately and invoke initialize_megatron() with
no arguments.
Phlip79 added a commit to Phlip79/Megatron-LM that referenced this pull request Apr 19, 2026
Backward-compat stubs and aliases remain for external libraries, but
all internal Megatron-LM code now uses the canonical Hybrid names.

Changes:
- mamba_builders.py: backward-compat re-export stub with deprecation
  warning (delegates to hybrid_builders)
- pretrain_mamba.py: backward-compat wrapper that runpy's pretrain_hybrid
  with deprecation warning
- pretrain_hybrid.py: port upstream parse_and_validate_args refactor
  (from upstream NVIDIA#4225 which modified pretrain_mamba.py)
- megatron/training/arguments.py: import Symbols from canonical
  hybrid_layer_allocation path instead of backward-compat stub
- nemotron3_super_release_g200/model_config.yaml: update --spec and
  --mtp-spec to use hybrid_layer_specs path
- test_dsa_gpt_mamba_equivalence.py: update all class references and
  imports to HybridModel / hybrid_stack_spec
- tools/checkpoint/remap_gpt_dsa_to_mamba.py: update docstrings and
  help text from MambaModel/MambaStack to HybridModel/HybridStack

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
lmcafee-nvidia added a commit to lmcafee-nvidia/Megatron-LM that referenced this pull request Apr 30, 2026
…init split

PR NVIDIA#4225 extracted argument parsing out of initialize_megatron(); call
parse_and_validate_args() separately and invoke initialize_megatron() with
no arguments.
yangbofun pushed a commit to xlm-research/Megatron-LM that referenced this pull request May 22, 2026
Signed-off-by: Maanu Grover <maanug@nvidia.com>
Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com>
yhgalaxy pushed a commit to yhgalaxy/Megatron-LM that referenced this pull request Jun 17, 2026
Signed-off-by: Maanu Grover <maanug@nvidia.com>
Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com>
Signed-off-by: yhgalaxy <yhgalaxy@outlook.com>
jon-barker pushed a commit to jon-barker/Megatron-LM that referenced this pull request Jul 10, 2026
Signed-off-by: Maanu Grover <maanug@nvidia.com>
Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com>
Signed-off-by: Jon Barker <jbarker@aws-cmh-slurm-1-vscode-02.cm.cluster>
terminator123 pushed a commit to 021ai/Megatron-LM that referenced this pull request Aug 3, 2026
Signed-off-by: Maanu Grover <maanug@nvidia.com>
Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

6 participants