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.

[mkhona-nvidia]

Hi @plugyawn thank you for this, this a long-overdue change to megatron.

But this is a hugely massive change that hits several parts of the codebase -- what do you think of the approach of splitting this PR into multiple PRs, each of which is modular and self-contained in the following way:

• First, build the infrastructure and plugs that can modify each layer's init and hyperparameter multipliers. Use this infra to show the specific choices made by megatron's defaults. You already have the right abstractions for this with the policy.

• Next PR, use this infra to provide a regular AdamW-style muP init + hyperparameter multiplier

• Next PR, use this infra to provide a Muon-friendly spectral muP

[plugyawn]

Hi @janEbert, @mkhona-nvidia! The three PR route makes sense to me: just wanted to clarify that we also have a functioning Depth-MuP now: not just AdamW width-MuP and Muon-specific spectral width-MuP.

In the meantime, I'll share the WIP design doc too.
It covers the depth-MuP recipe in detail as well, along with the relevant loss/depth plots.

Took a while to work with Codex to make a design-doc skill, I hope it's the format that you expected. The comments in red were my latest comments while iterating. Will remove it in the final version.
Design doc

For the 3(4?) PRs, I'll refactor the design doc to mention which parts are relevant to which stage of the implementation.

[plugyawn]

Added the first of the PRs! Have the next 3 PRs roughly scoped out as well.
Also, recently saw some work that suggested that only scaling the embedding layer might work as well as MuP in some cases, so perhaps that's also worth looking at?

[plugyawn]

:( @janEbert have we given up on this set of PRs?

[janEbert]

Hey @plugyawn, this is not given up at all. Other responsibilities came up, which started to take a substantial amount of time. I understand it's frustrating with the amount of work and constantly changing codebase (especially with all the work you put into addressing our wishes to make it easier to review), and I'm very sorry about the delay. However, please also understand that, even assuming that we had established common ground for these changes before their implementation, it's a lot of painstaking work to properly review this type of feature due to its highly theoretical nature and reliance on mathematical correctness. Add to that the fundamental design decisions, which are also important to get right.

That's why I'd prefer to review this properly instead of with a quick look. Sadly, this takes a lot of focussed time, which – sadly – I haven't found recently. I'm assuming that I'll find time toward the second half of June, so please bear with me until then.

[plugyawn]

No worries, I completely understand, sorry for the confusion!

[plugyawn]

@janEbert small reminder!