ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Studio diffusion (Phase 7): robust backend-flag snapshot/restore and restore on failed speeded load

Browse source 
- snapshot_backend_flags reads each flag defensively (getattr + hasattr), so a build/platform
  missing one (no cuda.matmul on CPU/MPS) still captures the rest instead of skipping the
  whole snapshot. restore_backend_flags restores each flag independently so one failure can't
  leave the others leaked process-wide.
- load_pipeline restores the flags (and clears the GPU cache) when the build fails after
  apply_speed_optims mutated the process-wide flags but before _state captured them for unload
  to restore -- otherwise a failed default/max load left cudnn.benchmark/TF32 on and
  contaminated later off generations.
This commit is contained in:
Daniel Han 2026-06-28 06:15:43 +00:00
commit e102c0f18e
3 changed files with 126 additions and 60 deletions
Download patch file Download diff file Expand all files Collapse all files

106
studio/backend/core/inference/diffusion.py
View file

@ -482,55 +482,67 @@ class DiffusionBackend:
# first so unload can restore them: TF32 / cudnn.benchmark are global,
# and a later `off` load must not inherit this load's settings.
backend_flags_before = snapshot_backend_flags()
speed_applied = apply_speed_optims(
pipe,
target,
is_gguf = bool(gguf_filename),
family = fam,
speed_mode = effective_speed,
logger = logger,
)
# Quantise the dense companion text encoder(s) (opt-in fp8 / nvfp4),
# also before placement so the offload hooks move the smaller weights.
te_quant = quantize_text_encoders(
pipe,
target,
mode = text_encoder_quant,
logger = logger,
)
# apply_speed_optims mutates PROCESS-WIDE flags (TF32 / cudnn.benchmark);
# they are only restored via _LoadState.backend_flags_before on unload. If
# the build fails after this but before _state commits (e.g. an OOM in
# apply_memory_plan / pipe.to), nothing would restore them and a later `off`
# generation would be contaminated, so restore on any non-committed exit.
committed = False
try:
speed_applied = apply_speed_optims(
pipe,
target,
is_gguf = bool(gguf_filename),
family = fam,
speed_mode = effective_speed,
logger = logger,
)
# Quantise the dense companion text encoder(s) (opt-in fp8 / nvfp4),
# also before placement so the offload hooks move the smaller weights.
te_quant = quantize_text_encoders(
pipe,
target,
mode = text_encoder_quant,
logger = logger,
)
# Decide placement from MEASURED free device memory vs the model's
# estimated resident size (transformer GGUF dequantised + the
# companion text-encoder / VAE already cached for `base`), then
# apply it. Computed here, after the build but before placement,
# because the weights are still on CPU so free VRAM is the real
# budget. `cpu_offload=True` stays an explicit override.
plan = self._plan_memory(
target, gguf_path, gguf_filename, base, fam, memory_mode, cpu_offload
)
# apply_memory_plan returns the (policy, tiling) ACTUALLY engaged (it
# may fall back to whole-module offload, and tiling is a no-op on a
# pipeline with no tiling control), so status stays honest.
effective_policy, effective_tiling = apply_memory_plan(
pipe, plan, device = device, logger = logger
)
# Decide placement from MEASURED free device memory vs the model's
# estimated resident size (transformer GGUF dequantised + the
# companion text-encoder / VAE already cached for `base`), then
# apply it. Computed here, after the build but before placement,
# because the weights are still on CPU so free VRAM is the real
# budget. `cpu_offload=True` stays an explicit override.
plan = self._plan_memory(
target, gguf_path, gguf_filename, base, fam, memory_mode, cpu_offload
)
# apply_memory_plan returns the (policy, tiling) ACTUALLY engaged (it
# may fall back to whole-module offload, and tiling is a no-op on a
# pipeline with no tiling control), so status stays honest.
effective_policy, effective_tiling = apply_memory_plan(
pipe, plan, device = device, logger = logger
)
self._state = _LoadState(
pipe = pipe,
family = fam,
repo_id = repo_id,
base_repo = base,
device = device,
dtype = str(dtype).replace("torch.", ""),
cpu_offload = effective_policy != OFFLOAD_NONE,
offload_policy = effective_policy,
vae_tiling = effective_tiling,
memory_mode = plan.requested_mode,
speed_mode = effective_speed,
speed_optims = tuple(k for k, v in speed_applied.items() if v),
backend_flags_before = backend_flags_before,
text_encoder_quant = te_quant,
)
self._state = _LoadState(
pipe = pipe,
family = fam,
repo_id = repo_id,
base_repo = base,
device = device,
dtype = str(dtype).replace("torch.", ""),
cpu_offload = effective_policy != OFFLOAD_NONE,
offload_policy = effective_policy,
vae_tiling = effective_tiling,
memory_mode = plan.requested_mode,
speed_mode = effective_speed,
speed_optims = tuple(k for k, v in speed_applied.items() if v),
backend_flags_before = backend_flags_before,
text_encoder_quant = te_quant,
)
committed = True
finally:
if not committed:
restore_backend_flags(backend_flags_before)
clear_gpu_cache()
logger.info(
"diffusion.loaded: repo=%s base=%s device=%s offload=%s tiling=%s reasons=%s",

44
studio/backend/core/inference/diffusion_speed.py
View file

@ -42,31 +42,49 @@ SPEED_MODES = (SPEED_OFF, SPEED_DEFAULT, SPEED_MAX)
def snapshot_backend_flags() -> Optional[dict]:
"""Capture the process-wide torch backend flags this layer may mutate, so the
caller can restore them on unload. None if torch is unavailable."""
caller can restore them on unload. None if torch is unavailable. Each flag is read
defensively so a build/platform missing one (e.g. no cuda.matmul on CPU/MPS) still
captures the rest -- otherwise a single missing attribute would skip the whole
snapshot and a real mutated flag would leak."""
try:
import torch
return {
"matmul_tf32": bool(torch.backends.cuda.matmul.allow_tf32),
"cudnn_tf32": bool(torch.backends.cudnn.allow_tf32),
"cudnn_benchmark": bool(torch.backends.cudnn.benchmark),
}
except Exception: # noqa: BLE001 — best-effort; no snapshot -> no restore
except Exception: # noqa: BLE001 — no torch -> nothing to snapshot/restore
return None
state: dict[str, bool] = {}
matmul = getattr(getattr(torch.backends, "cuda", None), "matmul", None)
if matmul is not None and hasattr(matmul, "allow_tf32"):
state["matmul_tf32"] = bool(matmul.allow_tf32)
cudnn = getattr(torch.backends, "cudnn", None)
if cudnn is not None:
if hasattr(cudnn, "allow_tf32"):
state["cudnn_tf32"] = bool(cudnn.allow_tf32)
if hasattr(cudnn, "benchmark"):
state["cudnn_benchmark"] = bool(cudnn.benchmark)
return state
def restore_backend_flags(state: Optional[dict]) -> None:
"""Restore the flags captured by ``snapshot_backend_flags``. No-op on None."""
"""Restore the flags captured by ``snapshot_backend_flags``. No-op on None. Each
flag is restored independently so one failure can't leave the others leaked."""
if not state:
return
try:
import torch
torch.backends.cuda.matmul.allow_tf32 = state["matmul_tf32"]
torch.backends.cudnn.allow_tf32 = state["cudnn_tf32"]
torch.backends.cudnn.benchmark = state["cudnn_benchmark"]
except Exception: # noqa: BLE001 — best-effort restore
except Exception: # noqa: BLE001 — no torch -> nothing to restore
return
def _set(obj: Any, attr: str, key: str) -> None:
if obj is not None and key in state and hasattr(obj, attr):
try:
setattr(obj, attr, state[key])
except Exception: # noqa: BLE001 — best-effort per-flag restore
pass
_set(getattr(getattr(torch.backends, "cuda", None), "matmul", None), "allow_tf32", "matmul_tf32")
cudnn = getattr(torch.backends, "cudnn", None)
_set(cudnn, "allow_tf32", "cudnn_tf32")
_set(cudnn, "benchmark", "cudnn_benchmark")
def normalize_speed_mode(value: Optional[str]) -> str:
"""Lower/strip a requested speed mode (dashes ok); None / "" -> off."""

36
studio/backend/tests/test_diffusion_speed.py
View file

@ -116,6 +116,42 @@ def test_restore_backend_flags_tolerates_none():
restore_backend_flags(None) # no torch needed, no-op
def test_snapshot_partial_when_some_backends_missing(monkeypatch):
# A build/platform without cuda.matmul (e.g. CPU/MPS) must still snapshot + restore the
# flags it does have, rather than skipping the whole snapshot on one missing attribute.
torch = types.ModuleType("torch")
torch.backends = types.SimpleNamespace(
cuda = types.SimpleNamespace(), # no .matmul
cudnn = types.SimpleNamespace(benchmark = True), # no .allow_tf32
)
monkeypatch.setitem(sys.modules, "torch", torch)
snap = snapshot_backend_flags()
assert snap == {"cudnn_benchmark": True}
torch.backends.cudnn.benchmark = False
restore_backend_flags(snap)
assert torch.backends.cudnn.benchmark is True
def test_restore_is_independent_per_flag(monkeypatch):
# A read-only / failing attribute must not abort restoring the remaining flags.
torch = _stub_torch(monkeypatch)
class _NoMatmulSet:
@property
def allow_tf32(self):
return False
@allow_tf32.setter
def allow_tf32(self, value):
raise RuntimeError("read-only on this build")
torch.backends.cuda.matmul = _NoMatmulSet()
snap = {"matmul_tf32": False, "cudnn_tf32": False, "cudnn_benchmark": False}
torch.backends.cudnn.benchmark = True
restore_backend_flags(snap) # matmul setter raises, cudnn still restored
assert torch.backends.cudnn.benchmark is False
# ── applier ───────────────────────────────────────────────────────────────────