[bugfix] fix multi-LoRA load assertion failure on idle adapter slots

[Yunnglin]

Problem

In multi-LoRA mode (e.g. 5 adapter slots: lora_0 - lora_4), save() only persists the active adapter's weights (e.g. lora_0). When load() is called to reload a checkpoint, the missing_keys assertion in GPTBridge._set_module() incorrectly checks ALL LoRA-related keys, causing an AssertionError because idle slots (lora_1 - lora_4) are naturally not in the checkpoint.

AssertionError: incompatible_keys.missing_keys: ['blocks.0.attn.qkv.lora_A.lora_1.weight', ...]

Root Cause

The filter at line 273-276 only checks if a key contains .lora_A. / .lora_B. / .modules_to_save., but does not distinguish between the adapter being loaded vs other idle slots:

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
]

Fix

Add f'.{self._adapter_name}.' in k to only validate keys belonging to the current adapter slot being loaded:

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)
    and f'.{self._adapter_name}.' in k
]

This is safe because:

• Each adapter is saved/loaded independently
• self._adapter_name identifies the target slot (e.g. lora_0)
• If the checkpoint is incomplete for the target adapter, the assert still fires correctly
• Idle slots being absent from checkpoint is expected behavior, not an error

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

Affected Versions

Tested on mcore-bridge 1.1.0 ~ 1.5.1 (logic unchanged across versions).

[gemini-code-assist]

Code Review

This pull request modifies the checkpoint loading logic in gpt_bridge.py to only validate missing keys for the active adapter when in multi-LoRA mode. The reviewer identified a potential issue where the substring check for the adapter name could cause false positives (e.g., matching layer indices if the adapter name is a number) and suggested a more robust check that explicitly targets the PEFT suffix structures.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[Copilot]

Pull request overview

Fixes a false AssertionError when reloading PEFT/LoRA checkpoints in multi-adapter (multi-slot) setups by narrowing the missing_keys assertion to only the adapter slot being loaded, rather than all idle adapter slots.

Changes:

• Update _set_module()’s PEFT missing_keys filtering to validate only keys for self._adapter_name.
• Add clarifying comments explaining expected missing keys for idle LoRA slots during load.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.