diff --git a/swift/megatron/trainers/base.py b/swift/megatron/trainers/base.py index a4f1b225af..da6a53129d 100644 --- a/swift/megatron/trainers/base.py +++ b/swift/megatron/trainers/base.py @@ -205,7 +205,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 diff --git a/swift/megatron/trainers/rlhf_mixin.py b/swift/megatron/trainers/rlhf_mixin.py index 5cbd83a36f..f38451fad5 100644 --- a/swift/megatron/trainers/rlhf_mixin.py +++ b/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'): diff --git a/swift/megatron/utils/megatron_lm_utils.py b/swift/megatron/utils/megatron_lm_utils.py index 476b4eef2e..8d641bb02c 100644 --- a/swift/megatron/utils/megatron_lm_utils.py +++ b/swift/megatron/utils/megatron_lm_utils.py @@ -193,7 +193,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'. @@ -214,7 +215,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: