[main] feat(moe): Support apply wd to qk layernorm for Qwen3-Next (4/4)

megatron/core/optimizer/__init__.py

@@ -60,40 +60,48 @@
     OptimizerConfig,
     ParamKey,
     ParamPredicate,
+    ParamWithNamePredicate,
     SGDOptimizerConfig,
 )
 
 logger = logging.getLogger(__name__)
 
 
-def get_standard_config_overrides(
-    decoupled_lr: float | None = None, decoupled_min_lr: float | None = None
-) -> Dict[ParamKey, ParamGroupOverride]:
+def get_standard_config_overrides(config: OptimizerConfig) -> Dict[ParamKey, ParamGroupOverride]:
     """Get standard config overrides for the optimizer, handling decoupled LR and common wd skips.
 
     Args:
-        decoupled_lr (float | None): decoupled learning rate.
-        decoupled_min_lr (float | None): decoupled minimum learning rate.
+        config (OptimizerConfig): optimizer configuration object.
 
     Returns:
         Dict[ParamKey, ParamGroupOverride]: standard config overrides.
     """
     config_overrides: Optional[Dict[ParamKey, ParamGroupOverride]] = {}
-    if decoupled_lr is not None:
-        decoupled_lr_config: ParamGroupOverride = {"max_lr": decoupled_lr}
-        decoupled_param_key = ParamKey(attr="is_embedding_or_output_parameter")
-        if decoupled_min_lr is not None:
-            decoupled_lr_config["min_lr"] = decoupled_min_lr
-        config_overrides[decoupled_param_key] = decoupled_lr_config
+    # First, figure out how we are going to do wd skipping. The two main approaches are:
+    #  1. The classic megatron approach of skipping all len 1 and bias parameters.
+    #  2. The Qwen3-Next approach of doing 1, other than qk layernorm parameters.
+    if config.apply_wd_to_qk_layernorm:
+        shape_1_not_qkln_param = ParamWithNamePredicate(
+            name="s1_not_qkln",
+            fn=lambda param, name: (len(param.shape) == 1 or name.endswith(".bias"))
+            and not ("q_layernorm." in name or "k_layernorm." in name),
+        )
+        param_wd_mult_key = ParamKey(with_name_predicate=shape_1_not_qkln_param)
+    else:
+        param_length_1_match = ParamPredicate(
+            name="param_len_1", fn=lambda param: len(param.shape) == 1
+        )
+        param_wd_mult_key = ParamKey(name="*.bias", predicate=param_length_1_match)
 
-    # Next construct the standard param group overrides for no weight decay on bias parameters
-    #  as well as any length 1 parameters.
-    param_length_1_match = ParamPredicate(
-        name="param_len_1", fn=lambda param: len(param.shape) == 1
-    )
-    param_wd_mult_key = ParamKey(name="*.bias", predicate=param_length_1_match)
     config_overrides[param_wd_mult_key] = ParamGroupOverride(wd_mult=0.0)
 
+    if config.decoupled_lr is not None:
+        decoupled_lr_config: ParamGroupOverride = {"max_lr": config.decoupled_lr}
+        decoupled_param_key = ParamKey(attr="is_embedding_or_output_parameter")
+        if config.decoupled_min_lr is not None:
+            decoupled_lr_config["min_lr"] = config.decoupled_min_lr
+        config_overrides[decoupled_param_key] = decoupled_lr_config
+
     return config_overrides
 
 
