Skip to content

refactor(training): migrate SchedulerConfig - #8211

Merged
dimapihtar merged 53 commits into
NVIDIA:mainfrom
dimapihtar:scheduler_config_migration
Oct 21, 2026
Merged

dimapihtar merged 53 commits into
NVIDIA:mainfrom
dimapihtar:scheduler_config_migration

Conversation

@dimapihtar

@dimapihtar dimapihtar commented Oct 14, 2026 •

Copy link
Copy Markdown
Contributor
  • I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Uses cfg.scheduler reference instead of args for SchedulerConfig params.

⚠️ 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 @NVIDIA/mcore-oncall.

Issue tracking

For PRs from open-source community contributors:

  • New features: a linked issue is required. Please open a feature request and reference it here before submitting the PR.
  • Small updates (bug fixes, minor improvements): a linked issue is recommended and will accelerate the PR review process.

Linked issue:

Contribution process

Pre-checks

  • I have added relevant unit tests
  • I have added relevant functional tests
  • If this PR adds or changes a GPU kernel (Triton, jit_fuser/torch.compile, CUDA extension, TE or external-library dispatch, or a scatter/index accumulation), I have added or updated its bit-exact determinism test and registered it in tests/unit_tests/determinism/kernels/manifest.py (guide)
  • 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 @NVIDIA/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.

@maanug-nv maanug-nv 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.

lgtm now

@maanug-nv

Copy link
Copy Markdown
Contributor

/ok to test 14cb9c8

Signed-off-by: Maanu Grover <maanug@nvidia.com>
@maanug-nv

Copy link
Copy Markdown
Contributor

/ok to test 1abfe6f

dimapihtar and others added 2 commits October 20, 2026 18:22
Signed-off-by: Dmytro Pykhtar <dpykhtar@nvidia.com>
@dimapihtar

Copy link
Copy Markdown
Contributor Author

/ok to test bf2b9d3

@nemo-automation-bot

Copy link
Copy Markdown

🔄 Merge queue validation started!

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

Merged via the queue into NVIDIA:main with commit e4eb89b Oct 21, 2026
154 checks passed
@dimapihtar
dimapihtar deleted the scheduler_config_migration branch October 21, 2026 14:57
asolergi-nv added a commit to asolergi-nv/Megatron-LM that referenced this pull request Oct 22, 2026
Forty upstream commits; five files conflicted, all against the data-section
work:

- megatron/training/config/container.py: NVIDIA#8149 / NVIDIA#8211 restructured
  PretrainConfigContainer.validate() into _validate_rl_config() and gave
  _finalize() the batch-size and scheduler resolution. Kept both: validate()
  ends with _validate_data_section(); _finalize() checks the dataset mirrors
  first, runs upstream's batch-size, RNG and scheduler steps, then writes the
  mirrors and finalizes the data section before the RL DDP finalize.
- megatron/training/global_vars.py: NVIDIA#8149 made set_global_variables()
  finalize every PretrainConfigContainer after registering args and the run
  config. Taken as is; the branch's narrower "finalize when a dataset section
  exists" is subsumed (the dataset hook still requires a finalized container).
- pretrain_gpt.py, pretrain_hybrid.py: upstream edited is_dataset_built_on_rank
  (NVIDIA#8378: the sequence-packing scheduler keeps data iterators on the first and
  last stages only) and core_gpt_dataset_config_from_args (NVIDIA#8149: validation
  flags read from cfg.validation) in place; the branch had moved both. The
  removal stands, the scheduler branch is ported into
  megatron.training.datasets.providers.is_dataset_built_on_rank, and the
  validation flags already come from cfg.validation through the mirrors.
- tools/prepare_cache.py: both import lines kept; upstream's container
  registration (pretrain_cfg_container_from_args, finalize, set_run_config,
  update_train_iters(cfg.train)) stays next to the branch's standalone section.

Tests: the ownership fixture registers the parsed arguments with a
batch-parallel size and a micro batch, which _finalize() now resolves from
the registered args; the bootstrap test expects loader-only containers to be
finalized too; the loader-only fixture test registers a minimal namespace as
upstream's finalize tests do.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Signed-off-by: asolergi-nv <asolergibert@nvidia.com>
nemo-autobot-origin Bot pushed a commit to svcnemo-autobot/Megatron-LM that referenced this pull request Oct 24, 2026
Repairs semantic interactions between the imported main slice (tip
f9d9bdc) and dev
dd8f93b found by the baseline sweep
and review:

- experimental_attention_variant/dsa_gqa.py: mark main NVIDIA#7156's
  DSGQACoreAttention as a DSA indexer-loss metric writer
  (logs_dsa_indexer_loss = True, as DSAttention and CSA do) so dev's
  fixed-capacity DSA metric tracker reserves capacity for it.
- moe/experts.py: assign TEGroupedMLP.activation_recompute (and the
  fp8/fp4 set_save_original_input(linear_fc2) hook) before the fused-impl
  support check reads it; read main's _grouped_linear_kwargs with a
  getattr default in _make_fused_ops, matching the existing
  getattr(self, "activation_recompute", False) convention.
- moe/router.py: restore main NVIDIA#7394's global-batch branch in
  TopKRouter.quantile_balancing(), gated on dev's use_quantile_balancing.
- multi_token_prediction.py: black formatting (AST-equivalent).
- test_attention_variant_dsa_gqa.py: cover metric-tracker initialization
  followed by a DSA-GQA training forward, with and without indexer loss.
- test_fine_grained_activation_offloading.py: map expert_fc1 + moe_act to
  main's fused_group_mlp offload-summary group.
- test_multi_token_prediction.py: apply main NVIDIA#8211's cfg.finalize()
  pattern in the dev-only TestMHCMTPIntegration.

Signed-off-by: svcnemo-autobot <svcnemo-autobot@nvidia.com>

This branch was previously deployed

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

Labels

Approved All necessary approvals have been made complexity: low

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants