Refactor the optimizer override function so that users can swap in their own (#2010)

Signed-off-by: John St. John <jstjohn@nvidia.com>

Author: jstjohn

src/megatron/bridge/training/config.py

@@ -24,7 +24,12 @@
 from megatron.core.datasets.gpt_dataset import GPTDatasetConfig as MCoreGPTDatasetConfig
 from megatron.core.distributed import DistributedDataParallelConfig as MCoreDistributedDataParallelConfig
 from megatron.core.optimizer import OptimizerConfig as MCoreOptimizerConfig
+from megatron.core.optimizer import (
+    ParamGroupOverride,
+    ParamKey,
+)
 from megatron.core.transformer.enums import AttnBackend
+from megatron.core.transformer.module import MegatronModule
 
 from megatron.bridge.data.datasets.packed_sequence import PackedSequenceSpecs
 from megatron.bridge.models import GPTModelProvider, T5ModelProvider
@@ -258,6 +263,81 @@ class DatasetBuildContext:
     tokenizer: Optional[MegatronTokenizer] = None
 
 
+@dataclass(frozen=True)
+class OptimizerConfigOverrideProviderContext:
+    """Context for providing config overrides."""
+
+    scheduler_config: "SchedulerConfig"
+    optimizer_config: OptimizerConfig
+    model: Union[MegatronModule, list[MegatronModule]]
+
+
+@dataclass
+class OptimizerConfigOverrideProvider:
+    """Abstract base class for providing config overrides."""
+
+    def build_config_overrides(
+        self, context: OptimizerConfigOverrideProviderContext
+    ) -> dict[ParamKey, ParamGroupOverride] | None:
+        """Build config overrides for weight decay based on scheduler configuration.
+
+        This function creates parameter-specific overrides for weight decay behavior.
+        By default, weight decay is skipped for bias parameters and 1D parameters.
+        For Qwen3-Next models, weight decay is applied to q_layernorm and k_layernorm.
+
+        Args:
+            context: OptimizerConfigOverrideProviderContext which packages the scheduler
+                configuration, optimizer configuration, and model.
+
+        Returns:
+            Dictionary of ParamKey to ParamGroupOverride for the optimizer
+        """
+        model = context.model
+        scheduler_config = context.scheduler_config
+        optimizer_config = context.optimizer_config
+
+        config_overrides: dict[ParamKey, ParamGroupOverride] = {}
+
+        # Collect param names that should skip weight decay
+        # NOTE: this can be simplified once https://github.com/NVIDIA/Megatron-LM/pull/2753
+        #  is merged into dev. Then we can re-use megatron's apply_wd_to_qk_layernorm option
+        #  and call megatron.core.optimizer.get_standard_config_overrides(optimizer_config)
+        #  directly for standard settings, replacing the custom logic below for qwen3-next.
+        no_wd_names: list[str] = []
+        is_qwen3_next = scheduler_config.no_weight_decay_cond_type == "qwen3_next"
+
+        model_list = model if isinstance(model, list) else [model]
+        for model_chunk in model_list:
+            for name, param in model_chunk.named_parameters():
+                # Skip weight decay for bias parameters
+                if name.endswith(".bias"):
+                    no_wd_names.append(name)
+                    continue
+
+                # Skip weight decay for 1D parameters
+                if len(param.shape) == 1:
+                    if is_qwen3_next:
+                        # Qwen3-Next: apply weight decay to qk layernorm (don't add to skip list)
+                        if "q_layernorm" in name or "k_layernorm" in name:
+                            continue
+                    no_wd_names.append(name)
+
+        # Create a single ParamKey with all names that should skip weight decay
+        if no_wd_names:
+            no_wd_key = ParamKey(name=tuple(no_wd_names))
+            config_overrides[no_wd_key] = ParamGroupOverride(wd_mult=0.0)
+
+        # Now handle decoupled LR:
+        if optimizer_config.decoupled_lr is not None:
+            decoupled_lr_config: ParamGroupOverride = {"max_lr": optimizer_config.decoupled_lr}
+            decoupled_param_key = ParamKey(attr="is_embedding_or_output_parameter")
+            if optimizer_config.decoupled_min_lr is not None:
+                decoupled_lr_config["min_lr"] = optimizer_config.decoupled_min_lr
+            config_overrides[decoupled_param_key] = decoupled_lr_config
+
+        return config_overrides if config_overrides else None
+
+
 @dataclass
 class DatasetProvider(DataloaderConfig, ABC):
     """Abstract base class for custom dataset configurations.
@@ -1163,6 +1243,9 @@ class ConfigContainer(Container):
     train: TrainingConfig
     model: GPTModelProvider | T5ModelProvider | MambaModelProvider
     optimizer: OptimizerConfig
+    optimizer_config_override_provider: OptimizerConfigOverrideProvider = field(
+        default_factory=OptimizerConfigOverrideProvider
+    )
     ddp: DistributedDataParallelConfig = field(default_factory=DistributedDataParallelConfig)
     scheduler: SchedulerConfig
     dataset: GPTDatasetConfig | FinetuningDatasetConfig | DatasetProvider

src/megatron/bridge/training/optim.py

@@ -12,73 +12,30 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Dict, Optional, Union
+from typing import Optional, Union
 
 from megatron.core.optimizer import (
     MegatronOptimizer,
     OptimizerConfig,
-    ParamKey,
     get_megatron_optimizer,
 )
 from megatron.core.optimizer.muon import get_megatron_muon_optimizer
-from megatron.core.optimizer_param_scheduler import OptimizerParamScheduler, ParamGroupOverride
+from megatron.core.optimizer_param_scheduler import OptimizerParamScheduler
 from megatron.core.transformer.module import MegatronModule
 
-from megatron.bridge.training.config import SchedulerConfig
-
-
-def _build_config_overrides(
-    scheduler_config: SchedulerConfig,
-    model: Union[MegatronModule, list[MegatronModule]],
-) -> Optional[Dict[ParamKey, ParamGroupOverride]]:
-    """Build config overrides for weight decay based on scheduler configuration.
-
-    This function creates parameter-specific overrides for weight decay behavior.
-    By default, weight decay is skipped for bias parameters and 1D parameters.
-    For Qwen3-Next models, weight decay is applied to q_layernorm and k_layernorm.
-
-    Args:
-        scheduler_config: Scheduler configuration containing weight decay settings
-        model: The model or list of model chunks to collect parameter names from
-
-    Returns:
-        Dictionary of ParamKey to ParamGroupOverride for the optimizer
-    """
-    config_overrides: Dict[ParamKey, ParamGroupOverride] = {}
-
-    # Collect param names that should skip weight decay
-    no_wd_names: list[str] = []
-    is_qwen3_next = scheduler_config.no_weight_decay_cond_type == "qwen3_next"
-
-    model_list = model if isinstance(model, list) else [model]
-    for model_chunk in model_list:
-        for name, param in model_chunk.named_parameters():
-            # Skip weight decay for bias parameters
-            if name.endswith(".bias"):
-                no_wd_names.append(name)
-                continue
-
-            # Skip weight decay for 1D parameters
-            if len(param.shape) == 1:
-                if is_qwen3_next:
-                    # Qwen3-Next: apply weight decay to qk layernorm (don't add to skip list)
-                    if "q_layernorm" in name or "k_layernorm" in name:
-                        continue
-                no_wd_names.append(name)
-
-    # Create a single ParamKey with all names that should skip weight decay
-    if no_wd_names:
-        no_wd_key = ParamKey(name=tuple(no_wd_names))
-        config_overrides[no_wd_key] = ParamGroupOverride(wd_mult=0.0)
-
-    return config_overrides if config_overrides else None
+from megatron.bridge.training.config import (
+    OptimizerConfigOverrideProvider,
+    OptimizerConfigOverrideProviderContext,
+    SchedulerConfig,
+)
 
 
 def setup_optimizer(
     optimizer_config: OptimizerConfig,
     scheduler_config: SchedulerConfig,
     model: Union[MegatronModule, list[MegatronModule]],
     use_gloo_process_groups: bool = False,
+    optimizer_config_override_provider: Optional[OptimizerConfigOverrideProvider] = None,
 ) -> tuple[MegatronOptimizer, OptimizerParamScheduler]:
     """Set up the optimizer and scheduler.
 
@@ -91,8 +48,13 @@ def setup_optimizer(
     Returns:
         tuple containing the optimizer and scheduler
     """
+    if optimizer_config_override_provider is None:
+        optimizer_config_override_provider = OptimizerConfigOverrideProvider()
+
     # Build config overrides for weight decay based on scheduler config and model params
-    config_overrides = _build_config_overrides(scheduler_config, model)
+    config_overrides = optimizer_config_override_provider.build_config_overrides(
+        OptimizerConfigOverrideProviderContext(scheduler_config, optimizer_config, model)
+    )
 
     if "muon" not in optimizer_config.optimizer and "soap" not in optimizer_config.optimizer:
         optimizer = get_megatron_optimizer(

src/megatron/bridge/training/setup.py

@@ -226,6 +226,7 @@ def modelopt_pre_wrap_hook(model):
         scheduler_config=cfg.scheduler,
         model=model,
         use_gloo_process_groups=cfg.dist.use_gloo_process_groups,
+        optimizer_config_override_provider=cfg.optimizer_config_override_provider,
     )
     timers("model-and-optimizer-setup").stop()
     barrier_and_log("after model, optimizer, and learning rate scheduler are built")