[plugyawn]

What does this PR do ?

PR 1/4 of the planned scaling-recipes stack.
Design doc: PR_MuP.pdf
Adds

• internal scaling-policy infrastructure used by the follow-up mup, depth_mup, and Muon interaction PRs, without introducing the public scaling-recipe surface yet, as planned in General scaling policy for HP transfer + Depth MuP recipe #4381.

This PR:

• Adds megatron.core.parameterization with resolved scaling context, model policy, training policy, and parameter-role helpers.
• Routes existing model-side MuP-sensitive behavior through policy hooks for embeddings, logits, attention scaling, initialization, and transformer-layer construction.
• Routes optimizer override construction through the resolved training policy while preserving existing behavior.
• Preserves scaling metadata through model and FSDP parameter paths.
• Makes optimizer param-group resume matching stable for older checkpoints by:
  • excluding mutable scheduler values such as max_lr / min_lr from group identity;
  • default-normalizing missing legacy identifier fields to runtime defaults.
• Adds internal documentation and regression coverage for the policy seam and optimizer group compatibility.

This is intentionally the infrastructure slice only. Follow-up PRs layer on:

1. feature/scaling-recipe-mup: canonical width-MuP recipe API.
2. feature/depth-mup: experimental dense-GPT depth-MuP support.
3. feature/muon-spectral-mup: Muon scalar optimizer interaction.

Issue tracking

Linked issue: Related to #4381.

Testing

• Added unit coverage for scaling-policy infrastructure behavior.
• Added optimizer regression coverage for stable group matching and legacy missing identifier defaults.
• Local smoke: optimizer group identifiers from legacy checkpoint-style groups match runtime defaults.
• BASE_REF=main SKIP_DOCS=true ./tools/autoformat.sh
• git diff --check
• git diff --name-only upstream/main...HEAD -- '*.py' | xargs python3 -m py_compile

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

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

Hi @janEbert @mkhona-nvidia! Planning to split this into 4 PRs as described. Design doc is attached.

[plugyawn]

Tagging @NVIDIA/mcore-oncall; the design discussion for this PR happened in #4381.

[plugyawn]

@janEbert should I start adding the next PRs too?

[janEbert]

If they're already done, why not? It's nice to see the overarching "story" and wrap-up, too. Thank you for the additional work in making this more manageable!

[plugyawn]

Added!

[plugyawn]

Made a few changes: some naming changes (ResolvedScalingContext --> ScalingContext, ResolvedScalingPolicy --> ModelScalingPolicy, and a few others), optimizer-group resume tests.
One question: the old #3058 MuP implementation can be completely deprecated after #4998; should we have a separate PR to drop it, or remove it in #4998?

[plugyawn]

Hi @mkhona-nvidia @janEbert! Have added the rest of the PRs, waiting on the reviews!

[plugyawn]

@janEbert just a bump!

[janEbert]

Hey, what I like about the initial design here is that it has a very explicit goal, hardcoding things where it makes sense to enable parametrization support but nothing else, which keeps things simple.

However, I think that's also a limitation that we'd like to avoid integrating before other design decisions. I'd prefer we start with general user-facing optimizer and initialization configuration APIs, which would allow other modern related configurations as well. I've had the feeling that modern LLM trainings really like to make use of per-layer configuration and enabling more general configuration options would benefit both goals in the long-term, enabling another dimension of user-freedom. The proposed scaling recipes could then be worked on top of those general APIs. (Although we'd prefer these to stay as examples. We would probably not maintain a list of parametrization policies in the core codebase.)

I think it makes sense to first consider these questions, which were also raised in #5732, and work from there.