Fix/lora save bloat (#293)

CHANGELOG.md

@@ -5,6 +5,15 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.13.3] - 2025-12-06
+
+### 🐛 Bug Fixes
+
+- **LoRA save bloat prevention**: Bake and strip LoRA wrappers before sharding to avoid exploding shard counts/sizes when saving quantized models with multiple/mismatched LoRAs (see [issue #217 comment](https://github.com/filipstrand/mflux/issues/217#issuecomment-3615321206)).
+- **Regression test hardening**: LoRA model-saving tests now include size guardrails (5% tolerance) while using the bundled local LoRA fixtures to catch shard bloat regressions early.
+
+---
+
 ## [0.13.2] - 2025-12-05
 
 ### ✨ Improvements

pyproject.toml

@@ -14,7 +14,7 @@ source-exclude = [
 
 [project]
 name = "mflux"
-version = "0.13.2"
+version = "0.13.3"
 description = "A MLX port of FLUX based on the Huggingface Diffusers implementation."
 readme = "README.md"
 keywords = ["diffusers", "flux", "mlx"]

src/mflux/models/common/lora/mapping/lora_saver.py

@@ -0,0 +1,82 @@
+import mlx.core as mx
+import mlx.nn as nn
+
+from mflux.models.common.lora.layer.fused_linear_lora_layer import FusedLoRALinear
+from mflux.models.common.lora.layer.linear_lora_layer import LoRALinear
+
+
+class LoRASaver:
+    @staticmethod
+    def bake_and_strip_lora(module: nn.Module) -> nn.Module:
+        def _assign(parent, attr_name, idx, new_child):
+            if parent is None:
+                return
+            if isinstance(parent, list) and idx is not None:
+                parent[idx] = new_child
+            elif isinstance(parent, dict) and attr_name is not None:
+                parent[attr_name] = new_child
+            elif attr_name is not None:
+                setattr(parent, attr_name, new_child)
+
+        def _bake_single(lora_layer: LoRALinear) -> nn.Module:
+            base_linear = lora_layer.linear
+            LoRASaver._apply_lora_delta(base_linear, lora_layer)
+            return base_linear
+
+        def _bake_fused(fused_layer: FusedLoRALinear) -> nn.Module:
+            base_linear = fused_layer.base_linear
+            for lora in fused_layer.loras:
+                if isinstance(lora, LoRALinear):
+                    LoRASaver._apply_lora_delta(base_linear, lora)
+            return base_linear
+
+        def _walk(obj, parent=None, attr_name=None, idx=None):
+            # Replace wrappers first
+            if isinstance(obj, FusedLoRALinear):
+                new_child = _bake_fused(obj)
+                _assign(parent, attr_name, idx, new_child)
+                obj = new_child
+            elif isinstance(obj, LoRALinear):
+                new_child = _bake_single(obj)
+                _assign(parent, attr_name, idx, new_child)
+                obj = new_child
+
+            # Recurse into containers/modules
+            if isinstance(obj, list):
+                for i, child in enumerate(list(obj)):
+                    _walk(child, obj, None, i)
+            elif isinstance(obj, tuple):
+                temp_list = list(obj)
+                for i, child in enumerate(temp_list):
+                    _walk(child, temp_list, None, i)
+                if parent is not None:
+                    _assign(parent, attr_name, idx, type(obj)(temp_list))
+            elif isinstance(obj, dict):
+                for key, child in list(obj.items()):
+                    _walk(child, obj, key, None)
+            elif isinstance(obj, nn.Module):
+                for name, child in vars(obj).items():
+                    if isinstance(child, (nn.Module, list, tuple, dict)):
+                        _walk(child, obj, name, None)
+
+        _walk(module, None, None, None)
+        return module
+
+    @staticmethod
+    def _apply_lora_delta(base_linear: nn.Module, lora_layer: LoRALinear) -> None:
+        if not hasattr(base_linear, "weight"):
+            return
+
+        weight = base_linear.weight
+        delta = mx.matmul(lora_layer.lora_A, lora_layer.lora_B)  # shape: [in, out]
+        delta = mx.transpose(delta)  # shape: [out, in]
+        delta = lora_layer.scale * delta
+
+        if weight.shape != delta.shape:
+            print(f"⚠️  Skipping LoRA bake due to shape mismatch: weight {weight.shape} vs delta {delta.shape}")
+            return
+
+        try:
+            base_linear.weight = weight + delta.astype(weight.dtype)
+        except Exception as e:  # noqa: BLE001
+            print(f"⚠️  Failed to bake LoRA into base layer: {e}")

src/mflux/models/common/weights/saving/model_saver.py

@@ -8,6 +8,7 @@ from mlx.utils import tree_flatten
 from tqdm import tqdm
 from transformers import PreTrainedTokenizer
 
+from mflux.models.common.lora.mapping.lora_saver import LoRASaver
 from mflux.utils.version_util import VersionUtil
 
 if TYPE_CHECKING:
@@ -35,6 +36,8 @@ class ModelSaver:
         for attr_name, subdir in tqdm(components, desc="Saving components", unit="component"):
             component = getattr(model, attr_name, None)
             if component is not None:
+                # Bake and strip any LoRA wrappers to avoid duplicating shared weights
+                LoRASaver.bake_and_strip_lora(component)
                 ModelSaver._save_weights(base_path, bits, component, subdir)
 
     @staticmethod

tests/model_saving/test_model_saving_lora.py

@@ -9,15 +9,25 @@ from mflux.models.common.config import ModelConfig
 from mflux.models.flux.variants.txt2img.flux import Flux1
 
 PATH = "tests/4bit/"
+SIZE_TOLERANCE_RATIO = 0.05  # allow small metadata/header differences
 
 
 class TestModelSavingLora:
+    LORA_FILES = [
+        "FLUX-dev-lora-MiaoKa-Yarn-World.safetensors",
+        "Flux_-_Renaissance_art_style.safetensors",
+    ]
+    LORA_SCALES = [0.6, 0.4]
+
     @pytest.mark.slow
     def test_save_and_load_4bit_model_with_lora(self):
         # Clean up any existing temporary directories from previous test runs
-        TestModelSavingLora.delete_folder_if_exists(PATH)
+        TestModelSavingLora._delete_folder_if_exists(PATH)
 
         try:
+            lora_paths = TestModelSavingLora._get_lora_paths()
+            assert len(lora_paths) == len(TestModelSavingLora.LORA_SCALES)
+
             # given a saved quantized model on disk (without LoRA)...
             fluxA = Flux1(
                 model_config=ModelConfig.schnell(),
@@ -30,8 +40,8 @@ class TestModelSavingLora:
             fluxB = Flux1(
                 model_config=ModelConfig.schnell(),
                 quantize=4,
-                lora_paths=TestModelSavingLora.get_lora_path(),
-                lora_scales=[1.0],
+                lora_paths=lora_paths,
+                lora_scales=TestModelSavingLora.LORA_SCALES,
             )
             image1 = fluxB.generate_image(
                 seed=44,
@@ -46,8 +56,8 @@ class TestModelSavingLora:
             fluxC = Flux1(
                 model_config=ModelConfig.schnell(),
                 model_path=PATH,
-                lora_paths=TestModelSavingLora.get_lora_path(),
-                lora_scales=[1.0],
+                lora_paths=lora_paths,
+                lora_scales=TestModelSavingLora.LORA_SCALES,
             )
 
             # ...and generating the identical image
@@ -68,14 +78,51 @@ class TestModelSavingLora:
 
         finally:
             # cleanup
-            TestModelSavingLora.delete_folder(PATH)
+            TestModelSavingLora._delete_folder_if_exists(PATH)
+
+    @pytest.mark.slow
+    def test_save_with_lora_has_same_shard_count_as_base(self):
+        base_path = "tests/4bit_base/"
+        lora_path = "tests/4bit_with_lora/"
+
+        TestModelSavingLora._delete_folder_if_exists(base_path)
+        TestModelSavingLora._delete_folder_if_exists(lora_path)
+
+        try:
+            flux_base = Flux1(
+                model_config=ModelConfig.schnell(),
+                quantize=4,
+            )
+            flux_base.save_model(base_path)
+            del flux_base
+
+            flux_lora = Flux1(
+                model_config=ModelConfig.schnell(),
+                quantize=4,
+                lora_paths=TestModelSavingLora._get_lora_paths(),
+                lora_scales=TestModelSavingLora.LORA_SCALES,
+            )
+            flux_lora.save_model(lora_path)
+            del flux_lora
+
+            base_shards = list((Path(base_path) / "transformer").glob("*.safetensors"))
+            lora_shards = list((Path(lora_path) / "transformer").glob("*.safetensors"))
+
+            assert len(base_shards) == len(lora_shards), "LoRA save should not inflate transformer shard count"
+
+            # Also assert the total saved size stays within a tight bound to catch shard bloat
+            base_size = TestModelSavingLora._dir_size_bytes(base_path)
+            lora_size = TestModelSavingLora._dir_size_bytes(lora_path)
+
+            # LoRA baking should leave sizes effectively unchanged; allow a small tolerance for metadata
+            max_allowed = base_size * (1 + SIZE_TOLERANCE_RATIO)
+            assert lora_size <= max_allowed, f"LoRA save size grew unexpectedly: base={base_size}B vs lora={lora_size}B"
+        finally:
+            TestModelSavingLora._delete_folder_if_exists(base_path)
+            TestModelSavingLora._delete_folder_if_exists(lora_path)
 
     @staticmethod
-    def delete_folder(path: str) -> None:
-        return shutil.rmtree(path)
-
-    @staticmethod
-    def delete_folder_if_exists(path: str) -> None:
+    def _delete_folder_if_exists(path: str) -> None:
         if os.path.exists(path):
             shutil.rmtree(path)
             print(f"Deleted folder: {path}")
@@ -83,12 +130,23 @@ class TestModelSavingLora:
             print(f"Folder does not exist: {path}")
 
     @staticmethod
-    def resolve_path(path) -> Path | None:
+    def _resolve_path(path) -> Path | None:
         if path is None:
             return None
         return Path(__file__).parent.parent / "resources" / path
 
     @staticmethod
-    def get_lora_path() -> list[str]:
-        path = TestModelSavingLora.resolve_path("FLUX-dev-lora-MiaoKa-Yarn-World.safetensors")
-        return [str(path)]
+    def _get_lora_paths() -> list[str]:
+        resolved_paths = [TestModelSavingLora._resolve_path(path) for path in TestModelSavingLora.LORA_FILES]
+        missing = [p for p in resolved_paths if p is None or not p.exists()]
+        if missing:
+            missing_names = ", ".join(sorted(p.name if p else "unknown" for p in missing))
+            pytest.skip(f"Missing local LoRA test asset(s): {missing_names}")
+        return [str(path) for path in resolved_paths if path is not None]
+
+    @staticmethod
+    def _dir_size_bytes(path: str | Path) -> int:
+        p = Path(path)
+        if not p.exists():
+            return 0
+        return sum(f.stat().st_size for f in p.rglob("*") if f.is_file())