fix(model-loaders): stop materializing scaled-fp8 checkpoints in float32

invokeai/backend/model_manager/load/model_loaders/krea2.py

@@ -2,7 +2,7 @@
 """Class for Krea-2 model loading in InvokeAI."""
 
 from pathlib import Path
-from typing import Any, Optional
+from typing import TYPE_CHECKING, Any, Optional
 
 import accelerate
 from transformers import AutoConfig, AutoTokenizer
@@ -27,6 +27,10 @@
 from invokeai.backend.quantization.gguf.loaders import gguf_sd_loader
 from invokeai.backend.util.devices import TorchDevice
 
+if TYPE_CHECKING:
+    # torch is imported lazily inside the helpers below; this is annotations-only.
+    import torch
+
 
 def _normalize_qwen3vl_rope_config(config: Any) -> Any:
     """Mirror Qwen3-VL rope_parameters into rope_scaling for Transformers compatibility."""
@@ -89,12 +93,23 @@ def _is_native_krea2_format(sd: dict[str, Any]) -> bool:
     )
 
 
-def _dequantize_scaled_fp8(sd: dict[str, Any]) -> dict[str, Any]:
+def _dequantize_scaled_fp8(sd: dict[str, Any], dtype: "torch.dtype") -> dict[str, Any]:
     """Dequantize ComfyUI 'scaled fp8' weights: ``dequant = weight.float() * weight_scale``.
 
     Each quantized layer stores an fp8 ``<name>.weight`` plus a (usually scalar) ``<name>.weight_scale``.
-    Returns a new dict with the weights dequantized to float and the ``.weight_scale`` keys removed.
-    No-op if there are no scale keys.
+    Returns a new dict with the weights dequantized and the ``.weight_scale`` keys removed. No-op if
+    there are no scale keys.
+
+    The multiply runs in float32 for precision, but each result is stored as ``dtype`` immediately so
+    the *whole model* is never materialized in float32. Krea-2's ~12 GB fp8 checkpoint would otherwise
+    peak at ~50 GB of RAM (4 bytes/param) before the caller's later bf16 cast brings it down to ~25 GB,
+    which puts a 32 GB machine into swap during a cold load. This mirrors the same fix already applied
+    to the FLUX.2 loader.
+
+    ``dtype`` is required on purpose. It used to default to bfloat16, which is wrong on a device where
+    ``choose_bfloat16_safe_dtype`` picks float16: the weights would land in bf16 and then take a second
+    rounding step on the caller's later float16 cast. Callers already know the compute dtype, so there
+    is no reason to guess one here.
     """
     import torch
 
@@ -107,7 +122,8 @@ def _dequantize_scaled_fp8(sd: dict[str, Any]) -> dict[str, Any]:
         if weight_key in out:
             weight = torch.as_tensor(_to_plain_tensor(out[weight_key])).float()
             scale = torch.as_tensor(_to_plain_tensor(out[scale_key])).float()
-            out[weight_key] = weight * scale
+            out[weight_key] = (weight * scale).to(dtype)
+            del weight
         del out[scale_key]
     return out
 
@@ -331,17 +347,19 @@ def _load_from_singlefile(self, config: AnyModelConfig) -> AnyModel:
             raise TypeError(f"Expected Main_Checkpoint_Krea2_Config, got {type(config).__name__}.")
         model_path = Path(config.path)
 
+        target_device = TorchDevice.choose_torch_device()
+        model_dtype = TorchDevice.choose_bfloat16_safe_dtype(target_device)
+
         sd = load_file(model_path)
         sd = _strip_comfyui_prefix(sd)
-        # ComfyUI 'scaled fp8' checkpoints: fold the per-tensor weight_scale into the weights (→ float).
-        sd = _dequantize_scaled_fp8(sd)
+        # ComfyUI 'scaled fp8' checkpoints: fold the per-tensor weight_scale into the weights. The
+        # compute dtype is resolved first so the dequantized weights land there directly instead of
+        # transiently materializing the whole model in float32.
+        sd = _dequantize_scaled_fp8(sd, model_dtype)
         # Native/ComfyUI key naming → diffusers Krea2Transformer2DModel keys.
         if _is_native_krea2_format(sd):
             sd = _convert_krea2_native_to_diffusers(sd)
 
