Scaling for MuP over Muon optimizer.

[plugyawn]

What does this PR do ?

Addresses #3177, make default Muon scaling when conducting MuP transfer behave as it should.
By default, Megatron scales Muon's matrix sign directions by $\sqrt{\max(\text{fan}_{in}, \text{fan}_{out})}$ (spectral). An alternate option is to use unit_rms_norm, or $\sqrt{\frac{\text{fan}_{out}}{\text{fan}_{in}}}$.

• Earlier, as a bug, turning MuP on would also apply Adam-wise LR/eps updates on the Muon-managed matrices, making MuP unprincipled on Muon (the matrices would be scaled both by Muon and by MuP). This has been fixed.

• In addition, even without double scaling, Muon-managed matrices scaled with the default spectral, loses width-invariance due to the max(). For square-matrices, for example, it grows as $\sqrt{n}$

The principled way to go forward seems to be to set Muon scaling to unit_rms_norm, i.e, $\sqrt{\frac{\text{fan}_{out}}{\text{fan}_{in}}}$, and so, we default to Muon with unit_rms_norm when MuP is active.

L to R: SP, MuP+spectral, MuP+unit_rms_norm (new default), old MuP double-scaling + spectral (old default)

L: SP, R: MuP+spectral vs MuP + unit_rms_norm. As expected, unit_rms_norm shows MuP transfer better than spectral on these sweeps.

Experiments on a slice of wikitext-103, 12 layer transformer, 100 iterations, if not stated otherwise. Confidence bounds from 5 runs with random seeds.

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

flowchart LR
    A[Pre-checks] --> B[PR Tests]
    subgraph Code Review/Approval
        C1[Expert Review] --> C2[Final Review]
    end
    B --> C1
    C2 --> D[Merge]

Loading

Pre-checks

• I want this PR in a versioned release and have added the appropriate Milestone (e.g., Core 0.8)
• 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 (bugfix, assuming further documentation not necessary?)
• I have run the autoformatter.sh on my PR

Code review

The following process is enforced via the CODEOWNERS file for changes into megatron/core. For changes outside of megatron/core, it is up to the PR author whether or not to tag the Final Reviewer team.

For MRs into `main` branch

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!

(Step 1): Add PR label Expert Review

(Step 2): Collect the expert reviewers reviews

1. Attach the Expert Review label when your PR is ready for review.
2. GitHub auto-assigns expert reviewers based on your changes. They will get notified and pick up your PR soon.

⚠️ Only proceed to the next step once all reviewers have approved, merge-conflict are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

(Step 3): Final Review

1. Add Final Review label
2. GitHub auto-assigns final reviewers based on your changes. They will get notified and pick up your PR soon.

(Optional Step 4): Cherry-pick into release branch

If this PR also needs to be merged into core_r* release branches, after this PR has been merged, select Cherry-pick to open a new PR into the release branch.

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.

Merging your PR

Any member of core-adlr and core-nemo will be able to merge your PR.

[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.

[BoxiangW]

Thanks for your PR! The loss v. lr graph looks great!

[BoxiangW]

LGTM

[plugyawn]

@BoxiangW thank you for the review! Added the arguments.py help string. Should be good for merge now, I think.

[plugyawn]

@skyw, I've updated the warn behavior on spectral vs unit_rms_norm. It now explicitly throws a warning whenever MuP is used without unit_rms_norm, but doesn't change the defaults.

[skyw]

some left over needs to be cleaned. otherwise LGTM

[plugyawn]

@skyw thanks! Made the changes.

[plugyawn]

Ran a few more exps to see how compensatory scaling on the Muon (2D) matrices on spectral behaves.

Left is MuP while letting Muon handle the 2D scaling in spectral. As we see, it still does fine.

Right is MuP + Muon with 2D scaling in spectral, with an extra multiplicative compensation term of 1/sqrt(width_mult) on the 2D matrices, to cancel the width growth of spectral scaling and restore width-invariant update magnitude across widths.

The best achieved loss seems to be lower, and the clustering of the minima appears to be tighter, but that might be an artifact of the steps to train.

I've kept it out of the PR for now since I've not found it used in published literature (so far) and it seems to make sense to have a standing warning (considering it does still work rather fine).

[plugyawn]

@BoxiangW not sure why the CI is still failing... I actually ran out of compute 😭, but I'll try to get some more and update this.

Update: I asked Codex, and it suggested that the newer GB200 workflow is getting skipped, since this is a non-maintainer PR, but I'm unsure if that's the case.

---

Edit: the exact diagnosis is:

  - The internal CI branch aac1e1b14 merged a newer main.
  - That newer workflow added GB200 integration jobs:
      - cicd-parse-integration-tests-gb200
      - cicd-integration-tests-latest-gb200
  - Those jobs are maintainer-only in .github/workflows/cicd-main.yml:
      - needs.is-not-external-contributor.outputs.is_maintainer == 'true'
  - For this PR they get skipped.
  - Nemo_CICD_Test then counts skipped jobs globally and treats them as fatal because SKIPPING_IS_ALLOWED=false.

[BoxiangW]

/ok to test 9849e84

[plugyawn]

The MLA-Yarn unit test fail looks unrelated to the MuP changes. Perhaps a transient error?

[BoxiangW]

/ok to test a86e420

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

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

[plugyawn]

Thanks everyone!