@@ -132,7 +140,7 @@ def _get_param_groups(
         #  the config_overrides argument by default lead to bias parameters and length 1 parameters.
         #  We assume that users of decoupled LR already provide config overrides so will adapt
         #  to the new API.
-        config_overrides = get_standard_config_overrides()
+        config_overrides = get_standard_config_overrides(config=config)
 
     for model_chunk in model_chunks:
         for name, param in model_chunk.named_parameters():

megatron/core/optimizer/optimizer_config.py

@@ -33,6 +33,34 @@ def __call__(self, param: torch.nn.Parameter) -> bool:
         return self.fn(param)
 
 
+@dataclass(frozen=True)
+class ParamWithNamePredicate:
+    """Wraps a matching function to make it hashable for ParamKey.
+    Example:
+        >>> shape_1_not_qkln_param = ParamWithNamePredicate(
+                name="s1_not_qkln",
+                fn=lambda param, name: (
+                    len(param.shape) == 1 or name.endswith(".bias")
+                    and not ("q_layernorm." in name or "k_layernorm." in name)
+                )
+            )
+        >>> shape_1_not_qkln_param(torch.empty(10), "interesting.bias")
+        True
+        >>> shape_1_not_qkln_param(torch.empty(10), "interesting.q_layernorm.bias")
+        False
+
+    NOTE:
+        __hash__ and __eq__ are automatically generated by @dataclass(frozen=True)
+        based solely on 'name' because we set compare=False/hash=False on 'fn'.
+    """
+
+    name: str
+    fn: Callable[[torch.nn.Parameter, str], bool] = field(compare=False, hash=False)
+
+    def __call__(self, param: torch.nn.Parameter, name: str) -> bool:
+        return self.fn(param, name)
+
+
 @dataclass(frozen=True, slots=True)
 class ParamKey:
     """Key to group parameters by. All such grouped parameters can share an
@@ -49,6 +77,15 @@ class ParamKey:
     predicate: Union[ParamPredicate, Tuple[ParamPredicate]] = field(default_factory=tuple)
     """Predicate(s) to match parameters by. If multiple predicates are provided, any must match."""
 
+    with_name_predicate: Union[ParamWithNamePredicate, Tuple[ParamWithNamePredicate]] = field(
+        default_factory=tuple
+    )
+    """
+    Predicate(s) to match parameters with their name. If multiple predicates are provided, 
+      any must match. This is useful if you need to filter out some parameters from an otherwise 
+      positive match by their name.
+    """
+
     def matches(self, param: torch.nn.Parameter, param_name: str) -> bool:
         """Returns true if passed-in parameter (with name) matches `param_key`.
 
@@ -86,6 +123,15 @@ def matches(self, param: torch.nn.Parameter, param_name: str) -> bool:
             for predicate in self.predicate:
                 if predicate(param):
                     return True
+
+        # Check if with_name_predicate matches.
+        if isinstance(self.with_name_predicate, ParamWithNamePredicate):
+            if self.with_name_predicate(param, param_name):
+                return True
+        else:
+            for predicate in self.with_name_predicate:
+                if predicate(param, param_name):
+                    return True
         return False
 
 
@@ -105,9 +151,20 @@ class OptimizerConfig:
     min_lr: Optional[float] = None
     """Minumum value for learning rate. The scheduler clip values below this threshold."""
 
+    decoupled_lr: Optional[float] = None
+    """Separate learning rate for the input and output layer."""
+
+    decoupled_min_lr: Optional[float] = None
+    """Minimum value for learning rate for the input and output layer. The scheduler clip values
+       below this threshold.
+    """
+
     weight_decay: float = 0.01
     """Weight decay coefficient for L2 regularization."""
 
+    apply_wd_to_qk_layernorm: bool = False
+    """If true, apply weight decay to qk layernorm as a special case."""
+
     ##############
     # Precision
     ##############

megatron/training/arguments.py

@@ -925,6 +925,17 @@ def validate_args(args, defaults={}):
         dc = torch.cuda.get_device_capability()
         assert dc[0] >= 8, "Unsupported compute capability for GroupedGEMM kernels."
 
+    if args.no_weight_decay_cond_type is not None:
+        print_rank_0(
+            'WARNING: --no-weight-decay-cond-type is deprecated. Please use --apply-wd-to-qk-layernorm instead.',
+            args.rank,
+        )
+        if args.no_weight_decay_cond_type == "apply_wd_to_qk_layernorm":
+            args.apply_wd_to_qk_layernorm = True
+        else:
+            raise ValueError(f"Invalid no_weight_decay_cond_type: {args.no_weight_decay_cond_type}")
+        args.no_weight_decay_cond_type = None
+
     if args.weight_decay_incr_style == 'constant':
         assert args.start_weight_decay is None
         assert args.end_weight_decay is None
@@ -2077,6 +2088,8 @@ def _add_regularization_args(parser):
                        help='Dropout probability for hidden state transformer.')
     group.add_argument('--weight-decay', type=float, default=0.01,
                        help='Weight decay coefficient for L2 regularization.')