-        target_device = TorchDevice.choose_torch_device()
-        model_dtype = TorchDevice.choose_bfloat16_safe_dtype(target_device)
-
         with accelerate.init_empty_weights():
             model = Krea2Transformer2DModel(**KREA2_TRANSFORMER_CONFIG)
 
@@ -568,7 +586,7 @@ def _load_text_encoder(self, config: Qwen3VLEncoder_Checkpoint_Config) -> AnyMod
             getattr(t, "dtype", None) in (torch.float8_e4m3fn, torch.float8_e5m2) for t in sd.values()
         )
         # ComfyUI 'scaled fp8': fold weight_scale into the weights, then drop quantization metadata.
-        sd = _dequantize_scaled_fp8(sd)
+        sd = _dequantize_scaled_fp8(sd, model_dtype)
         for k in list(sd.keys()):
             if isinstance(k, str) and (k.endswith(".comfy_quant") or "scale_input" in k):
                 del sd[k]

invokeai/backend/model_manager/load/model_loaders/z_image.py

@@ -620,7 +620,13 @@ def _load_from_singlefile(
                             if block_size > 1:
                                 # Repeat scale along this dimension to match weight shape
                                 scale = scale.repeat_interleave(block_size, dim=dim)
-                sd[weight_key] = weight_float * scale
+                # Multiply in float32 for precision, but store the compute dtype immediately so the
+                # *whole model* is never materialized in float32. Keeping every dequantized weight as
+                # float32 until the caller's later cast quadruples the per-parameter cost (4 bytes vs
+                # 1 on disk) and dominates the cold-load RAM peak — enough to swap a 32 GB machine.
+                # Same fix as in the FLUX.2 and Krea-2 loaders.
+                sd[weight_key] = (weight_float * scale).to(model_dtype)
+                del weight_float
                 dequantized_count += 1
 
         if dequantized_count > 0:

tests/backend/model_manager/load/test_krea2_state_dict_utils.py

@@ -90,19 +90,38 @@ def test_folds_scale_into_weight_and_drops_scale_key(self) -> None:
             "layer.weight": torch.tensor([2.0, 4.0]),
             "layer.weight_scale": torch.tensor(0.5),
         }
-        out = _dequantize_scaled_fp8(sd)
+        out = _dequantize_scaled_fp8(sd, torch.bfloat16)
         assert "layer.weight_scale" not in out
-        assert torch.allclose(out["layer.weight"], torch.tensor([1.0, 2.0]))
+        assert torch.allclose(out["layer.weight"].float(), torch.tensor([1.0, 2.0]))
+
+    def test_result_is_stored_in_the_compute_dtype_not_float32(self) -> None:
+        """The whole model must never be materialized in float32.
+
+        The multiply runs in float32 for precision, but holding every dequantized weight there
+        costs 4 bytes per parameter: Krea-2's ~12 GB fp8 checkpoint peaked at ~50 GB of RAM before
+        the caller's later bf16 cast, which swaps a 32 GB machine during a cold load.
+        """
+        sd = {
+            "layer.weight": torch.tensor([2.0, 4.0]),
+            "layer.weight_scale": torch.tensor(0.5),
+        }
+        assert _dequantize_scaled_fp8(dict(sd), torch.bfloat16)["layer.weight"].dtype is torch.bfloat16
+        assert _dequantize_scaled_fp8(dict(sd), torch.float16)["layer.weight"].dtype is torch.float16
+
+    def test_dtype_is_required(self) -> None:
+        """No implicit bfloat16 fallback: on a float16-only device that would cost an extra rounding step."""
+        with pytest.raises(TypeError):
+            _dequantize_scaled_fp8({"layer.weight": torch.tensor([2.0])})  # type: ignore[call-arg]
 
     def test_noop_without_scale_keys(self) -> None:
         sd = {"layer.weight": torch.tensor([2.0, 4.0])}
-        out = _dequantize_scaled_fp8(sd)
+        out = _dequantize_scaled_fp8(sd, torch.bfloat16)
         assert out is sd
 
     def test_orphan_scale_key_is_dropped(self) -> None:
         # A scale key with no matching weight is simply removed (nothing to multiply).
         sd = {"other.weight": torch.tensor([1.0]), "layer.weight_scale": torch.tensor(0.5)}
-        out = _dequantize_scaled_fp8(sd)
+        out = _dequantize_scaled_fp8(sd, torch.bfloat16)
         assert "layer.weight_scale" not in out
         assert "other.weight" in out