[megatron] support Qwen3-Next apply_wd_to_qk_layernorm (modelscope#7954)

Author: Jintao-Huang

docs/source/Megatron-SWIFT/Command-line-parameters.md

@@ -301,6 +301,7 @@ Megatron训练参数继承自Megatron参数和基本参数（**与ms-swift共用
   - 提示：在日志中打印的"learning rate"为llm的学习率。
 - aligner_lr: 当训练多模态大模型时，该参数指定aligner的学习率，默认为None，等于learning_rate。
 - gradient_checkpointing_kwargs: 传入`torch.utils.checkpoint`中的参数。例如设置为`--gradient_checkpointing_kwargs '{"use_reentrant": false}'`。默认为None。该参数只对`vit_gradient_checkpointing`生效。
+- apply_wd_to_qk_layernorm: 用于Qwen3-Next全参数训练，对 qk layernorm 应用权重衰减。默认为False。
 - 🔥packing: 使用`padding_free`的方式将不同长度的数据样本打包成**近似**统一长度的样本（packing能保证不对完整的序列进行切分），实现训练时各节点与进程的负载均衡（避免长文本拖慢短文本的训练速度），从而提高GPU利用率，保持显存占用稳定。当使用 `--attention_backend flash` 时，可确保packed样本内的不同序列之间相互独立，互不可见（除Qwen3-Next，因为含有linear-attention）。该参数默认为`False`。Megatron-SWIFT的所有训练任务都支持该参数。注意：**packing会导致数据集样本数减少，请自行调节梯度累加数和学习率**。
 - packing_length: packing的长度。默认为None，设置为max_length。
 - packing_num_proc: packing的进程数，默认为1。需要注意的是，不同的`packing_num_proc`，最终形成的packed数据集是不同的。（该参数在流式packing时不生效）。通常不需要修改该值，packing速度远快于tokenize速度。

docs/source_en/Megatron-SWIFT/Command-line-parameters.md

@@ -320,6 +320,7 @@ Megatron training parameters are inherited from Megatron parameters and basic pa
   - Note: The "learning rate" printed in the logs is the learning rate of the LLM.
 - aligner_lr: Specifies the learning rate for the aligner module in multimodal models. Default is `None`, same as `learning_rate`.
 - gradient_checkpointing_kwargs: Arguments passed to `torch.utils.checkpoint`. For example: set `--gradient_checkpointing_kwargs '{"use_reentrant": false}'`. Defaults to `None`. This parameter only takes effect when `vit_gradient_checkpointing` is enabled.
+- apply_wd_to_qk_layernorm: Used for Qwen3-Next full-parameter training to apply weight decay to qk layernorm. Defaults to False.
 - 🔥packing: Use the `padding_free` method to pack data samples of different lengths into samples of **approximately** uniform length (packing ensures that complete sequences are not split), achieving load balancing across nodes and processes during training (preventing long texts from slowing down short text training), thereby improving GPU utilization and maintaining stable memory usage. When using `--attention_backend flash`, it ensures that different sequences within packed samples remain independent and invisible to each other (except for Qwen3-Next, which contains linear-attention). This parameter defaults to `False`. All training tasks in Megatron-SWIFT support this parameter. Note: **packing will reduce the number of dataset samples, please adjust gradient accumulation steps and learning rate accordingly**.
 - packing_length: the length to use for packing. Defaults to None, in which case it is set to max_length.
 - packing_num_proc: Number of processes for packing, default is 1. Note that different values of `packing_num_proc` will result in different packed datasets. (This parameter does not take effect during streaming packing). Usually there is no need to modify this value, as packing speed is much faster than tokenization speed.

swift/megatron/arguments/megatron_args.py

@@ -379,6 +379,7 @@ class ExtraMegatronArguments(RLHFMegatronArgumentsMixin, MegatronTunerMixin):
     linear_value_head_dim: Optional[int] = None
     linear_conv_kernel_dim: Optional[int] = None
     layer_types: Optional[List[str]] = None
+    apply_wd_to_qk_layernorm: bool = False
     # qwen3_vl, qwen3_omni
     mrope_interleaved: Optional[bool] = None
 
@@ -730,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 self.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

@@ -6,7 +6,7 @@
 import shutil
 import time
 from abc import ABC, abstractmethod
-from contextlib import contextmanager
+from contextlib import contextmanager, nullcontext
 from datetime import datetime
 from functools import partial
 from typing import Callable, Dict, List, Literal, Optional
@@ -317,7 +317,9 @@ def _get_param_groups(
         Returns:
             List of parameter groups.
         """
+        args = get_args()
         if self.args.vit_lr is not None or self.args.aligner_lr is not None:
+            assert self.args.megatron_model_meta.is_multimodal
             vit_lr = self.args.vit_lr if self.args.vit_lr is not None else self.args.lr
             aligner_lr = self.args.aligner_lr if self.args.aligner_lr is not None else self.args.lr
             logger.info(f'vit_lr: {vit_lr}, aligner_lr: {aligner_lr}, llm_lr: {self.args.lr}')
@@ -335,6 +337,9 @@ 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']):
+                    no_wd = False
                 else:
                     # Do not regularize biases and norm parameters.
                     #  optionally, also skip weight decay for embedding parameters if requested
@@ -423,10 +428,6 @@ def _get_param_groups(
 
     @contextmanager
     def _patch_get_param_groups(self):
-        if not self.args.megatron_model_meta.is_multimodal or (self.args.vit_lr is None
-                                                               and self.args.aligner_lr is None):
-            yield
-            return
         from megatron.core import optimizer
 
         _get_param_groups = optimizer._get_param_groups
@@ -497,7 +498,12 @@ def new_model_provider_func(*_args, **kwargs):
         # read iteration
         if not args.finetune:
             args.iteration, args.num_floating_point_operations_so_far = self._load_iteration()
-        with self._patch_load_state_dict(self._load_base_checkpoint), self._patch_get_param_groups():
+
+        if args.apply_wd_to_qk_layernorm or self.args.vit_lr is not None or self.args.aligner_lr is not None:
+            param_groups_context = self._patch_get_param_groups()
+        else:
+            param_groups_context = nullcontext()
+        with self._patch_load_state_dict(self._load_base_checkpoint), param_groups_context:
             model, optimizer, opt_param_scheduler = self._origin_setup_model_and_optimizer(
                 new_model_provider_func, model_type, *_args, **kwargs)
         self.wrapped_models = model