[megatron] support lora modules_to_save (#4916)

Author: Jintao-Huang

examples/train/megatron/lora/dpo.sh

@@ -1,4 +1,4 @@
-# 2 * 55GiB; 4.50s/it
+# 2 * 60GiB; 4.50s/it
 PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
 NPROC_PER_NODE=2 \
 CUDA_VISIBLE_DEVICES=0,1 \
@@ -10,6 +10,7 @@ megatron rlhf \
     --lora_rank 8 \
     --lora_alpha 32 \
     --target_modules all-linear \
+    --modules_to_save word_embeddings output_layer \
     --split_dataset_ratio 0.01 \
     --expert_model_parallel_size 2 \
     --moe_grouped_gemm true \
@@ -29,7 +30,7 @@ megatron rlhf \
     --save megatron_output/Qwen3-30B-A3B-Base \
     --eval_interval 100 \
     --save_interval 100 \
-    --max_length 8192 \
+    --max_length 2048 \
     --num_workers 8 \
     --dataset_num_proc 8 \
     --no_save_optim true \

examples/train/megatron/lora/loss_scale.sh

@@ -0,0 +1,40 @@
+# 2 * 60GiB, 3.4s/it
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+NPROC_PER_NODE=2 \
+CUDA_VISIBLE_DEVICES=0,1 \
+megatron sft \
+    --load Qwen3-30B-A3B-Base-mcore \
+    --train_type lora \
+    --dataset AI-ModelScope/function-calling-chatml#10000 \
+    --loss_scale hermes \
+    --agent_template hermes \
+    --lora_rank 8 \
+    --lora_alpha 32 \
+    --target_modules all-linear \
+    --modules_to_save word_embeddings output_layer \
+    --split_dataset_ratio 0.01 \
+    --expert_model_parallel_size 2 \
+    --moe_grouped_gemm true \
+    --moe_shared_expert_overlap true \
+    --moe_aux_loss_coeff 0.01 \
+    --micro_batch_size 8 \
+    --global_batch_size 16 \
+    --recompute_granularity full \
+    --recompute_method uniform \
+    --recompute_num_layers 1 \
+    --max_epochs 1 \
+    --finetune true \
+    --cross_entropy_loss_fusion true \
+    --lr 1e-4 \
+    --lr_warmup_fraction 0.05 \
+    --min_lr 1e-5 \
+    --save megatron_output/Qwen3-30B-A3B-Base \
+    --eval_interval 200 \
+    --save_interval 200 \
+    --max_length 2048 \
+    --num_workers 8 \
+    --dataset_num_proc 8 \
+    --no_save_optim true \
+    --no_save_rng true \
+    --sequence_parallel true \
+    --attention_backend flash

swift/megatron/init.py

@@ -639,6 +639,49 @@ def sharded_state_dict(
     TEGroupedLinear.sharded_state_dict = sharded_state_dict
 
 
+def _patch_peft_ModulesToSaveWrapper():
+    from peft.tuners import tuners_utils
+    from megatron.core.dist_checkpointing.mapping import ShardedStateDict
+    from .utils import tuners_sharded_state_dict
+
+    ModulesToSaveWrapper = tuners_utils.ModulesToSaveWrapper
+
+    class NewModulesToSaveWrapper(ModulesToSaveWrapper):
+
+        def __init__(self, module_to_save, *args, **kwargs):
+            tp_group = getattr(module_to_save, 'tp_group', None)
+            if tp_group is not None:
+                module_to_save.tp_group = None
+            super().__init__(module_to_save, *args, **kwargs)
+            if tp_group is not None:
+                module_to_save.tp_group = tp_group
+                for module in self.modules_to_save.values():
+                    module.tp_group = tp_group
+
+        def sharded_state_dict(
+                self,
+                prefix: str = '',
+                sharded_offsets: Tuple[Tuple[int, int, int]] = (),
+                metadata: Optional[dict] = None,
+        ) -> ShardedStateDict:
+            sharded_state_dict = tuners_sharded_state_dict(self, prefix, sharded_offsets, metadata)
+            if prefix == 'output_layer.':
+                output_layer_extra_state_key = f'{prefix}modules_to_save.default._extra_state'
+
+                # Old GPT checkpoints only stored the output layer weight key. So we remove the
+                # _extra_state key but check that it doesn't contain any data anyway
+                output_extra_state = sharded_state_dict.pop(output_layer_extra_state_key, None)
+                assert not (output_extra_state and output_extra_state.data
+                            ), f'Expected output layer extra state to be empty, got: {output_extra_state}'
+                # fix error
+                if f'{prefix}modules_to_save.default.weight' in sharded_state_dict:
+                    sharded_state_dict[f'{prefix}weight'] = sharded_state_dict[
+                        f'{prefix}modules_to_save.default.weight']
+            return sharded_state_dict
+
+    tuners_utils.ModulesToSaveWrapper = NewModulesToSaveWrapper
+
+
 def _patch_megatron():
     _patch_transformer_engine()
     _patch__batched_p2p_ops()
@@ -647,7 +690,8 @@ def _patch_megatron():
     from swift.megatron import tuners  # patch lora
     try:
         _patch_peft_BaseTuner()
-        logger.info('Patch peft_BaseTuner successfully applied.')
+        _patch_peft_ModulesToSaveWrapper()
+        logger.info('Patch peft successfully applied.')
     except Exception:
         pass
     try:

swift/megatron/trainers/base.py

@@ -20,7 +20,7 @@
 from packaging import version
 
 from swift.utils import JsonlWriter, get_logger, is_master
-from ..utils import adapter_state_dict_context, prepare_mcore_model
+from ..utils import adapter_state_dict_context, copy_original_module_weight, prepare_mcore_model
 from .utils import get_swift_datasets_provider
 
 logger = get_logger()
@@ -124,12 +124,20 @@ def _load_base_checkpoint(*_args, **kwargs):
             state_dict_model = {}
             mapping = {}
             for k, v in sharded_state_dict['model'].items():
-                if 'lora_A' in k or 'lora_B' in k:
+                if 'lora_A' in k or 'lora_B' in k or 'original_module' in k:
                     continue
-                origin_k = k
-                k = k.replace('.base_layer', '')
-                mapping[k] = origin_k
-                v.key = v.key.replace('.base_layer', '')
+                # lora
+                if '.base_layer' in k:
+                    origin_k = k
+                    k = k.replace('.base_layer', '')
+                    mapping[k] = origin_k
+                    v.key = v.key.replace('.base_layer', '')
+                elif '.modules_to_save' in k:
+                    # modules to save
+                    origin_k = k
+                    k = k.replace('.modules_to_save.default', '')
+                    mapping[k] = origin_k
+                    v.key = v.key.replace('.modules_to_save.default', '')
                 state_dict_model[k] = v
             sharded_state_dict['model'] = state_dict_model
             res = origin__load_base_checkpoint(*_args, **kwargs)
@@ -168,6 +176,8 @@ def new_model_provider_func(*args, **kwargs):
         if args.adapter_load is not None:
             with adapter_state_dict_context():
                 load_checkpoint(model, optimizer, opt_param_scheduler, load_arg='adapter_load', strict=False)
+        if args.train_type != 'full' and args.modules_to_save:
+            copy_original_module_weight(self.unwrapped_model)
         return model, optimizer, opt_param_scheduler
 
     def train_step(self, forward_step_func, data_iterator, model, optimizer, opt_param_scheduler, config):

swift/megatron/tuners/lora.py

@@ -14,13 +14,14 @@
 from megatron.core.models.common.embeddings.language_model_embedding import LanguageModelEmbedding
 from megatron.core.transformer.mlp import apply_swiglu_sharded_factory
 from megatron.core.transformer.module import MegatronModule
-from megatron.core.transformer.utils import make_sharded_tensors_for_checkpoint, sharded_state_dict_default
 from packaging import version
 from peft.tuners.lora import model
 from peft.tuners.lora.layer import LoraLayer
 from peft.tuners.tuners_utils import BaseTunerLayer, check_adapters_to_merge
 from peft.utils.other import transpose
 
+from ..utils import tuners_sharded_state_dict
+
 
 class LoraParallelLinear(MegatronModule, LoraLayer):
 
@@ -271,21 +272,7 @@ def sharded_state_dict(
             sharded_offsets: Tuple[Tuple[int, int, int]] = (),
             metadata: Optional[dict] = None,
     ) -> ShardedStateDict:
-        sharded_state_dict = {}
-        # Save parameters
-        self._save_to_state_dict(sharded_state_dict, '', keep_vars=True)
-        sharded_state_dict = make_sharded_tensors_for_checkpoint(
-            sharded_state_dict, prefix, sharded_offsets=sharded_offsets)
-        # Recurse into submodules
-        for name, module in self.named_children():
-            if 'Dict' in module.__class__.__name__:
-                modules = module.named_children()
-            else:
-                modules = [(None, module)]
-            for n, m in modules:
-                _prefix = f'{prefix}{name}.' if n is None else f'{prefix}{name}.{n}.'
-                sharded_state_dict.update(sharded_state_dict_default(m, _prefix, sharded_offsets, metadata))
-
+        sharded_state_dict = tuners_sharded_state_dict(self, prefix, sharded_offsets, metadata)
         if prefix.endswith('linear_fc1.'):
             if isinstance(self.base_layer, TEGroupedLinear) and self.config.gated_linear_unit:
                 num_global_experts = (parallel_state.get_expert_model_parallel_world_size() * self.base_layer.num_gemms)

swift/megatron/utils/__init__.py

@@ -2,4 +2,5 @@
 
 from .convert import convert_hf2mcore, convert_mcore2hf
 from .patcher import patch_megatron_tokenizer
-from .utils import adapter_state_dict_context, prepare_mcore_model
+from .utils import (adapter_state_dict_context, copy_original_module_weight, prepare_mcore_model,
+                    tuners_sharded_state_dict)

swift/megatron/utils/convert.py

@@ -221,5 +221,4 @@ def convert_mcore2hf(args: ExportArguments) -> None:
         model_dirs=[ckpt_dir, args.model_dir],
         max_shard_size=args.max_shard_size,
         additional_saved_files=hf_model.model_meta.additional_saved_files)
-    args.save_args()
     logger.info(f'Successfully saved HF model weights in `{args.output_dir}`.')

swift/megatron/utils/utils.py

@@ -1,11 +1,13 @@
 # Copyright (c) Alibaba, Inc. and its affiliates.
 from contextlib import contextmanager
+from typing import Optional, Tuple
 
 import torch.distributed as dist
 from megatron.core import mpu
 from megatron.core.extensions.transformer_engine import TEGroupedLinear, TELayerNormColumnParallelLinear, TELinear
 from megatron.core.models.common.embeddings.language_model_embedding import LanguageModelEmbedding
-from megatron.training import get_args
+from megatron.core.transformer.utils import make_sharded_tensors_for_checkpoint, sharded_state_dict_default
+from megatron.training import checkpointing, get_args
 
 from swift.utils import activate_parameters, find_layers, freeze_parameters, get_logger, get_model_parameter_info
 
@@ -96,7 +98,6 @@ def adapter_state_dict_context():
     if args.train_type == 'full':
         yield
         return
-    from megatron.training import checkpointing
     _origin_generate_state_dict = checkpointing.generate_state_dict
 
     def generate_state_dict(args, model, *_args, **kwargs):
@@ -121,3 +122,35 @@ def generate_state_dict(args, model, *_args, **kwargs):
         yield
     finally:
         checkpointing.generate_state_dict = _origin_generate_state_dict
+
+
+def tuners_sharded_state_dict(
+        module,
+        prefix: str = '',
+        sharded_offsets: Tuple[Tuple[int, int, int]] = (),
+        metadata: Optional[dict] = None,
+):
+    sharded_state_dict = {}
+    # Save parameters
+    module._save_to_state_dict(sharded_state_dict, '', keep_vars=True)
+    sharded_state_dict = make_sharded_tensors_for_checkpoint(
+        sharded_state_dict, prefix, sharded_offsets=sharded_offsets)
+    # Recurse into submodules
+    for name, module in module.named_children():
+        if 'Dict' in module.__class__.__name__:
+            modules = module.named_children()
+        else:
+            modules = [(None, module)]
+        for n, m in modules:
+            _prefix = f'{prefix}{name}.' if n is None else f'{prefix}{name}.{n}.'
+            sharded_state_dict.update(sharded_state_dict_default(m, _prefix, sharded_offsets, metadata))
+    return sharded_state_dict
+
+
+def copy_original_module_weight(model):
+    for module in model.modules():
+        if 'ModulesToSaveWrapper' in module.__class__.__name__ and hasattr(module, 'original_module'):
+            original_module = module.original_module
+            modules_to_save = module.modules_to_save
+            if 'default' in modules_to_save:
+                original_module.load_state_dict(modules_to_save['default'].state_dict())

tests/megatron/test_lora.py

@@ -12,14 +12,14 @@ def test_sft():
             loss_scale='hermes',
             split_dataset_ratio=0.01,
             tensor_model_parallel_size=2,
-            load_from_cache_file=False,
             train_type='lora',
             recompute_granularity='full',
             recompute_method='uniform',
             recompute_num_layers=1,
             # pipeline_model_parallel_size=2,
             # freeze_parameters_ratio=0.5,
             train_iters=100,
+            modules_to_save=['word_embeddings', 'output_layer'],
             eval_iters=5,
             save_interval=5,
             no_save_optim=True,
@@ -41,6 +41,7 @@ def test_moe():
             # expert_model_parallel_size=2,
             train_type='lora',
             recompute_granularity='full',
+            modules_to_save=['word_embeddings', 'output_layer'],
             recompute_method='uniform',
             recompute_num_layers=1,
             # pipeline_model_parallel_size=2,
@@ -67,15 +68,11 @@ def test_embedding():
     pass
 
 
-def test_modules_to_save():
-    pass
-
-
 def test_resume():
     pass
 
 
 if __name__ == '__main__':
-    # test_sft()
+    test_sft()
     # test_moe()
-    test_convert()
+    # test_convert()