update

Author: Jintao-Huang

swift/megatron/arguments/megatron_args.py

@@ -731,6 +731,8 @@ def __post_init__(self):
         os.environ.setdefault('CUDA_DEVICE_MAX_CONNECTIONS', '1')
         if self.recompute_granularity == 'none':
             self.recompute_granularity = None
+        if self.apply_wd_to_qk_layernorm and args.hf_model_type != 'qwen3_next':
+            raise ValueError('apply_wd_to_qk_layernorm is only supported for qwen3_next')
         self._set_default()
         self.model_info, self.model_meta = get_model_info_meta(
             self.model, model_type=self.model_type, use_hf=self.use_hf, hub_token=self.hub_token)

swift/megatron/trainers/base.py

@@ -337,8 +337,8 @@ def _get_param_groups(
 
                 if no_weight_decay_cond is not None:
                     no_wd: bool = no_weight_decay_cond(name, param)
-                elif args.apply_wd_to_qk_layernorm and any(name.endswith(k) for k in ['q_layernorm.weight', 'k_layernorm.weight']):
-                    # assert args.hf_model_type == 'qwen3_next', 'currently only support qwen3_next'
+                elif args.apply_wd_to_qk_layernorm and any(
+                        name.endswith(k) for k in ['q_layernorm.weight', 'k_layernorm.weight']):
                     no_wd = False
                 else:
                     # Do not regularize biases and norm parameters.