+    group.add_argument('--apply-wd-to-qk-layernorm', action='store_true',
+                       help='Apply weight decay to qk layernorm as a special case.')
     group.add_argument('--clip-grad', type=float, default=1.0,
                        help='Gradient clipping based on global L2 norm.')
     group.add_argument('--adam-beta1', type=float, default=0.9,
@@ -2111,6 +2124,12 @@ def _add_regularization_args(parser):
     group.add_argument('--muon-extra-scale-factor', type=float, default=1.0,
                        help='Additional scale factor for the muon update')
 
+    group.add_argument('--no-weight-decay-cond-type', type=str, choices=['apply_wd_to_qk_layernorm'],
+                       help='Type of no weight decay condition. Choices: '
+                       'None (default): apply weight decay to 1D weights and biases.'
+                       '"apply_wd_to_qk_layernorm": additionally apply weight decay to '
+                       'qk layernorm as a special case.'
+                       'DEPRECATED. Please use --apply-wd-to-qk-layernorm instead. ')
     return parser
 
 

megatron/training/training.py

@@ -1329,7 +1329,7 @@ def get_megatron_optimizer_config(args: Any) -> OptimizerConfig:
 
     # Construct the appropriate config_overrides object. This default handles many cases, but
     #  can be added to as needed by the user, or replaced entirely with a custom override.
-    config_overrides = get_standard_config_overrides(args.decoupled_lr, args.decoupled_min_lr)
+    config_overrides = get_standard_config_overrides(config=config)
 
     return config, config_overrides
 

tests/unit_tests/test_optimizer.py

@@ -21,6 +21,7 @@
     _get_param_groups,
     check_config_overrides_consistency,
     get_megatron_optimizer,
+    get_standard_config_overrides,
 )
 from megatron.core.optimizer_param_scheduler import ParamGroupOverride
 from megatron.core.process_groups_config import ProcessGroupCollection
@@ -45,14 +46,18 @@
 
 
 class Net(nn.Module):
-    def __init__(self):
+    def __init__(self, add_layernorm=False):
         super().__init__()
         self.conv1 = nn.Conv2d(3, 6, 5)
         self.pool = nn.MaxPool2d(2, 2)
         self.conv2 = nn.Conv2d(6, 16, 5)
         self.fc1 = nn.Linear(16 * 5 * 5, 120)
         self.fc2 = nn.Linear(120, 84)
         self.fc3 = nn.Linear(84, 10)
+        if add_layernorm:
+            self.q_layernorm = nn.LayerNorm(10, bias=False)
+            self.k_layernorm = nn.LayerNorm(10, bias=False)
+            self.layernorm = nn.LayerNorm(10, bias=False)
 
     def forward(self, x):
         x = self.pool(F.relu(self.conv1(x)))
@@ -206,6 +211,81 @@ def test_get_param_groups_overlapping_matches(mock_get_world_size):
     assert param_groups[2]['max_lr'] == 0.01
 
 
