"""Tests for `ModelLoader` FP8 helpers. Covers: - `_should_use_fp8` excludes ControlLoRA (the LoRA loader never runs the layerwise casting helper, and a LoRA isn't a standalone forward module — so a persisted `fp8_storage=true` must be a no-op). - `_wrap_forward_with_fp8_cast` uses pre/post hooks with `always_call=True`, so it is exception-safe AND survives `apply_custom_layers_to_model`'s instance swap. Without hooks, an instance-level `forward` override would be carried into the new CustomLinear via the shared `__dict__` and silently bypass `CustomLinear.forward` — breaking LoRA patch dispatch for FP8 checkpoint models. - `_apply_fp8_to_nn_module` skips precision-sensitive layers (norm, pos_embed, etc.) so FLUX RMSNorm.scale and friends aren't crushed to FP8. """ from logging import getLogger from types import SimpleNamespace from unittest.mock import patch import pytest import torch from invokeai.backend.model_manager.load.load_default import ModelLoader from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.custom_modules.custom_linear import ( CustomLinear, ) from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.torch_module_autocast import ( apply_custom_layers_to_model, ) from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelType def _make_loader(device: str = "cuda") -> ModelLoader: """Build a ModelLoader without going through dependency injection. `_should_use_fp8` and `_wrap_forward_with_fp8_cast` only depend on `_torch_device`, so we instantiate via __new__ and set the minimum state directly. """ loader = ModelLoader.__new__(ModelLoader) loader._torch_device = torch.device(device) loader._torch_dtype = torch.float16 loader._logger = getLogger("test") return loader def _make_config(model_type: ModelType, fp8: bool, base: BaseModelType = BaseModelType.Flux): return SimpleNamespace( type=model_type, base=base, name="test", default_settings=SimpleNamespace(fp8_storage=fp8), ) def test_should_use_fp8_excludes_control_lora(): """ControlLoRA gets the FP8 toggle in the UI history but the LoRA loader never applies layerwise casting (the model isn't run as a standalone forward pass — it patches into a base model). The loader must silently ignore a persisted `fp8_storage=true` to avoid misleading users who toggled it under a prior version. """ loader = _make_loader(device="cuda") with patch("torch.cuda.is_available", return_value=True): assert loader._should_use_fp8(_make_config(ModelType.ControlLoRa, fp8=True)) is False def test_should_use_fp8_excludes_lora(): loader = _make_loader(device="cuda") assert loader._should_use_fp8(_make_config(ModelType.LoRA, fp8=True)) is False def test_should_use_fp8_returns_true_for_main_with_fp8(): loader = _make_loader(device="cuda") assert loader._should_use_fp8(_make_config(ModelType.Main, fp8=True)) is True def test_should_use_fp8_returns_false_for_main_without_fp8(): loader = _make_loader(device="cuda") assert loader._should_use_fp8(_make_config(ModelType.Main, fp8=False)) is False def test_should_use_fp8_returns_false_on_cpu(): loader = _make_loader(device="cpu") assert loader._should_use_fp8(_make_config(ModelType.Main, fp8=True)) is False class _RaisingModule(torch.nn.Module): """A module whose forward unconditionally raises — used to test that the FP8 wrapper's storage-dtype cleanup runs even when forward fails.""" def __init__(self): super().__init__() self.weight = torch.nn.Parameter(torch.zeros(4)) self.bias = torch.nn.Parameter(torch.zeros(4)) def forward(self, x: torch.Tensor) -> torch.Tensor: raise RuntimeError("boom") def _fp8_supported() -> bool: return hasattr(torch, "float8_e4m3fn") @pytest.mark.skipif(not _fp8_supported(), reason="torch.float8_e4m3fn not available") def test_wrap_forward_restores_storage_dtype_on_exception(): """When forward raises, params must be returned to storage dtype. Otherwise FP8 storage savings silently revert to fp16/bf16 and the cache's size accounting becomes stale. """ storage_dtype = torch.float8_e4m3fn compute_dtype = torch.bfloat16 module = _RaisingModule() for p in module.parameters(recurse=False): p.data = p.data.to(storage_dtype) ModelLoader._wrap_forward_with_fp8_cast(module, storage_dtype, compute_dtype) # Sanity: params start in storage dtype. assert module.weight.dtype == storage_dtype assert module.bias.dtype == storage_dtype with pytest.raises(RuntimeError, match="boom"): module(torch.zeros(4, dtype=compute_dtype)) # Critical assertion: cleanup ran despite the exception. assert module.weight.dtype == storage_dtype assert module.bias.dtype == storage_dtype @pytest.mark.skipif(not _fp8_supported(), reason="torch.float8_e4m3fn not available") def test_wrap_forward_casts_to_compute_then_back_on_success(): """Happy-path sanity check: params are in compute dtype during forward, storage dtype after.""" storage_dtype = torch.float8_e4m3fn compute_dtype = torch.bfloat16 seen_dtypes: list[torch.dtype] = [] class _CaptureModule(torch.nn.Module): def __init__(self): super().__init__() self.weight = torch.nn.Parameter(torch.zeros(4)) def forward(self, x: torch.Tensor) -> torch.Tensor: seen_dtypes.append(self.weight.dtype) return x + self.weight module = _CaptureModule() for p in module.parameters(recurse=False): p.data = p.data.to(storage_dtype) ModelLoader._wrap_forward_with_fp8_cast(module, storage_dtype, compute_dtype) module(torch.zeros(4, dtype=compute_dtype)) assert seen_dtypes == [compute_dtype] assert module.weight.dtype == storage_dtype def test_apply_fp8_to_nn_module_uses_wrapper(): """`_apply_fp8_to_nn_module` should delegate per-module wrapping to `_wrap_forward_with_fp8_cast`, which encapsulates the hook registration. """ module = torch.nn.Linear(4, 4) with patch.object(ModelLoader, "_wrap_forward_with_fp8_cast") as mock_wrap: ModelLoader._apply_fp8_to_nn_module(module, torch.float16, torch.float32) mock_wrap.assert_called_once_with(module, torch.float16, torch.float32) def test_apply_fp8_to_nn_module_skips_norm_modules(): """Modules whose path matches `norm` must not be cast — diffusers' `enable_layerwise_casting` does the same. FLUX RMSNorm.scale is the canonical example: a tiny learned scalar that breaks badly in FP8. """ class _Model(torch.nn.Module): def __init__(self): super().__init__() self.norm1 = torch.nn.LayerNorm(4) self.linear = torch.nn.Linear(4, 4) storage_dtype = torch.float16 compute_dtype = torch.float32 model = _Model() for p in model.parameters(): p.data = p.data.to(compute_dtype) ModelLoader._apply_fp8_to_nn_module(model, storage_dtype, compute_dtype) # Linear params get cast to storage dtype. assert model.linear.weight.dtype == storage_dtype # Norm params stay in compute dtype — they must not be cast. assert model.norm1.weight.dtype == compute_dtype assert model.norm1.bias.dtype == compute_dtype def test_apply_fp8_to_nn_module_skips_pos_embed_and_proj_in_out(): """Position embeddings and the in/out projection of transformer blocks are also on the diffusers default skip list — they're precision-sensitive. """ class _Model(torch.nn.Module): def __init__(self): super().__init__() self.pos_embed = torch.nn.Linear(4, 4) self.proj_in = torch.nn.Linear(4, 4) self.proj_out = torch.nn.Linear(4, 4) self.attn = torch.nn.Linear(4, 4) storage_dtype = torch.float16 compute_dtype = torch.float32 model = _Model() for p in model.parameters(): p.data = p.data.to(compute_dtype) ModelLoader._apply_fp8_to_nn_module(model, storage_dtype, compute_dtype) assert model.attn.weight.dtype == storage_dtype assert model.pos_embed.weight.dtype == compute_dtype assert model.proj_in.weight.dtype == compute_dtype assert model.proj_out.weight.dtype == compute_dtype def test_apply_fp8_to_nn_module_skips_unsupported_layer_types(): """Only the layer classes in `_FP8_SUPPORTED_PYTORCH_LAYERS` are cast — matches diffusers' behavior. A custom RMSNorm-style module with a raw Parameter must be left alone, otherwise its learned scalar gets clobbered. """ class _ScaleModule(torch.nn.Module): """Mimics FLUX RMSNorm — a tiny learned scalar that must not be cast to FP8.""" def __init__(self): super().__init__() self.scale = torch.nn.Parameter(torch.ones(4)) def forward(self, x: torch.Tensor) -> torch.Tensor: return x * self.scale class _Model(torch.nn.Module): def __init__(self): super().__init__() self.rms = _ScaleModule() self.linear = torch.nn.Linear(4, 4) storage_dtype = torch.float16 compute_dtype = torch.float32 model = _Model() for p in model.parameters(): p.data = p.data.to(compute_dtype) ModelLoader._apply_fp8_to_nn_module(model, storage_dtype, compute_dtype) assert model.linear.weight.dtype == storage_dtype # Critical: the RMS-style scalar lives on a custom module type, not in the supported list. assert model.rms.scale.dtype == compute_dtype def test_wrap_forward_reaches_custom_linear_after_apply_custom_layers(): """Production order: `_load_model` applies FP8 wrapping, THEN `ModelCache.put()` calls `apply_custom_layers_to_model` which constructs a NEW `CustomLinear` object via `CustomLinear.__new__` and points its `__dict__` at the original `Linear.__dict__` (see `wrap_custom_layer`). The new object is installed on the parent in place of the original Linear. An instance-level `forward` override would be carried into the new CustomLinear via the shared dict but would close over the OLD Linear instance — so calls to the new CustomLinear would silently route to `Linear.forward(old_instance, ...)` and bypass `CustomLinear.forward`, where LoRA/ControlLoRA patches are applied. This is the bug a reviewer reproduced on a fresh worktree. Hooks fix this because `nn.Module._call_impl` dispatches them with the *actual* called instance, and `self.forward(...)` is resolved by normal class lookup — reaching `CustomLinear.forward`. This test exercises the production wrapping path (real `apply_custom_layers_to_model`) and asserts CustomLinear.forward is reached by attaching a sentinel patch list and observing that the patch-aware branch runs. """ class Parent(torch.nn.Module): def __init__(self): super().__init__() self.child = torch.nn.Linear(4, 4, bias=False) parent = Parent() original_linear = parent.child ModelLoader._wrap_forward_with_fp8_cast(original_linear, torch.float16, torch.float32) apply_custom_layers_to_model(parent) new_child = parent.child # Sanity: production wrapping replaced the child with a NEW CustomLinear instance. assert isinstance(new_child, CustomLinear) assert new_child is not original_linear # Attach a sentinel patch so CustomLinear.forward routes through the LoRA-aware branch # (see custom_linear.py: `if len(self._patches_and_weights) > 0`). If that branch fires, # our FP8 wrapping is correctly dispatched through CustomLinear.forward. patch_was_invoked = {"hit": False} class _SentinelPatch: def __init__(self): self.hit = patch_was_invoked def __call__(self, *_args, **_kwargs): # not actually called pass # Patch the CustomLinear's patch-handling branch to record that it was reached. original_patch_branch = CustomLinear._autocast_forward_with_patches def tracked_patch_branch(self, input): patch_was_invoked["hit"] = True # Return a same-shape tensor so the outer caller doesn't choke. return torch.zeros_like(input @ self.weight.t()) new_child._patches_and_weights = [(_SentinelPatch(), 1.0)] try: CustomLinear._autocast_forward_with_patches = tracked_patch_branch _ = new_child(torch.zeros(1, 4, dtype=torch.float32)) finally: CustomLinear._autocast_forward_with_patches = original_patch_branch new_child._patches_and_weights = [] assert patch_was_invoked["hit"] is True, ( "FP8-wrapped forward did not reach CustomLinear.forward — LoRA/ControlLoRA patches " "would be silently bypassed on FP8 checkpoint models." ) def test_apply_fp8_layerwise_casting_uses_hook_path_for_model_mixin(): """Regression test for the FLUX.2 Klein 9B partial-load device-mismatch crash. Diffusers' `enable_layerwise_casting()` registers a `LayerwiseCastingHook` whose `pre_forward` only casts dtype (not device) and whose hook system replaces `Linear.forward` with a wrapper that calls the *original* `Linear.forward` captured before the hook was installed. `ModelCache.put()` later wraps Linear as CustomLinear sharing `__dict__`, so the diffusers wrapper is carried into the new CustomLinear and routes calls to the captured original Linear.forward — bypassing `CustomLinear.forward`'s `cast_to_device`. On partial load (some weights on CPU, input on cuda), this raises a device-mismatch error. The fix routes ModelMixin through `_apply_fp8_to_nn_module` (hook-based, `forward`-preserving). This test asserts that path is taken even when the model inherits from ModelMixin. """ from diffusers.models.modeling_utils import ModelMixin class _FakeModelMixin(ModelMixin): # ModelMixin requires a config_name class attribute and a config dict for serialization. # We never serialize, so we only need to satisfy isinstance() checks. config_name = "config.json" def __init__(self): super().__init__() self.linear = torch.nn.Linear(4, 4, bias=False) def forward(self, x): return self.linear(x) loader = _make_loader(device="cuda") config = _make_config(ModelType.Main, fp8=True) model = _FakeModelMixin() with ( patch.object(ModelLoader, "_should_use_fp8", return_value=True), patch.object(ModelLoader, "_apply_fp8_to_nn_module") as mock_to_nn, patch.object(_FakeModelMixin, "enable_layerwise_casting") as mock_enable, ): loader._apply_fp8_layerwise_casting(model, config) mock_to_nn.assert_called_once() mock_enable.assert_not_called()