Design Discussion For Parameter-Level Initialization/Optimization Configurability

## High Level Problem Description

P0: I want to assign each parameter its own initialization and its own optimizer.

P0: I want to be able to configure the values passed to the initialization function and optimizer.

P1: I want the initialization function and optimizer constructor to have access to training process information like model width, model depth, parameter layer index, training duration, etc. If necessary I can supply all of this by managing my external config.

## Extended Problem Description

Principled scaling schemes such as ABC-parameterizations (muP) or norm-constrained optimizers (Muon/Shampoo/Scion/etc) have spawned large research areas to support stable and performant optimization across scales. A general trend in both optimization research and large-scale training is to enable per-parameter optimizers and per-parameter initializations. In optimal recipes, both initialization and optimizer are functions of layer type and layer characteristics.

These recipes are hard to express in Megatron-LM because per-parameter overrides are not exposed through [OptimizerConfig](https://github.com/NVIDIA/Megatron-LM/blob/cf2f07d7b1315c96c05554c670c43207c6783e5e/megatron/core/optimizer/optimizer_config.py#L139C7-L139C22) but rather through internal functions [see Emerging Optimizers Examples](https://github.com/NVIDIA/Megatron-LM/blob/cf2f07d7b1315c96c05554c670c43207c6783e5e/megatron/core/optimizer/emerging_optimizers.py). Per-parameter overrides need to be directly configurable from training entrypoints [link](https://github.com/NVIDIA/Megatron-LM/blob/cf2f07d7b1315c96c05554c670c43207c6783e5e/pretrain_gpt.py) because Megatron-LM does not expose any detailed configuration override options for the optimizer through TransformerConfig.

The same issue is present for initialization. Granular initialization options are not accessible in `TransformerConfig` [link](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/transformer/transformer_config.py). Recent changes introduced [muP scaled inits](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/transformer/transformer_config.py#L39) which enable one scaling init recipe, but not flexible granular recipes.

See also https://github.com/NVIDIA/Megatron-LM/pull/4381 for great existing work trying to get these scaling ideas into Megatron-LM.

## Motivating Examples

For an example of "mixed optimizers", consider the current Megatron-LM Muon optimizer. This optimizer partitions parameter groups based on per-parameter optimizer type [link](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/optimizer/__init__.py#L975). which is actually Muon on some layers and Adam on others.

Per-parameter learning rates and per-parameter initializations are required for stable MoE learning across scales when using Adam [link](https://arxiv.org/abs/2605.14200).

For another example of “mixed initializations” see current [muP work in Megatron-LM](https://github.com/NVIDIA/Megatron-LM/blob/cf2f07d7b1315c96c05554c670c43207c6783e5e/megatron/core/utils.py#L840) which provides a re-scaled initialization function and uses additional muP context. For an example of a broad range of required initializations to support advanced optimization work, see [linked Scion paper](https://arxiv.org/pdf/2502.07529) Table 3 and Table 4, which demonstrate how changing modeling/norm choices leads to different initialization requirements.

## What already exists in Megatron-LM

A mix of optimizer config overrides are enabled to support both emerging optimizers and muP scaling in https://github.com/NVIDIA/Megatron-LM/blob/cf2f07d7b1315c96c05554c670c43207c6783e5e/megatron/core/optimizer/__init__.py.

A related configuration pattern appears for mixed quantization strategies: https://github.com/NVIDIA/Megatron-LM/tree/main/megatron/core/quantization

If we follow the quantization config approach above we would de-couple the matching logic from the override assignment. I would not duplicate the quantization approach because layer types which often share scaling recipes like FC1 and FC2 are still likely to have different optimal weight/init multipliers.

## Describe the solution you'd like

Below is one solution candidate (for optimizer) and some open design decisions for initialization. Suggestions/modifications are very welcome especially on initialization. See also [Existing MR for existing ideas on implementation](https://github.com/NVIDIA/Megatron-LM/pull/4381/changes#diff-4550b4a545b6aeb7ef82616f380f3a189343727bc13c3fab38b9762554810d40)

### Proposed Optimization solution
Extend `OptimizerConfig` with a `config_overrides` field. This new field can accept either a yaml filepath or a config object. The config structure will be match_glob→param_override_name→param_override_value. We could enable both `param_override_mult` and/or `param_override_value`. Then if the key contains the suffix `_mult` we multiply and `_value` we replace. I prefer to either support both, or `_value`-only to avoid users having to carefully track base parameters to resolve multiplied values. The `config_overrides` subconfig should enable per-parameter overrides like optimizer, lr, epsilon etc. This approach must support arbitrary valid optimizer overrides. “Valid” means that we cannot toggle global choices layerwise optimization on/off at the subconfig override level. Enabling `config_overrides` should be mutually exclusive to other approaches like the current muP config overrides or layernorm weight decay 0 overrides.

### Sketch of Initialization Challenges
We should support configurable initialization through an `InitializationConfig` object. Similar to the new `config_overrides` field above, the `InitializationConfig` object should enable glob or regex-style matching and allow the user to specify either a current Megatron function or a custom function.

**Design Decision 1: Do we initialize or re-initialize parameters?**
We can either attempt to initialize all parameters correctly from scratch, or call a re-initialization function after construction unless we are checkpoint loading. Direct initialization requires a lot more plumbing changes. Re-initialization after model construction completes can run into issues like accidental parameter de-synchronization across DP ranks if not checked carefully. So we need to carefully track RNG states.

**Decision Decision 2: How should parameters role and parameter context be identified?**
If parameters are responsible for declaring their context (layer id, model total depth, layer input/output width, layer role, desired layer input/output norms) then we can rely on individual modules to store their context and the user only needs to assign functions which apply to that context. If a “global brain” must identify every parameter and its role, we assume that we can correctly deduce model parameter attributes like “un-sharded weight width” post-initialization.

Example `OptimizerConfig`
```
# Existing OptimizerConfig fields remain the defaults for unmatched parameters.
optimizer: adam
lr: 2.0e-4
min_lr: 2.0e-5
weight_decay: 0.1
adam_eps: 1.0e-8

config_overrides:
require_context_for_parameter_modules: true # require that modules provide their own inputs
unmatched: use_global_defaults # default to global optimization parameters if unmatched
conflicts: error # no accidental multiple group assignment

groups:
- name: norm
match: '(^|\.)(norm|layer_norm|layernorm)(\.|$)'
set:
optimizer: adam
max_lr: 1.0e-3
min_lr: 1.0e-4
eps: 1.0e-14
weight_decay: 0.0

- name: fc1_hidden
match: '.*fc1.*'
set:
optimizer: muon
optimizer_kwargs:
momentum: 0.95
muon_scale_mode: fan_in
- name: output
match: '.*output.*'
set:
optimizer: muon
optimizer_kwargs:
momentum: 0.99

```

Example `InitializationConfig`

```
# Custom initializer targets receive the tensor, ParameterContext,
# TransformerConfig, and a deterministic RNG generator.

init_config:
require_context_for_parameter_modules: true
unmatched: use_default_initializer
conflicts: error
apply: after_materialization_before_precision_or_distributed_wrapping

default_initializer:
target: torch.nn.init.normal_
kwargs:
mean: 0.0
std: 0.002

groups:
- name: biases
match:
name_regex:
- '\.bias$'
initializer:
target: torch.nn.init.zeros_

- name: norm_weights
match:
context:
parameter_role:
- rmsnorm_weight
- qk_norm_weight
initializer:
target: megatron.core.utils.fancy_norm_init # function of layer index and depth
```