+@patch('torch.distributed.get_world_size', return_value=1)
+@patch(
+    'torch.distributed.all_gather_object', lambda output_list, obj: output_list.__setitem__(0, obj)
+)
+def test_get_param_groups_with_standard_config_overrides(apply_wd_to_qk_layernorm: bool):
+    """In this test, we see if the standard config overrides are applied correctly."""
+
+    # Initialize the model with layernorm
+    net = Net()
+
+    config = OptimizerConfig(optimizer='adam', lr=0.01)
+    config_overrides = get_standard_config_overrides(config=config)
+    param_groups = _get_param_groups([net], config, config_overrides)
+
+    assert len(param_groups) == 2
+    p_set = set(net.parameters())
+
+    assert p_set == set(param_groups[0]['params']) | set(param_groups[1]['params'])
+    assert len(p_set) == len(param_groups[0]['params']) + len(param_groups[1]['params'])
+    assert param_groups[0]['wd_mult'] == 0.0 or param_groups[1]['wd_mult'] == 0.0
+    assert param_groups[0]['wd_mult'] == 1.0 or param_groups[1]['wd_mult'] == 1.0
+    assert len(param_groups[0]['params']) > 0 and len(param_groups[1]['params']) > 0
+
+    # Both param groups should have 5 parameters.
+    # Param group A (wd_mult=1.0): conv1.weight, conv2.weight, fc1.weight, fc2.weight, fc3.weight
+    # Param group B (wd_mult=0.0): conv1.bias, conv2.bias, fc1.bias, fc2.bias, fc3.bias
+    assert len(param_groups[0]['params']) == 5, (
+        f"Expected 5 parameters in the first param group, "
+        f"but got {len(param_groups[0]['params'])}"
+    )
+    assert len(param_groups[1]['params']) == 5, (
+        f"Expected 5 parameters in the second param group, "
+        f"but got {len(param_groups[1]['params'])}"
+    )
+
+
+@patch('torch.distributed.get_world_size', return_value=1)
+@patch(
+    'torch.distributed.all_gather_object', lambda output_list, obj: output_list.__setitem__(0, obj)
+)
+def test_get_param_groups_appling_wd_to_qk_layernorm(apply_wd_to_qk_layernorm: bool):
+    """In this test, we see if the `apply_wd_to_qk_layernorm` config is applied correctly."""
+
+    # Initialize the model with layernorm
+    net = Net(add_layernorm=True)
+
+    config = OptimizerConfig(
+        optimizer='adam', lr=0.01, apply_wd_to_qk_layernorm=apply_wd_to_qk_layernorm
+    )
+    config_overrides = get_standard_config_overrides(config=config)
+    param_groups = _get_param_groups([net], config, config_overrides)
+
+    assert len(param_groups) == 2
+    p_set = set(net.parameters())
+
+    assert p_set == set(param_groups[0]['params']) | set(param_groups[1]['params'])
+    assert len(p_set) == len(param_groups[0]['params']) + len(param_groups[1]['params'])
+    assert param_groups[0]['wd_mult'] == 1.0
+    assert param_groups[1]['wd_mult'] == 0.0
+
+    # There are two param groups, having 7, and 6 parameters respectively.
+    # Param group A (wd_mult=1.0): conv1.weight, conv2.weight, fc1.weight, fc2.weight, fc3.weight,
+    #    q_layernorm.weight, k_layernorm.weight
+    # Param group B (wd_mult=0.0): conv1.bias, conv2.bias, fc1.bias, fc2.bias, fc3.bias,
+    #    layernorm.weight
+    assert len(param_groups[0]['params']) == 7, (
+        f"Expected 5 parameters in the first param group, "
+        f"but got {len(param_groups[0]['params'])}"
+    )
+    assert len(param_groups[1]['params']) == 6, (
+        f"Expected 6 parameters in the second param group, "
+        f"but got {len(param_groups[1]['params'])}"
+    )
+
+
 def test_chained_optimizer():
     net = Net()
     optimizer_1 = Adam(list(net.parameters())[:2], lr=0.01)