General scaling policy for HP transfer + Depth MuP recipe

[plugyawn]

What does this PR do ?

Addresses part of #4088, introducing a scaling policy for transfer recipes.
This also subsumes and refactors #3058 and #3715 into a mup recipe.
For now, I've kept the --use-mup flag connected (since it was part of the last release), but I'll add in a warning for deprecation?

Current recipes:

• none (default) being the standard Megatron parameterization
• mup being the refactored existing MuP behavior, preserved exactly.
• depth_mup (draft) with support for Adam/AdamW + dense residual transformers to demonstrate support for other scaling recipes.

depth_mup recipe

We do the following:

• dense residual branch multiplier: depth_mult^-1
• hidden matrix-like LR depth multiplier: depth_mult^0
• hidden Adam epsilon depth multiplier: depth_mult^-1

And also,

• a dense block output-projection init compensation: depth_mult^+0.5
  • Megatron already uses layer-count-based residual output-projection initialization. This allows us to avoid accidentally double-counting depth at initialization while making the residual branch scaling explicit in the forward graph.

Chiefly, this PR adds a new megatron.core.parameterization package, which handles:

• canonical scaling recipe configs
• final resolved model layer and trainer/optimizer scaling policies
• parameter-role helpers for embedding/output/shared/vector/matrix-like parameters, for support extending to Muon.

For unsupported optimizers like SGD, we raise an error for now. Some math discovery might be necessary here. SGD depth transfer appears to need explicit hidden-weight, hidden-bias laws that would be complex to implement in one go.

Plots for MuP: Adam, SGD, Muon respectively:



240 iterations, wikitext8, on an A100 80 GB.

Needs some discussion on the YAML changing and the resume loop.

Some more thought: I think a scaling policy and training recipes could be a great feature moving forward as we can do more with the CLI, considering codex seems quite adept at using CLIs in general! I understand it's a more convenience-based idea, but would be exciting to see first-class support for autonomous and principled pretraining runs.

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

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.

[github-actions]

This PR has been automatically converted to draft because all PRs must start as drafts.

When you are ready for review, click Ready for Review to begin the review process. This will:

1. Add the oncall reviewer (optional reviewer)
2. Add required review teams based on your changes

See the contribution guide for more details.

[plugyawn]

Adding the plots for depth-MuP soon.

[plugyawn]

@Phlip79 @janEbert would love for you to have a look!

[plugyawn]

Shifting Depth-MuP implementation from TP VI (Yang's) to the Bytedance version, since it covers GPT-style transformers more comfortably. The paper is about MuonClip/Kimi-Muon, but also talks about Adam/AdamW.

[plugyawn]

Preliminary depth-MuP results on wikitext8, 300M tokens; depths 4, 8, 16, and 48 underway!

Edit:

For width: 2048, across 4 GPUs (4xA100 80GB)

[plugyawn]

Thanks for the review @janEbert! Could you take a look at the depth-MuP plots too? The optimas seem to align well, and the curves seem bundled in the right way, but I'm not sure about the tail behaviour. Lower depths here prefer slightly higher LRs, which makes sense going by general wisdom.

[janEbert]

Hey, this is a huge PR, so it will take some time to review. :)
Thanks for your patience in advance!

[janEbert]

Hey @plugyawn, we discussed in a smaller round how to tackle this PR and came to the conclusion that a design doc would be extremely helpful in order to better understand and review the PR. Ideally, you would include explanatory diagrams, considerations of edge cases, and explanations on why certain parts of the code needed to be touched. For example, the plots you already produced would also be perfect for this kind of document.

This is also part of the PR template, so I hope you don't feel too thrown off by this request:

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

[plugyawn]

Fixed some bugs and redid the depth-MuP plots with width 2048, and verified across 4xA100s. Preparing the design doc, will share it soon!

[janEbert]

@plugyawn really appreciate you taking time for the additional input! Just tagging @NVIDIA/mcore-oncall should be good enough (I'll update the PR template with this; "@mcore-oncall" seems to be outdated). Oncall is a rotating member of our team and, coincidentally, it's me right now, so feel free to just ask here/reach out via mail. :)

Stacked PRs are not currently enabled in this repository.