From 0d866998f06db4dcbe91838dea923c88d188f5ca Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 18 Jul 2026 03:49:05 +0000 Subject: [PATCH] Fix silent LoRA drop and wasted transformer prefetch on GGUF quant loads Two live-test findings on the images load path: - transformer_quant with baked LoRAs, when the dense quantized build is declined for memory or fails: the load completed as a plain GGUF with the adapters silently dropped (HTTP success, supports_lora=false after the fact) -- wrong output with no signal. The load now fails with the recovery options (drop the adapters, free VRAM, or pick a smaller model). Weight-0 adapters still count as no bake request, and the plain no-LoRA decline keeps its silent GGUF fallback. - A fresh GGUF load on a small GPU prefetched the base repo's full bf16 transformer shards (~47 GB on Qwen-Image) because the dense-quant prefetch widening only checked scheme viability, not whether the device could ever hold the candidate resident. Gate the widening on total device capacity (reserve + 0.85 margin, the plan_fits_total_capacity bar) so a card that is certain to decline the dense build never pays the download; capable devices keep the prefetch. --- studio/backend/core/inference/diffusion.py | 40 +++++- .../backend/tests/test_diffusion_backend.py | 114 ++++++++++++++++++ 2 files changed, 153 insertions(+), 1 deletion(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 364df883ec..e49e29c492 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -580,7 +580,26 @@ class DiffusionBackend: ) # A prequant candidate loads a small checkpoint, not the dense transformer/ shards, # so widening for it defeats the savings and can disk-full the fallback-less begin_load pull. - return candidate is not None and not candidate.prequant + if candidate is None or candidate.prequant: + return False + # Capacity gate: on a device that cannot hold even the candidate's post-quant + # resident set, load_pipeline's re-plan is CERTAIN to decline the dense path, so + # widening would fetch the multi-GB base transformer/ shards (47 GB on Qwen-Image) + # only to run the GGUF as-is. Mirror plan_fits_total_capacity's bar against TOTAL + # capacity (not instantaneous free, which a transient tenant could undercount). + from .diffusion_memory import ( + _reserve_mib, + snapshot_device_memory, + ) + + memory = snapshot_device_memory(target) + total = memory.total_mib + steady = getattr(candidate, "steady_total_mib", None) + if total is not None and steady is not None: + budget = int((int(total) - _reserve_mib(memory.memory_kind, int(total))) * 0.85) + if int(steady) > budget: + return False + return True except Exception: # noqa: BLE001 — widening the prefetch is best-effort only return False @@ -1379,6 +1398,25 @@ class DiffusionBackend: # The engaged dense build uses the re-planned placement; the GGUF-size plan stays for fallback. plan = quant_plan + if ( + pipe is None + and kind == "gguf" + and normalize_transformer_quant(transformer_quant) is not None + and any(w != 0 for (_lid, w) in (loras or ())) + ): + # The client asked for adapters BAKED into a quantized build, and that + # build was declined (memory plan) or failed. The GGUF fallback cannot + # carry the adapters, so completing it would return HTTP success while + # silently generating without the requested LoRAs -- wrong output with + # no signal. Fail the load with the recovery options instead. + raise RuntimeError( + "The requested LoRA adapters could not be applied: baking adapters " + "requires the quantized (int8/fp8) transformer build, which was " + "declined or failed on this device (see the server log), and the " + "GGUF fallback cannot carry them. Retry without transformer_quant " + "adapters, free VRAM, or pick a smaller model." + ) + if pipe is None: if kind == "pipeline": # Full diffusers repo: from_pretrained pulls every component and re-applies diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index e2f49d2c6f..152620c909 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -2842,6 +2842,82 @@ def test_dense_quant_replan_no_retry_when_capacity_truly_short( assert replan_calls == [True] # genuine capacity shortfall: declined without a retry +def _decline_dense_quant(backend, monkeypatch, tmp_path): + """Configure the harness so the dense-quant fast path is declined for capacity + (mirrors test_dense_quant_replan_no_retry_when_capacity_truly_short).""" + import dataclasses + + from core.inference import diffusion as dmod + + _force_cuda_target(backend, monkeypatch) + monkeypatch.setattr(dmod, "dense_transformer_supported", lambda target: True) + monkeypatch.setattr( + dmod, "select_transformer_quant_scheme", lambda target, mode, family = None: "int8" + ) + monkeypatch.setattr( + dmod, + "resolve_dense_quant_candidate", + lambda **kw: types.SimpleNamespace( + transient_transformer_mib = 33_831, companions_mib = 46_157, prequant = False + ), + ) + orig_plan = DiffusionBackend._plan_memory + + def spy_plan(self, *a, transformer_resident_override_mib = None, **k): + real = orig_plan( + self, *a, transformer_resident_override_mib = transformer_resident_override_mib, **k + ) + if transformer_resident_override_mib is None: + return dataclasses.replace(real, offload_policy = "model") + return types.SimpleNamespace( + offload_policy = "model", + estimates = {"resident_required_mib": 150_000, "safe_device_budget_mib": 40_000}, + device_memory = types.SimpleNamespace( + total_mib = 183_359, memory_kind = "discrete_vram", free_mib = 60_000 + ), + reasons = ("companions exceed budget",), + ) + + monkeypatch.setattr(DiffusionBackend, "_plan_memory", spy_plan) + (tmp_path / "m.gguf").write_bytes(b"x") + + +def test_declined_dense_with_baked_loras_fails_instead_of_silent_drop( + fake_runtime, tmp_path, monkeypatch +): + # transformer_quant + adapters, dense build declined for capacity: the GGUF fallback + # cannot bake the adapters, so completing it would silently generate WITHOUT the + # requested LoRAs behind an HTTP success. The load must fail with the recovery options. + backend = DiffusionBackend() + _decline_dense_quant(backend, monkeypatch, tmp_path) + with pytest.raises(RuntimeError, match = "LoRA adapters could not be applied"): + backend.load_pipeline( + str(tmp_path), + gguf_filename = "m.gguf", + family_override = "z-image", + transformer_quant = "int8", + loras = [("adapter", 1.0)], + ) + + +def test_declined_dense_without_loras_still_falls_back_to_gguf( + fake_runtime, tmp_path, monkeypatch +): + # The plain decline (no adapters requested) keeps the silent GGUF fallback: weight-0 + # adapters count as "none" (an explicit disable is not a bake request). + backend = DiffusionBackend() + _decline_dense_quant(backend, monkeypatch, tmp_path) + result = backend.load_pipeline( + str(tmp_path), + gguf_filename = "m.gguf", + family_override = "z-image", + transformer_quant = "int8", + loras = [("adapter", 0.0)], + ) + assert result is not None + assert backend.status()["transformer_quant"] is None # GGUF-as-is fallback + + class _BakePipe: def __init__(self): self.calls: list = [] @@ -3125,6 +3201,44 @@ def test_base_file_downloaded_include_transformer_flag(): assert _base_file_downloaded("README.md", include_transformer = True) is False +def test_dense_quant_prefetch_capacity_gate(fake_runtime, monkeypatch): + # On a device that cannot hold even the candidate's post-quant resident set, the re-plan + # is certain to decline the dense path -- widening would fetch the multi-GB base + # transformer/ shards (measured: ~47 GB on a Qwen-Image GGUF load on a 24 GB card) only + # to run the GGUF as-is. The gate compares steady_total against TOTAL capacity (reserve + + # 0.85 margin), never the instantaneous free reading. + from core.inference import diffusion as dmod + from core.inference import diffusion_memory as dmem + + backend = DiffusionBackend() + _force_cuda_target(backend, monkeypatch) + fam = detect_family("unsloth/Qwen-Image-GGUF") + + def candidate_with(steady): + return lambda **kw: types.SimpleNamespace(prequant = False, steady_total_mib = steady) + + monkeypatch.setattr( + dmem, + "snapshot_device_memory", + lambda target: types.SimpleNamespace( + total_mib = 24_564, free_mib = 24_000, memory_kind = "discrete_vram" + ), + ) + # int8 qwen steady (~22 GB DiT + 17 GB companions = 39 GB) cannot fit a 24 GB card. + monkeypatch.setattr(dmod, "resolve_dense_quant_candidate", candidate_with(39_900)) + assert backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "int8"}) is False + # A candidate that fits total capacity still widens. + monkeypatch.setattr(dmod, "resolve_dense_quant_candidate", candidate_with(12_000)) + assert backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "int8"}) is True + # Unknown sizes keep the old behaviour (widen: the loader may still take the dense path). + monkeypatch.setattr( + dmod, + "resolve_dense_quant_candidate", + lambda **kw: types.SimpleNamespace(prequant = False), + ) + assert backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "int8"}) is True + + def test_dense_quant_prefetch_needed_gates(fake_runtime, monkeypatch): # The transformer/ prefetch widens exactly when load_pipeline takes the dense-quant path: it # defers to resolve_dense_quant_candidate (quant requested + device supported + scheme