Fix Megatron reference adapter loading

Author: taking-lying-flat

swift/megatron/trainers/base.py

@@ -198,7 +198,7 @@ def _prepare_peft_model(self, models):
         if args.mcore_model is None:
             self.bridge.load_weights(models, args.model_dir)
         peft_models = [prepare_mcore_model(args, model) for model in models]
-        if args.tuner_type == 'lora' and args.adapters and args.mcore_adapter is None:
+        if args.tuner_type in {'lora', 'lora_llm'} and args.adapters and args.mcore_adapter is None:
             assert len(args.adapters) == 1, 'Currently only support one adapter.'
             self.bridge.load_weights(models, args.adapters[0], peft_format=True, adapter_name='default')
         return peft_models

swift/megatron/trainers/rlhf_mixin.py

@@ -21,7 +21,7 @@ def _load_checkpoint(self):
         if args.mcore_ref_model is not None:
             load_mcore_checkpoint(args, self.ref_models, load_arg='mcore_ref_model')
         if args.mcore_ref_adapter is not None:
-            load_mcore_checkpoint(args, self.wrapped_models, load_arg='mcore_ref_adapter')
+            load_mcore_checkpoint(args, self.wrapped_models, load_arg='mcore_ref_adapter', adapter_name='ref_adapter')
         super()._load_checkpoint()
 
     def prepare_model(self):
@@ -39,10 +39,10 @@ def prepare_model(self):
             ref_model_id_or_path = args.ref_model or args.model
             ref_model_dir = safe_snapshot_download(ref_model_id_or_path, use_hf=args.use_hf, hub_token=args.hub_token)
             self.bridge.load_weights(self.ref_models, ref_model_dir)
-        if args.tuner_type == 'lora' and args.ref_adapters and args.mcore_ref_adapter is None:
+        if args.tuner_type in {'lora', 'lora_llm'} and args.ref_adapters and args.mcore_ref_adapter is None:
             assert len(args.ref_adapters) == 1, 'Currently only support one adapter.'
             self.bridge.load_weights(
-                self.ref_models, args.ref_adapters[0], peft_format=True, adapter_name='ref_adapter')
+                self.unwrapped_models, args.ref_adapters[0], peft_format=True, adapter_name='ref_adapter')
 
     def _get_data_collator(self):
         if self.args.rlhf_type in ('grpo', 'gkd'):

swift/megatron/utils/megatron_lm_utils.py

@@ -170,7 +170,8 @@ def _load_optimizer_state_dict(optimizer, state_dict):
 
 def _filter_adapter_state_dict(state_dict, peft_format: bool, adapter_name: str = 'default'):
     """
-    When peft_format is True, keep only the PEFT format state_dict;
+    When peft_format is True, keep only the requested PEFT adapter and map its
+    checkpoint lookup keys from the default adapter slot;
     when False, remove the PEFT format state_dict.
 
     This function ensures it is called when tuner_type != 'full'.
@@ -191,7 +192,13 @@ def _filter_adapter_state_dict(state_dict, peft_format: bool, adapter_name: str
         state_dict_model = state_dict[model_key]
         for k, v in state_dict_model.items():
             if peft_format:
-                if '.lora_A.' in k or '.lora_B.' in k or '.modules_to_save.' in k:
+                adapter_modules = ('lora_A', 'lora_B', 'modules_to_save')
+                if any(f'.{module}.{adapter_name}.' in k for module in adapter_modules):
+                    if adapter_name != 'default':
+                        # Keep the state-dict key for the target adapter, but read the tensor from the
+                        # default adapter slot used by the source checkpoint.
+                        for module in adapter_modules:
+                            v.key = v.key.replace(f'.{module}.{adapter_name}.', f'.{module}.default.')
                     new_state_dict[k] = v
             else:
                 if '.lora_A.' in k or '.lora_B.' in k or 'original_module.' in k: