diff --git a/invokeai/backend/model_manager/load/model_cache/model_cache.py b/invokeai/backend/model_manager/load/model_cache/model_cache.py index 3691f374425..226bdb5767c 100644 --- a/invokeai/backend/model_manager/load/model_cache/model_cache.py +++ b/invokeai/backend/model_manager/load/model_cache/model_cache.py @@ -34,7 +34,7 @@ ) from invokeai.backend.model_manager.load.model_util import calc_model_size_by_data from invokeai.backend.model_manager.taxonomy import AnyModel, SubModelType -from invokeai.backend.util.devices import TorchDevice +from invokeai.backend.util.devices import TorchDevice, device_has_dedicated_vram from invokeai.backend.util.level_zero import xpu_device_is_integrated from invokeai.backend.util.logging import InvokeAILogger from invokeai.backend.util.prefix_logger_adapter import PrefixedLoggerAdapter @@ -507,25 +507,10 @@ def __call__( _MIN_RELEASABLE_FOR_RESERVE = 256 * MB -def _has_dedicated_vram(device: torch.device) -> bool: - """Whether a device has its own VRAM rather than sharing system memory with the CPU. - - Intel integrated GPUs share memory with the CPU exactly as MPS does, so budgeting them as - if they had dedicated VRAM double-counts the same DRAM (once as "VRAM", once as RAM). When - the Level Zero probe cannot determine the answer, the previous discrete-GPU behaviour is - kept rather than guessing. - """ - if device.type == "cuda": - return True - if device.type == "xpu": - return xpu_device_is_integrated(device) is not True - return False - - def _is_integrated_xpu(device: torch.device) -> bool: """Whether ``device`` is an Intel integrated GPU, i.e. one whose VRAM *is* system RAM. - Narrower than ``not _has_dedicated_vram()``: CPU and MPS also share memory, but their + Narrower than ``not device_has_dedicated_vram()``: CPU and MPS also share memory, but their keep-RAM-copy behaviour is long-standing and out of scope here. """ return device.type == "xpu" and xpu_device_is_integrated(device) is True @@ -1071,7 +1056,7 @@ def put( # loading path there. partial_load_to_vram() respects vram_available, whereas the full-load path ignores it # ("no choice but to try and fit it all"), and on Linux that overshoot is an uncatchable OOM-kill rather than # a degraded load. Combined with keep_ram_copy=False below, streaming moves weights rather than copying them. - running_with_dedicated_vram = _has_dedicated_vram(effective_execution_device) + running_with_dedicated_vram = device_has_dedicated_vram(effective_execution_device) supports_partial_loading = running_with_dedicated_vram or _is_integrated_xpu(effective_execution_device) # On an integrated GPU the "VRAM" copy and the RAM copy are the same DRAM, so keeping a RAM @@ -2368,7 +2353,7 @@ def _get_vram_available(self, working_mem_bytes: Optional[int], honor_cap: bool # for activations). Credit the reclaimable reserve, excluding intra-segment # fragmentation slack that a large contiguous allocation could not use. vram_available_to_process = vram_free + vram_allocated + self._get_reclaimable_allocator_bytes() - elif self._execution_device.type == "xpu" and _has_dedicated_vram(self._execution_device): + elif self._execution_device.type == "xpu" and device_has_dedicated_vram(self._execution_device): vram_allocated = torch.xpu.memory_allocated(self._execution_device) vram_free, _vram_total = TorchDevice.xpu_mem_get_info(self._execution_device) vram_available_to_process = vram_free + vram_allocated + self._get_reclaimable_allocator_bytes() @@ -2495,7 +2480,7 @@ def _calc_ram_available_to_model_cache(self) -> int: total_vram_bytes: int | None = None if self._execution_device.type == "cuda": total_vram_bytes = torch.cuda.get_device_properties(self._execution_device).total_memory - elif self._execution_device.type == "xpu" and _has_dedicated_vram(self._execution_device): + elif self._execution_device.type == "xpu" and device_has_dedicated_vram(self._execution_device): # A total of 0 means "unknown" (the device couldn't report total_memory). Treat it like # the CPU case (None) so the RAM heuristic below doesn't compute a negative cache cap. # @@ -3143,7 +3128,7 @@ def offload_models_from_vram_except(self, keep_model_keys: Collection[str]) -> i Returns the number of VRAM bytes freed. """ - if not _has_dedicated_vram(self._execution_device): + if not device_has_dedicated_vram(self._execution_device): # CPU, MPS, integrated GPUs: their "VRAM" is system RAM, so moving models out frees nothing. return 0 keep = set(keep_model_keys) diff --git a/invokeai/backend/quantization/gguf/ggml_tensor.py b/invokeai/backend/quantization/gguf/ggml_tensor.py index 7bb5532194a..c7f0108575f 100644 --- a/invokeai/backend/quantization/gguf/ggml_tensor.py +++ b/invokeai/backend/quantization/gguf/ggml_tensor.py @@ -9,6 +9,7 @@ TORCH_COMPATIBLE_QTYPES, dequantize, ) +from invokeai.backend.util.devices import TorchDevice, device_has_dedicated_vram def dequantize_and_run(func, args, kwargs): @@ -206,8 +207,8 @@ def __torch_dispatch__(cls, func, types, args, kwargs): # What one dequantization allocates beyond its result, per weight element: the block kernels in # `utils.py` widen codes and scales through integer and float16 intermediates before the final cast. # Measured on CUDA (torch 2.13) for a 4608x53248 weight and rounded up; F32/F16 are a plain cast. Types -# not listed (the IQ family decodes through the numpy fallback) get the worst measured value. The -# host-side counterpart, for decoding a whole tensor at load, is `_PEAK_COPIES` in +# not listed get the worst measured value when they have a device kernel. The host-side counterpart, +# for decoding a whole tensor at load, is `_PEAK_COPIES` in # `model_manager/load/quantized_embedding.py`. _GGMLQ = gguf.GGMLQuantizationType _DEQUANT_INTERMEDIATE_BYTES_PER_ELEMENT: dict[gguf.GGMLQuantizationType, float] = { @@ -227,6 +228,10 @@ def __torch_dispatch__(cls, func, types, args, kwargs): } _DEQUANT_INTERMEDIATE_BYTES_WORST = 8.0 +# NumPy decoding can retain float32 chunks alongside its full float32 output. Keep the existing +# conservative shared-memory allowance separate from the device-kernel fallback. +_HOST_DEQUANT_INTERMEDIATE_BYTES_PER_ELEMENT = 8.0 + def peak_ggml_linear_dequant_transient_bytes(model: torch.nn.Module) -> int: """Peak bytes one forward transiently needs to dequantize this model's GGUF-packed Linear weights. @@ -239,15 +244,24 @@ def peak_ggml_linear_dequant_transient_bytes(model: torch.nn.Module) -> int: the device for the call. Zero when no Linear weight is packed. """ peak = 0 + host_intermediates: float | None = None for module in model.modules(): if not isinstance(module, torch.nn.Linear): continue weight = module.weight if not isinstance(weight, GGMLTensor): continue - intermediates = _DEQUANT_INTERMEDIATE_BYTES_PER_ELEMENT.get( - weight._ggml_quantization_type, _DEQUANT_INTERMEDIATE_BYTES_WORST - ) + qtype = weight._ggml_quantization_type + if qtype in TORCH_COMPATIBLE_QTYPES or qtype in DEQUANTIZE_FUNCTIONS: + intermediates = _DEQUANT_INTERMEDIATE_BYTES_PER_ELEMENT.get(qtype, _DEQUANT_INTERMEDIATE_BYTES_WORST) + else: + if host_intermediates is None: + # Streaming can leave weights on CPU; budget against the session's execution device. + device = TorchDevice.choose_torch_device() + host_intermediates = ( + 0.0 if device_has_dedicated_vram(device) else _HOST_DEQUANT_INTERMEDIATE_BYTES_PER_ELEMENT + ) + intermediates = host_intermediates per_element = weight.compute_dtype.itemsize + intermediates peak = max(peak, math.ceil(weight.tensor_shape.numel() * per_element) + weight.quantized_data.nbytes) return peak diff --git a/invokeai/backend/util/devices.py b/invokeai/backend/util/devices.py index b648adc0c8f..a56524ef3b3 100644 --- a/invokeai/backend/util/devices.py +++ b/invokeai/backend/util/devices.py @@ -61,6 +61,19 @@ def device_is_integrated(device: torch.device) -> Optional[bool]: return None if flag is None else bool(flag) +def device_has_dedicated_vram(device: torch.device) -> bool: + """Whether to budget device memory as dedicated VRAM rather than shared system RAM. + + CUDA/ROCm retains its dedicated-memory budget. Integrated XPU shares system RAM; an unknown + XPU type keeps the discrete-device policy rather than guessing. + """ + if device.type == "cuda": + return True + if device.type == "xpu": + return xpu_device_is_integrated(device) is not True + return False + + def _default_cuda_device() -> torch.device: """The CUDA device ``device: auto`` means: the current one, unless it is an integrated GPU and a discrete one exists. diff --git a/tests/backend/model_manager/load/model_cache/test_model_cache_integrated_gpu.py b/tests/backend/model_manager/load/model_cache/test_model_cache_integrated_gpu.py index 2621aa0ca41..2e094b257bd 100644 --- a/tests/backend/model_manager/load/model_cache/test_model_cache_integrated_gpu.py +++ b/tests/backend/model_manager/load/model_cache/test_model_cache_integrated_gpu.py @@ -7,6 +7,7 @@ """ import logging +from contextlib import contextmanager from types import SimpleNamespace from unittest.mock import MagicMock, patch @@ -56,12 +57,17 @@ def _make_cache( ) +@contextmanager def _patch_integrated(is_integrated: bool | None): """Stub the Level Zero probe. None models a device whose type cannot be determined.""" - return patch( - "invokeai.backend.model_manager.load.model_cache.model_cache.xpu_device_is_integrated", - return_value=is_integrated, - ) + with ( + patch("invokeai.backend.util.devices.xpu_device_is_integrated", return_value=is_integrated), + patch( + "invokeai.backend.model_manager.load.model_cache.model_cache.xpu_device_is_integrated", + return_value=is_integrated, + ), + ): + yield def _logged_messages(logger: MagicMock) -> list[str]: diff --git a/tests/backend/model_manager/load/model_cache/test_model_cache_offload_from_vram.py b/tests/backend/model_manager/load/model_cache/test_model_cache_offload_from_vram.py index 084f6a616b4..368e0620c13 100644 --- a/tests/backend/model_manager/load/model_cache/test_model_cache_offload_from_vram.py +++ b/tests/backend/model_manager/load/model_cache/test_model_cache_offload_from_vram.py @@ -43,7 +43,7 @@ def move(cache_entry, vram_bytes_to_free, keep_required_weights_in_vram=None): with ( patch.object(cache, "_move_model_to_ram", side_effect=move), - patch(f"{MODULE}._has_dedicated_vram", return_value=True), + patch(f"{MODULE}.device_has_dedicated_vram", return_value=True), patch(f"{MODULE}.TorchDevice.empty_cache") as empty_cache, ): yield keys, empty_cache @@ -112,9 +112,15 @@ def test_offload_model_from_vram_still_moves_one_model_with_its_submodels(cache: assert sorted(keys) == ["flux", "flux:vae"] -def test_a_device_whose_vram_is_system_ram_moves_nothing(cache: ModelCache): +@pytest.mark.parametrize("device", ["cpu", "mps", "xpu:2"]) +def test_a_device_whose_vram_is_system_ram_moves_nothing(cache: ModelCache, device: str): """CPU, MPS and integrated GPUs: moving a model "to RAM" would copy it within the same memory.""" - cache.put("flux:transformer", torch.nn.Linear(MB // 4, 1, bias=False)) - with patch.object(cache, "_move_model_to_ram") as move: + _put(cache, "flux:transformer", 1) + with ( + patch.object(cache, "_execution_device", torch.device(device)), + patch("invokeai.backend.util.devices.xpu_device_is_integrated", return_value=True), + patch.object(cache, "_move_model_to_ram", return_value=MB) as move, + patch(f"{MODULE}.TorchDevice.empty_cache"), + ): assert cache.offload_models_from_vram_except(set()) == 0 move.assert_not_called() diff --git a/tests/backend/quantization/gguf/test_ggml_tensor.py b/tests/backend/quantization/gguf/test_ggml_tensor.py index 32b1d537fba..f3755e5213b 100644 --- a/tests/backend/quantization/gguf/test_ggml_tensor.py +++ b/tests/backend/quantization/gguf/test_ggml_tensor.py @@ -4,7 +4,9 @@ from invokeai.backend.quantization.gguf.ggml_tensor import GGMLTensor from invokeai.backend.quantization.gguf.utils import DEQUANTIZE_FUNCTIONS, TORCH_COMPATIBLE_QTYPES +from invokeai.backend.util import devices from invokeai.backend.util.calc_tensor_size import calc_tensor_size +from invokeai.backend.util.devices import TorchDevice def quantize_tensor(data: torch.Tensor, ggml_quantization_type: gguf.GGMLQuantizationType) -> GGMLTensor: @@ -154,10 +156,17 @@ def test_ggml_tensor_numpy_fallback_dequantize_uses_logical_shape(): assert t.get_dequantized_tensor().shape == logical_shape -def _gguf_linear(in_features: int, out_features: int) -> torch.nn.Linear: +def _gguf_linear( + in_features: int, + out_features: int, + qtype: gguf.GGMLQuantizationType = gguf.GGMLQuantizationType.Q8_0, + compute_dtype: torch.dtype = torch.float32, +) -> torch.nn.Linear: + block_size, type_size = gguf.GGML_QUANT_SIZES[qtype] + packed = torch.zeros((out_features, in_features // block_size * type_size), dtype=torch.uint8) linear = torch.nn.Linear(in_features, out_features, bias=False) linear.weight = torch.nn.Parameter( - quantize_tensor(torch.zeros(out_features, in_features), gguf.GGMLQuantizationType.Q8_0), requires_grad=False + GGMLTensor(packed, qtype, torch.Size((out_features, in_features)), compute_dtype), requires_grad=False ) return linear @@ -170,17 +179,15 @@ def test_the_shared_dequant_transient_counts_gguf_linears(): from invokeai.backend.quantization.int8_convrot import Int8ConvrotLinear gguf_model = torch.nn.Sequential(_gguf_linear(64, 128), _gguf_linear(64, 32)) - largest = gguf_model[0].weight gguf_peak = peak_dequant_transient_bytes(gguf_model, torch.float32) - # At least the copy of the larger layer, in the weight's own compute dtype, plus its packed bytes; - # never the sum of both layers. - assert gguf_peak > largest.tensor_shape.numel() * largest.compute_dtype.itemsize + largest.quantized_data.nbytes - assert gguf_peak == peak_dequant_transient_bytes(torch.nn.Sequential(gguf_model[0]), torch.float32) + # 8192 elements * (4-byte result + 2-byte Q8 intermediates), plus 256 packed 34-byte blocks. + assert gguf_peak == 57856 + assert peak_dequant_transient_bytes(torch.nn.Sequential(gguf_model[0]), torch.float32) == 57856 # Whichever scheme holds the largest layer decides, in either order. larger_int8 = Int8ConvrotLinear( - weight=torch.zeros(4096, 4096, dtype=torch.int8), weight_scale=torch.ones(4096, 1), convrot=False + weight=torch.zeros(128, 128, dtype=torch.int8), weight_scale=torch.ones(128, 1), convrot=False ) mixed = torch.nn.Sequential(gguf_model[0], larger_int8) assert peak_dequant_transient_bytes(mixed, torch.float32) == larger_int8.dequant_transient_bytes(torch.float32) @@ -190,3 +197,115 @@ def test_the_shared_dequant_transient_counts_gguf_linears(): assert smaller_int8.dequant_transient_bytes(torch.float32) < gguf_peak assert peak_dequant_transient_bytes(torch.nn.Sequential(smaller_int8, gguf_model[0]), torch.float32) == gguf_peak assert peak_dequant_transient_bytes(torch.nn.Sequential(torch.nn.Linear(64, 64)), torch.float32) == 0 + + +@pytest.mark.parametrize("compute_dtype", [torch.float16, torch.float32], ids=["fp16", "fp32"]) +def test_host_dequant_reservation_counts_only_result_and_packed_weight(compute_dtype, monkeypatch): + """Host decode intermediates do not consume dedicated VRAM, even with the weight still on CPU.""" + from invokeai.backend.quantization.dequantizing_linear import peak_dequant_transient_bytes + from invokeai.backend.quantization.gguf.ggml_tensor import peak_ggml_linear_dequant_transient_bytes + + monkeypatch.setattr(TorchDevice, "choose_torch_device", lambda: torch.device("cuda:1")) + linear = _gguf_linear(256, 2, gguf.GGMLQuantizationType.IQ4_NL, compute_dtype) + weight = linear.weight + assert weight.device.type == "cpu" + decoded = weight.get_dequantized_tensor() + assert decoded.dtype == compute_dtype + + model = torch.nn.Sequential(linear) + expected_bytes = decoded.nbytes + weight.quantized_data.nbytes + assert peak_ggml_linear_dequant_transient_bytes(model) == expected_bytes + # GGUF carries its own dtype, even if a caller passes a different module-level dtype. + assert peak_dequant_transient_bytes(model, torch.float64) == expected_bytes + + +@pytest.mark.parametrize( + ("execution_device", "integrated", "compute_dtype", "intermediate_bytes"), + [ + pytest.param("cpu", None, torch.float32, 8, id="cpu-fp32"), + pytest.param("mps", None, torch.float16, 8, id="mps-fp16"), + pytest.param("xpu:2", True, torch.float32, 8, id="integrated-xpu-fp32"), + pytest.param("xpu:2", False, torch.float16, 0, id="discrete-xpu-fp16"), + pytest.param("xpu:2", None, torch.float16, 0, id="unknown-xpu-fp16"), + ], +) +def test_host_dequant_reservation_follows_execution_memory_budget( + execution_device, integrated, compute_dtype, intermediate_bytes, monkeypatch +): + """Shared execution retains the prior conservative 8 bytes/element for host intermediates.""" + from invokeai.backend.quantization.dequantizing_linear import peak_dequant_transient_bytes + + device = torch.device(execution_device) + monkeypatch.setattr(TorchDevice, "choose_torch_device", lambda: device) + + def is_integrated(probed_device): + assert probed_device == device + return integrated + + monkeypatch.setattr(devices, "xpu_device_is_integrated", is_integrated) + linear = _gguf_linear(256, 2, gguf.GGMLQuantizationType.IQ4_NL, compute_dtype) + assert linear.weight.device.type == "cpu" + # 512 logical elements and sixteen 18-byte IQ4_NL blocks; the weight owns the result dtype. + expected_bytes = 512 * (compute_dtype.itemsize + intermediate_bytes) + 288 + assert peak_dequant_transient_bytes(linear, torch.float64) == expected_bytes + + +def test_shared_host_dequant_reservation_is_independent_of_device_kernel_bound(monkeypatch): + """Lowering the device-kernel allowance must not shrink shared host working memory.""" + from invokeai.backend.quantization.dequantizing_linear import peak_dequant_transient_bytes + from invokeai.backend.quantization.gguf import ggml_tensor + + monkeypatch.setattr(ggml_tensor, "_DEQUANT_INTERMEDIATE_BYTES_WORST", 0.0) + monkeypatch.setattr(TorchDevice, "choose_torch_device", lambda: torch.device("mps")) + linear = _gguf_linear(256, 2, gguf.GGMLQuantizationType.IQ4_NL, torch.float16) + assert linear.weight.device.type == "cpu" + # 512 elements * (2-byte result + 8-byte host allowance), plus sixteen 18-byte packed blocks. + assert peak_dequant_transient_bytes(linear, torch.float64) == 5408 + + +def test_bf16_device_dequant_reservation_keeps_measured_intermediates(): + from invokeai.backend.quantization.dequantizing_linear import peak_dequant_transient_bytes + + linear = _gguf_linear(256, 2, gguf.GGMLQuantizationType.BF16, torch.float16) + # 512 elements * (2-byte result + 6-byte BF16 intermediates), plus 1024 packed bytes. + assert peak_dequant_transient_bytes(linear, torch.float64) == 5120 + + +def test_dedicated_dequant_reservation_uses_largest_bytes_in_either_order(monkeypatch): + from invokeai.backend.quantization.gguf.ggml_tensor import peak_ggml_linear_dequant_transient_bytes + + monkeypatch.setattr(TorchDevice, "choose_torch_device", lambda: torch.device("cuda:1")) + larger_iq = _gguf_linear(256, 16, gguf.GGMLQuantizationType.IQ4_NL) + smaller_q8 = _gguf_linear(256, 12) + # IQ: 4096 * 4 + 2304 = 18688. Q8: 3072 * (4 + 2) + 3264 = 21696. + for layers in ((larger_iq, smaller_q8), (smaller_q8, larger_iq)): + assert peak_ggml_linear_dequant_transient_bytes(torch.nn.Sequential(*layers)) == 21696 + + +def test_host_dequant_reservation_follows_execution_device_changes(monkeypatch): + from invokeai.backend.quantization.gguf.ggml_tensor import peak_ggml_linear_dequant_transient_bytes + + device = torch.device("xpu:2") + monkeypatch.setattr(TorchDevice, "choose_torch_device", lambda: device) + monkeypatch.setattr(devices, "xpu_device_is_integrated", lambda _: True) + model = torch.nn.Sequential( + _gguf_linear(256, 2, gguf.GGMLQuantizationType.IQ4_NL, torch.float16), + _gguf_linear(256, 1, gguf.GGMLQuantizationType.IQ4_NL, torch.float16), + ) + + assert peak_ggml_linear_dequant_transient_bytes(model) == 5408 # 512 * (2 + 8) + 288. + device = torch.device("cuda:1") + assert peak_ggml_linear_dequant_transient_bytes(model) == 1312 # 512 * 2 + 288. + + +def test_new_device_dequant_kernel_keeps_conservative_intermediate_reservation(monkeypatch): + from invokeai.backend.quantization.gguf.ggml_tensor import peak_ggml_linear_dequant_transient_bytes + + qtype = gguf.GGMLQuantizationType.IQ4_NL + packed = torch.zeros((16, 18), dtype=torch.uint8) + weight = GGMLTensor(packed, qtype, torch.Size((2, 256)), torch.float16) + linear = torch.nn.Linear(256, 2, bias=False) + linear.weight = torch.nn.Parameter(weight, requires_grad=False) + # A newly registered device kernel has no measured table entry yet. Do not treat it as host-only. + monkeypatch.setitem(DEQUANTIZE_FUNCTIONS, qtype, DEQUANTIZE_FUNCTIONS[gguf.GGMLQuantizationType.Q8_0]) + assert peak_ggml_linear_dequant_transient_bytes(linear) == 512 * (2 + 8) + packed.nbytes diff --git a/tests/backend/util/test_devices.py b/tests/backend/util/test_devices.py index 5bca941e44b..9ee15b63462 100644 --- a/tests/backend/util/test_devices.py +++ b/tests/backend/util/test_devices.py @@ -15,6 +15,7 @@ TorchDevice, choose_precision, choose_torch_device, + device_has_dedicated_vram, device_is_integrated, torch_dtype, ) @@ -33,6 +34,30 @@ def test_device_choice(device_name): assert torch_device == torch.device(device_name) +@pytest.mark.parametrize( + ("device_name", "integrated", "expected"), + [ + pytest.param("cpu", None, False, id="cpu-shared"), + pytest.param("mps", None, False, id="mps-shared"), + pytest.param("cuda:1", True, True, id="rocm-apu-keeps-dedicated-budget"), + pytest.param("xpu:2", True, False, id="integrated-xpu-shared"), + pytest.param("xpu:2", False, True, id="discrete-xpu-dedicated"), + pytest.param("xpu:2", None, True, id="unknown-xpu-dedicated"), + ], +) +def test_device_memory_budget_classification(device_name, integrated, expected, monkeypatch): + device = torch.device(device_name) + + def probe(probed_device): + assert probed_device == device + return integrated + + monkeypatch.setattr("invokeai.backend.util.devices.xpu_device_is_integrated", probe) + monkeypatch.setattr(torch.cuda, "get_device_properties", lambda _: SimpleNamespace(is_integrated=1)) + monkeypatch.setattr(torch.version, "hip", "7.15.26333") + assert device_has_dedicated_vram(device) is expected + + # ===== per-thread session device (multi-GPU worker pinning) ================