From dd211bb6b04d82bccd020dedb7d46597758c7d4b Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Sat, 1 Aug 2026 06:08:55 +0200 Subject: [PATCH] fix(model-loaders): stop materializing scaled-fp8 checkpoints in float32 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The Krea-2 and Z-Image loaders dequantized ComfyUI 'scaled fp8' weights with `weight.float() * scale` and left the result in float32 until a much later cast to the compute dtype. That holds the entire model at 4 bytes per parameter: a ~12 GB Krea-2 fp8 checkpoint peaks at ~50 GB of RAM before dropping to ~25 GB, which puts a 32 GB machine into swap during a cold load — before anything moves toward VRAM. Both now multiply in float32 for precision but store the compute dtype immediately, halving the cold-load peak. This is the same fix the FLUX.2 loader already carries; its comment documents the identical symptom (~36 GB vs ~17 GB for a 9B model). The Qwen-Image loader was already correct. _dequantize_scaled_fp8 takes the target dtype as a parameter (defaulting to bfloat16), and the Krea-2 single-file loader resolves the compute dtype before calling it so the weights land in their final type directly instead of being cast twice. --- .../model_manager/load/model_loaders/krea2.py | 36 +++++++++++++------ .../load/model_loaders/z_image.py | 8 ++++- .../load/test_krea2_state_dict_utils.py | 16 ++++++++- 3 files changed, 47 insertions(+), 13 deletions(-) diff --git a/invokeai/backend/model_manager/load/model_loaders/krea2.py b/invokeai/backend/model_manager/load/model_loaders/krea2.py index 98cd4a2dec3..ea4c6e81f2d 100644 --- a/invokeai/backend/model_manager/load/model_loaders/krea2.py +++ b/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,25 +93,33 @@ 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 | None" = None) -> dict[str, Any]: """Dequantize ComfyUI 'scaled fp8' weights: ``dequant = weight.float() * weight_scale``. Each quantized layer stores an fp8 ``.weight`` plus a (usually scalar) ``.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. """ import torch scale_keys = [k for k in sd if isinstance(k, str) and k.endswith(".weight_scale")] if not scale_keys: return sd + dtype = dtype or torch.bfloat16 out = dict(sd) for scale_key in scale_keys: weight_key = scale_key.replace(".weight_scale", ".weight") 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 +343,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 +582,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] diff --git a/invokeai/backend/model_manager/load/model_loaders/z_image.py b/invokeai/backend/model_manager/load/model_loaders/z_image.py index 406d243c371..8e26b8f7e18 100644 --- a/invokeai/backend/model_manager/load/model_loaders/z_image.py +++ b/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: diff --git a/tests/backend/model_manager/load/test_krea2_state_dict_utils.py b/tests/backend/model_manager/load/test_krea2_state_dict_utils.py index c75873a8fdc..22fd59637fb 100644 --- a/tests/backend/model_manager/load/test_krea2_state_dict_utils.py +++ b/tests/backend/model_manager/load/test_krea2_state_dict_utils.py @@ -92,7 +92,21 @@ def test_folds_scale_into_weight_and_drops_scale_key(self) -> None: } out = _dequantize_scaled_fp8(sd) 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))["layer.weight"].dtype is torch.bfloat16 + assert _dequantize_scaled_fp8(dict(sd), torch.float16)["layer.weight"].dtype is torch.float16 def test_noop_without_scale_keys(self) -> None: sd = {"layer.weight": torch.tensor([2.0, 4.0])}