diff --git a/src/megatron/bridge/training/config.py b/src/megatron/bridge/training/config.py index 7698322b62..79aa3bc60f 100644 --- a/src/megatron/bridge/training/config.py +++ b/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 diff --git a/src/megatron/bridge/training/optim.py b/src/megatron/bridge/training/optim.py index 9722aadda7..c0e1f82517 100644 --- a/src/megatron/bridge/training/optim.py +++ b/src/megatron/bridge/training/optim.py @@ -12,66 +12,22 @@ # 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( @@ -79,6 +35,7 @@ def setup_optimizer( 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( diff --git a/src/megatron/bridge/training/setup.py b/src/megatron/bridge/training/setup.py index 7cfeea56d6..ca7b4f8fe1 100644 --- a/src/megatron/bridge/training/setup.py +++ b/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")