From fcfb24ab3459dc05d11d2d4141bd764afd358665 Mon Sep 17 00:00:00 2001 From: Yunlin Mao Date: Thu, 25 Jun 2026 11:57:06 +0800 Subject: [PATCH] [bugfix] fix multi-LoRA load assertion failure on idle adapter slots In multi-LoRA mode (e.g. 5 slots: lora_0~lora_4), save() only persists the active adapter (lora_0). When load() is called, the missing_keys check incorrectly asserts that ALL LoRA slots must be present in the checkpoint, causing an AssertionError for idle slots (lora_1~4). Fix: narrow the missing_keys filter to only validate keys belonging to the current adapter slot (self._adapter_name), since other slots being absent from the checkpoint is expected behavior. Reproduction: 1. Start Megatron backend with multi-LoRA config (5 slots) 2. Train and call model.save(name) 3. Call model.load(name) -> AssertionError on lora_1~4 missing keys --- src/mcore_bridge/bridge/gpt_bridge.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/mcore_bridge/bridge/gpt_bridge.py b/src/mcore_bridge/bridge/gpt_bridge.py index 1b90308..e48f0bd 100644 --- a/src/mcore_bridge/bridge/gpt_bridge.py +++ b/src/mcore_bridge/bridge/gpt_bridge.py @@ -270,9 +270,12 @@ def _set_module(self, mg_module, hf_state_dict, hf_prefix: str, to_mcore: bool): incompatible_keys = mg_module.load_state_dict(hf_state_dict, strict=False) missing_keys = incompatible_keys.missing_keys if self._peft_format: + # In multi-LoRA mode, only the current adapter slot is saved/loaded. + # Other idle slots (e.g. lora_1~N) are not in the checkpoint and + # their absence is expected — only validate the active adapter. missing_keys = [ - k for k in incompatible_keys.missing_keys - if '.lora_A.' in k or '.lora_B.' in k or '.modules_to_save.' in k + k for k in incompatible_keys.missing_keys if + ('.lora_A.' in k or '.lora_B.' in k or '.modules_to_save.' in k) and f'.{self._adapter_name}.' in k ] assert len(missing_keys) == 0, f'incompatible_keys.missing_keys: {missing_keys}' return {}