Repository: huggingface/transformers
Pull request: [WIP] FIX Make Mixtral LoRA loading work

Required fixes:

- some code was using unordered data structures, making weight order random
- adjust alpha to offset increased rank from fusion
- import functions from PEFT if available

See https://github.com/huggingface/peft/pull/3083.


Changed files:
filename: src/transformers/conversion_mapping.py
patch: @@ -464,6 +464,7 @@ def get_model_conversion_mapping(
     For a given `model`, obtain the weight conversion mapping if any are registered either as a simple renaming
     `_checkpoint_conversion_mapping` class argument, or in the general WeightConverter mapping.
     """
+    # note: this function is used in PEFT, so changing the API requires coordination
     weight_conversions = []
 
     # Load models with explicit, user-provided key mapping

filename: src/transformers/core_model_loading.py
patch: @@ -646,7 +646,7 @@ def materialize_tensors(self) -> dict[str, list[torch.Tensor]]:
         - saving: the tensors are already torch.Tensor instances (the existing model weights)
         """
         collected_tensors = {}
-        for key in set(self.collected_tensors.keys()):
+        for key in list(self.collected_tensors.keys()):
             # Remove from internal attribute
             tensors = self.collected_tensors.pop(key)
             # Async loading

filename: src/transformers/integrations/peft.py
patch: @@ -64,14 +64,25 @@
     from ..modeling_utils import LoadStateDictConfig
 
 
-def _block_diag_3d(*tensors):
+# TODO: remove once PEFT < 0.19 no longer supported
+def _block_diag_3d(tensors: list[torch.Tensor]) -> torch.Tensor:
+    if len(tensors) < 2:
+        raise ValueError(f"_block_diag_3d expects at least 2 tensors, got {len(tensors)}")
+
+    if any(t.dim() != 3 for t in tensors):
+        raise ValueError("_block_diag_3d expects all tensors to be 3d.")
+
+    num_experts = tensors[0].shape[0]
+    if any(t.shape[0] != num_experts for t in tensors):
+        raise ValueError("All tensors passed to _block_diag_3d must have the same number of experts.")
+
     lora_b_block_diag = []
-    for i in range(len(tensors[0])):
-        lora_b_block_diag.append(torch.block_diag(tensors[0][i], tensors[1][i]))
-    out = torch.stack(lora_b_block_diag, dim=0)
-    return out
+    for i in range(num_experts):
+        lora_b_block_diag.append(torch.block_diag(*[tensor[i] for tensor in tensors]))
+    return torch.stack(lora_b_block_diag, dim=0)
 
 
+# TODO: remove once PEFT < 0.19 no longer supported
 class PeftConcatenate(Concatenate):
     """Convert per-expert LoRA weights to merged weights.
 
@@ -81,16 +92,21 @@ class PeftConcatenate(Concatenate):
 
     To illustrate:
 
-    Before
+    Before:
+
     W0' = W0 + A0 @ B0
+
     W1' = W1 + A1 @ B1
 
-    After
+    After:
+
     W01' = W01 + A01 @ B01_bd
-        where
-        A01 = [A0, A1]
-        B01_bd = [[B0,  0],
-                  [0,  B1]]
+
+    where:
+
+    A01 = [A0, A1]
+
+    B01_bd = [[B0, 0], [0, B1]]
 
     This class is responsible for merging LoRA B in this block-diagonal fashion. Assuming that we fuse N weights, it
     should look like this:
@@ -119,12 +135,24 @@ def convert(
                 f"To convert this LoRA adapter, the LoRA weights all need to have either 2 or 3 dims, got {set(dims)}"
             )
 
+        # Keep source order stable (e.g. w1 before w3 for Mixtral) to preserve gate/up semantics.
+        ordered_tensors = [
+            input_dict[source_pattern] for source_pattern in source_patterns if source_pattern in input_dict
+        ]
+        if len(ordered_tensors) != len(input_dict):
+            missing = set(input_dict) - set(source_patterns)
+            raise ValueError(
+                "Collected tensors contain keys not present in source_patterns. "
+                f"Unexpected keys: {sorted(missing)}; source_patterns={source_patterns}"
+            )
+
         if set(dims) == {2}:
-            output_dict = {full_layer_name: torch.block_diag(*input_dict.values())}
+            output_dict = {full_layer_name: torch.block_diag(*ordered_tensors)}
         else:
-            out = _block_diag_3d(*input_dict.values())  # shape = experts, 2*out_feat, 2*r
-            out = torch.permute(out, (2, 0, 1))  # shape = 2*r, experts, 2*out_feat
-            out = out.flatten(0, 1)  # shape = 2*r * experts, 2*out_feat
+            # with r being the LoRA rank and n being the number of fused weights:
+            out = _block_diag_3d(ordered_tensors)  # shape = experts, n*out_feat, 2*r
+            out = torch.permute(out, (2, 0, 1))  # shape = 2*r, experts, n*out_feat
+            out = out.flatten(0, 1)  # shape = 2*r * experts, n*out_feat
             out = out.T
             output_dict = {full_layer_name: out}
         return output_dict
@@ -134,6 +162,7 @@ def reverse_op(self) -> ConversionOps:
         raise NotImplementedError("Reversing PEFT LoRA MoE conversions is not supported yet.")
 
 
+# TODO: remove once PEFT < 0.19 no longer supported
 class FlattenDims(ConversionOps):
     """
     Flatten the tensors along the given dimensions
@@ -164,6 +193,7 @@ def __repr__(self):
         return f"{self.__class__.__name__}(dims={self.dims})"
 
 
+# TODO: remove once PEFT < 0.19 no longer supported
 class PermuteDims(ConversionOps):
     """
     Permute the tensors along the given dimensions
@@ -192,11 +222,13 @@ def __repr__(self):
         return f"{self.__class__.__name__}(dims={self.dims})"
 
 
-def _build_peft_weight_mapping(
+# TODO: remove once PEFT < 0.19 no longer supported
+def build_peft_weight_mapping(
     weight_conversions: list[WeightConverter | WeightRenaming] | None, adapter_name: str, peft_config=None
 ) -> list[WeightConverter | WeightRenaming]:
     # We iterate over all the operations of the original model and simply edit them to apply to the PEFT adapter when
     # appropriate.
+    # Note: This function is used in PEFT, changing it requires coordination.
     if not weight_conversions:
         return []
 
@@ -311,6 +343,8 @@ def _build_peft_weight_mapping(
 # The main reason we have to explicit this is because the conversion mapping
 # has the full layer name, while the config do not. We coould regex match but
 # this is more explicit and less error prone.
+# Note: this is used in PEFT, changing it requires coordiation.
+# TODO: remove once PEFT < 0.19 no longer supported
 _MOE_TARGET_MODULE_MAPPING: dict[str, dict[str, str]] = {
     "mixtral": {
         "gate": "gate.weight",
@@ -326,12 +360,16 @@ def _build_peft_weight_mapping(
     },
 }
 
+# Note: this is used in PEFT, changing it requires coordiation.
+# TODO: remove once PEFT < 0.19 no longer supported
 _MOE_FUSED_TARGETS: dict[str, dict[str, set[str]]] = {
-    "mixtral": {"gate_up_proj": {"w1", "w3"}},
-    "qwen2_moe": {"gate_up_proj": {"gate_proj", "up_proj"}},
+    # use lists for dict values to ensure stable order
+    "mixtral": {"gate_up_proj": ["w1", "w3"]},
+    "qwen2_moe": {"gate_up_proj": ["gate_proj", "up_proj"]},
 }
 
 
+# TODO: remove once PEFT < 0.19 no longer supported
 def patch_moe_parameter_targeting(model, peft_config):
     """PEFT currently assumes that expert layers are of shape
         (expert, in, out)
@@ -349,11 +387,12 @@ def patch_moe_parameter_targeting(model, peft_config):
 
         @wraps(update_layer)
         def new_update_layer(layer, *args, **kwargs):
-            if not hasattr(layer, "_swapped_in_out") and layer.parameter_name in ("down_proj", "gate_up_proj"):
+            did_swap = getattr(layer, "_did_swap_in_out_features", False)
+            if not did_swap and layer.parameter_name in ("down_proj", "gate_up_proj"):
                 tmp_in_features = layer.in_features
                 layer.in_features = layer.out_features
                 layer.out_features = tmp_in_features
-                layer._swapped_in_out = True
+                layer._did_swap_in_out_features = True
             return update_layer(layer, *args, **kwargs)
 
         peft.tuners.lora.layer.ParamWrapper.update_layer = new_update_layer
@@ -520,12 +559,14 @@ def load_adapter(
             )
 
         weight_conversions = get_model_conversion_mapping(self)
+
+        # TODO: remove once PEFT < 0.19 is dropped, use peft.utils.transformers_weight_conversion
         peft_config = convert_peft_config_for_transformers(peft_config, model=self, conversions=weight_conversions)
 
         if hasattr(peft_config, "inference_mode"):
             peft_config.inference_mode = not is_trainable
 
-        peft_weight_conversions = _build_peft_weight_mapping(weight_conversions, adapter_name, peft_config=peft_config)
+        peft_weight_conversions = build_peft_weight_mapping(weight_conversions, adapter_name, peft_config=peft_config)
 
         patch_moe_parameter_targeting(model=self, peft_config=peft_config)
 
@@ -987,10 +1028,8 @@ def maybe_load_adapters(
 # PEFT adapters for these models, they also need to be updated. This may require updating the PEFT config too. The
 # logic for this is found below. Right now, only LoRA is supported.
 
-# TODO: These functions will be added to PEFT in release 0.19.0. Drop them here once 0.19.0 becomes the min PEFT
-# version.
-
 
+# TODO: remove once PEFT < 0.19 no longer supported
 def _convert_peft_config_moe(peft_config, model_type: str):
     base_model_type = _MODEL_TO_CONVERSION_PATTERN.get(model_type, None)
     if base_model_type is None:
@@ -1003,6 +1042,8 @@ def _convert_peft_config_moe(peft_config, model_type: str):
     peft_config.target_modules = set(peft_config.target_modules or [])
     if not hasattr(peft_config, "rank_pattern") or peft_config.rank_pattern is None:
         peft_config.rank_pattern = {}
+    if not hasattr(peft_config, "alpha_pattern") or peft_config.alpha_pattern is None:
+        peft_config.alpha_pattern = {}
 
     new_target_parameters = peft_config.target_parameters.copy()
     remaining_target_modules = set()
@@ -1036,15 +1077,25 @@ def _convert_peft_config_moe(peft_config, model_type: str):
 
         if len(present_targets) == len(required_old_targets) and len(required_old_targets) > 1:
             peft_config.rank_pattern[rf".*\.{re.escape(new_name)}"] = peft_config.r * len(required_old_targets)
+            # Preserve per-branch LoRA scaling after fusion.
+            # Example: w1 + w3 => r doubles, so alpha must also double to keep alpha/r unchanged.
+            peft_config.alpha_pattern[rf".*\.{re.escape(new_name)}"] = peft_config.lora_alpha * len(
+                required_old_targets
+            )
 
     peft_config.target_parameters = new_target_parameters
     peft_config.target_modules = remaining_target_modules
 
     return peft_config
 
 
+# TODO: remove once PEFT < 0.19 no longer supported
 def convert_peft_config_for_transformers(peft_config, model: torch.nn.Module, conversions: list[Any] | None):
-    # FIXME document this properly
+    """
+    Convert the PEFT config of models whose architecture changed from transformers v4 to v5.
+
+    For most models, this requires no changes, this mostly affects some MoE models like Mixtral.
+    """
     # If, for any reason, we cannot apply conversion, we just return the PEFT config as is.
     from peft import PeftType  # avoid circular import
 

filename: tests/peft_integration/test_peft_integration.py
patch: @@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import gc
+import importlib.metadata
 import json
 import os
 import re
@@ -21,6 +22,7 @@
 
 from datasets import Dataset, DatasetDict
 from huggingface_hub import hf_hub_download
+from packaging import version
 from torch import nn
 
 from transformers import (
@@ -930,6 +932,42 @@ def test_non_lora_add_adapter(self):
                 # should be different
                 assert not torch.allclose(output_base, output_peft, atol=atol, rtol=rtol)
 
+    def test_mixtral_lora_conversion(self):
+        if version.parse(importlib.metadata.version("peft")) < version.parse("0.19.0"):
+            self.skipTest("For this test to pass, PEFT 0.19 is required.")
+
+        inputs = torch.arange(10).view(1, -1).to(0)
+        model_name = "hf-internal-testing/Mixtral-tiny"
+        adapter_name = "peft-internal-testing/mixtral-pre-v5-lora"
+
+        # original logits were:
+        # tensor([[[ 0.2676,  0.3870,  0.2956,  ...,  0.4624,  0.1966,  0.2539],
+        #          [-0.6706, -0.0969, -0.6240,  ..., -0.0201,  0.7099, -0.3099],
+        #          [ 0.0663,  0.1653,  0.7189,  ...,  0.5905,  0.0649,  0.5839],
+        #          ...,
+        #          [-0.2712, -0.6451, -0.0219,  ..., -0.4344,  0.5471, -0.9355],
+        #          [-0.3607,  0.4526,  0.2750,  ...,  0.1082,  0.7179,  0.8487],
+        #          [ 0.5826, -0.1407, -0.3131,  ...,  0.1026,  0.6878, -0.3382]]],
+        #        device='cuda:0')
+        expected_logits_0_to_3 = torch.Tensor(
+            [
+                [0.2676, 0.3870, 0.2956],
+                [-0.6706, -0.0969, -0.6240],
+                [0.0663, 0.1653, 0.7189],
+            ]
+        ).to(device=torch_device, dtype=torch.float16)
+
+        model = AutoModelForCausalLM.from_pretrained(model_name)
+        model.load_adapter(adapter_name)
+        model.to(torch_device)
+        model.eval()
+        with torch.inference_mode():
+            output = model(inputs).logits
+
+        # a little bit of deviation but that's fine
+        atol, rtol = 1e-3, 1e-4
+        assert torch.allclose(output[0, :3, :3], expected_logits_0_to_3, atol=atol, rtol=rtol)
+
 
 @require_peft
 @require_torch