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

Serialize video-load GPU placement with eviction; skip LoRA adapters as base picks

Browse source 
Two evict/OOM fixes on the diffusion load paths:

- The video load moved a pipeline onto the GPU (apply_memory_plan) and
  committed it while holding no lock, so an unload / GPU-arbiter eviction --
  which bumps the load token and then barriers on _generate_lock before
  freeing -- could hand VIDEO to chat/images and let the new owner allocate
  concurrently with the in-flight placement, OOMing. Hold _generate_lock
  across placement + the locked commit, mirroring the image backend, so an
  evicting owner waits until this worker's placement is torn down or
  committed. Lock order stays _generate_lock -> _lock (unload takes _lock
  then releases it before the barrier), so there is no deadlock.

- resolve_local_single_file reinterpreted an On-Device folder as a base
  single_file load whenever it held exactly one .safetensors, so a PEFT LoRA
  adapter folder (adapter_config.json + adapter_model.safetensors) with a
  family-token name was picked as a base checkpoint, evicting the resident
  model before from_single_file failed on the adapter weights. Skip adapter
  folders (adapter_config.json) and the adapter_model basename so the pick
  stays a pipeline load and 400s in validation, before the GPU handoff.

Adds regression tests for both.
This commit is contained in:
Daniel Han 2026-07-07 20:57:44 +00:00
commit cf0f5d5504
4 changed files with 185 additions and 102 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -159,13 +159,24 @@ def resolve_local_single_file(model_path: str) -> Optional[str]:
advertised model is unusable. The images load route uses this to reinterpret such a pick as a
``single_file`` load of the sole checkpoint. A real pipeline dir (has ``model_index.json``) or
an ambiguous one (0 or more than 1 ``.safetensors``, e.g. a sharded pipeline) returns None and
loads unchanged. Never raises."""
loads unchanged. A PEFT LoRA adapter folder is also skipped (see below). Never raises."""
try:
root = Path(model_path).expanduser()
if not root.is_dir() or (root / "model_index.json").is_file():
return None
# A PEFT LoRA adapter folder (adapter_config.json + adapter_model.safetensors) is not a
# base checkpoint: from_single_file would fail on the adapter weights AFTER the route
# evicted the resident GPU model. Skip it so the pick stays a pipeline load and 400s in
# validation, before the GPU handoff. Also drop a bare adapter_model.safetensors so a
# config-less adapter export is never reinterpreted as the sole checkpoint.
if (root / "adapter_config.json").is_file():
return None
checkpoints = [
p.name for p in root.iterdir() if p.is_file() and p.suffix.lower() == ".safetensors"
p.name
for p in root.iterdir()
if p.is_file()
and p.suffix.lower() == ".safetensors"
and p.stem.lower() != "adapter_model"
]
except OSError:
return None

201
studio/backend/core/inference/video.py
View file

@ -1291,111 +1291,112 @@ class VideoBackend:
if view is pipe:
attention_engaged = engaged
speed_optims = tuple(k for k, v in applied.items() if v)
# A cancelled/superseded load must not place weights on the GPU the arbiter
# may already have handed to another backend; recheck right before placement
# (the commit below still does the final locked check).
if _load_token is not None and _load_token != self._load_token:
del pipe
clear_gpu_cache()
raise RuntimeError("Video load was cancelled or superseded.")
offload_policy, vae_tiling = apply_memory_plan(pipe, plan, device = device, logger = logger)
# A dual-DiT MoE pipe (Wan2.2-A14B) needs no extra per-expert offload pass here:
# apply_memory_plan's group tier (_apply_group_offload) already block-streams every
# DiT it finds on the pipe -- transformer AND transformer_2 -- and model/sequential
# offload hook every top-level module, so the second expert is covered under all tiers.
# A second _apply_group_offload on transformer_2 would re-register the group-offload
# hooks it already carries, which diffusers rejects with a duplicate-hook ValueError.
if not vae_tiling:
# Decode of a whole clip is the video memory peak; tiling is near-free
# in quality and keeps the decode bounded, so it is always on.
try:
pipe.vae.enable_tiling()
vae_tiling = True
except Exception as exc: # noqa: BLE001 -- tiling is an optimisation only
logger.warning("video.vae_tiling_failed: %s", exc)
resolved = build_resolved_record(
{
"memory_mode": (
memory_mode,
plan.requested_mode,
f"planned '{plan.offload_policy}' offload from the family size table",
),
"speed_mode": (
speed_mode,
effective_speed,
"quantized transformer requires compile"
if transformer_quant_engaged is not None
else "clip denoises amortise the one-time compile within a single run"
if speed_mode is None
else "requested",
),
"attention_backend": (
attention_backend,
attention_engaged or "native",
"cuDNN fused attention on NVIDIA when a speed profile is active",
),
"transformer_cache": (
None if cache_auto else transformer_cache,
cache_engaged or "off",
cache_reason,
),
"transformer_quant": (
transformer_quant,
transformer_quant_engaged or "off",
"dense DiT(s) torchao-quantised onto the low-precision tensor cores"
if transformer_quant_engaged is not None
else (
"skipped: offload moves the DiT, unsupported for torchao "
"tensors; pin a resident memory mode to combine them"
if quant_skipped_for_offload
else "not engaged (dense bf16 DiT loaded)"
),
),
"text_encoder_quant": (
text_encoder_quant,
text_encoder_quant_engaged or "off",
"dense text encoder quantised in place"
if text_encoder_quant_engaged is not None
else "not engaged (dense bf16 text encoder loaded)",
),
}
)
with self._lock:
with self._generate_lock:
# A cancelled/superseded load must not place weights on the GPU the arbiter
# may already have handed to another backend; recheck right before placement
# (the commit below still does the final locked check).
if _load_token is not None and _load_token != self._load_token:
del pipe
clear_gpu_cache()
raise RuntimeError("Video load was cancelled or superseded.")
self._state = _VideoLoadState(
pipe = pipe,
family = fam,
repo_id = repo_id,
base_repo = base,
device = device,
dtype = str(dtype).replace("torch.", ""),
kind = kind,
gguf_filename = gguf_filename,
offload_policy = offload_policy,
vae_tiling = vae_tiling,
memory_mode = plan.requested_mode,
speed_mode = effective_speed,
# Already filtered above to only the optimisations that engaged;
# apply_speed_optims returns every flag True/False and the view
# loop keeps just the True names.
speed_optims = speed_optims,
backend_flags = backend_flags,
attention_backend = attention_engaged,
transformer_cache = cache_engaged,
cache_auto = cache_may_toggle,
cache_quant_active = cache_quant_active,
cache_threshold = transformer_cache_threshold,
transformer_quant = transformer_quant_engaged,
text_encoder_quant = text_encoder_quant_engaged,
resolved = resolved,
offload_policy, vae_tiling = apply_memory_plan(pipe, plan, device = device, logger = logger)
# A dual-DiT MoE pipe (Wan2.2-A14B) needs no extra per-expert offload pass here:
# apply_memory_plan's group tier (_apply_group_offload) already block-streams every
# DiT it finds on the pipe -- transformer AND transformer_2 -- and model/sequential
# offload hook every top-level module, so the second expert is covered under all tiers.
# A second _apply_group_offload on transformer_2 would re-register the group-offload
# hooks it already carries, which diffusers rejects with a duplicate-hook ValueError.
if not vae_tiling:
# Decode of a whole clip is the video memory peak; tiling is near-free
# in quality and keeps the decode bounded, so it is always on.
try:
pipe.vae.enable_tiling()
vae_tiling = True
except Exception as exc: # noqa: BLE001 -- tiling is an optimisation only
logger.warning("video.vae_tiling_failed: %s", exc)
resolved = build_resolved_record(
{
"memory_mode": (
memory_mode,
plan.requested_mode,
f"planned '{plan.offload_policy}' offload from the family size table",
),
"speed_mode": (
speed_mode,
effective_speed,
"quantized transformer requires compile"
if transformer_quant_engaged is not None
else "clip denoises amortise the one-time compile within a single run"
if speed_mode is None
else "requested",
),
"attention_backend": (
attention_backend,
attention_engaged or "native",
"cuDNN fused attention on NVIDIA when a speed profile is active",
),
"transformer_cache": (
None if cache_auto else transformer_cache,
cache_engaged or "off",
cache_reason,
),
"transformer_quant": (
transformer_quant,
transformer_quant_engaged or "off",
"dense DiT(s) torchao-quantised onto the low-precision tensor cores"
if transformer_quant_engaged is not None
else (
"skipped: offload moves the DiT, unsupported for torchao "
"tensors; pin a resident memory mode to combine them"
if quant_skipped_for_offload
else "not engaged (dense bf16 DiT loaded)"
),
),
"text_encoder_quant": (
text_encoder_quant,
text_encoder_quant_engaged or "off",
"dense text encoder quantised in place"
if text_encoder_quant_engaged is not None
else "not engaged (dense bf16 text encoder loaded)",
),
}
)
# Ownership of the globals transferred to _state / _teardown_state.
self._precommit_globals = None
with self._lock:
if _load_token is not None and _load_token != self._load_token:
del pipe
clear_gpu_cache()
raise RuntimeError("Video load was cancelled or superseded.")
self._state = _VideoLoadState(
pipe = pipe,
family = fam,
repo_id = repo_id,
base_repo = base,
device = device,
dtype = str(dtype).replace("torch.", ""),
kind = kind,
gguf_filename = gguf_filename,
offload_policy = offload_policy,
vae_tiling = vae_tiling,
memory_mode = plan.requested_mode,
speed_mode = effective_speed,
# Already filtered above to only the optimisations that engaged;
# apply_speed_optims returns every flag True/False and the view
# loop keeps just the True names.
speed_optims = speed_optims,
backend_flags = backend_flags,
attention_backend = attention_engaged,
transformer_cache = cache_engaged,
cache_auto = cache_may_toggle,
cache_quant_active = cache_quant_active,
cache_threshold = transformer_cache_threshold,
transformer_quant = transformer_quant_engaged,
text_encoder_quant = text_encoder_quant_engaged,
resolved = resolved,
)
# Ownership of the globals transferred to _state / _teardown_state.
self._precommit_globals = None
logger.info(
"video.loaded: %s (%s, %s, offload=%s, speed=%s, quant=%s)",
repo_id,

15
studio/backend/tests/test_diffusion_backend.py
View file

@ -1272,6 +1272,21 @@ def test_resolve_local_single_file(tmp_path):
# A remote repo id (not a local dir) -> None.
assert resolve_local_single_file("unsloth/Qwen-Image-2512-GGUF") is None
# A PEFT LoRA adapter folder (adapter_config.json + adapter_model.safetensors), even with a
# family-token name, is NOT a base checkpoint: from_single_file would fail on the adapter
# weights AFTER the route evicted the resident GPU model, so it must not be reinterpreted as a
# single_file pick -> None (the pipeline pick then 400s in validation, before the handoff).
adapter = tmp_path / "flux-style-lora"
adapter.mkdir()
(adapter / "adapter_config.json").write_text("{}")
(adapter / "adapter_model.safetensors").write_bytes(b"w")
assert resolve_local_single_file(str(adapter)) is None
# A bare adapter_model.safetensors (no config) is likewise not treated as the sole checkpoint.
adapter2 = tmp_path / "z-image-lora"
adapter2.mkdir()
(adapter2 / "adapter_model.safetensors").write_bytes(b"w")
assert resolve_local_single_file(str(adapter2)) is None
def test_resolve_base_repo_drops_untrusted_card_tag(monkeypatch):
# When no base_repo is passed, the base is resolved from the GGUF repo's base_model card

56
studio/backend/tests/test_video_backend.py
View file

@ -699,6 +699,62 @@ def test_load_generate_unload_gguf(fake_runtime, tmp_path):
assert status["loaded"] is False
def test_load_holds_generate_lock_across_placement(fake_runtime, tmp_path, monkeypatch):
# The video load must hold _generate_lock across GPU placement (apply_memory_plan) so an
# unload / arbiter eviction -- which barriers on _generate_lock before freeing -- cannot hand
# the GPU to another backend while a multi-GB pipeline is still being moved onto it (mirrors
# the image backend, which places + commits under this lock). Verify unload() blocks until
# placement releases the lock, and the superseded load then aborts without committing.
import threading
from core.inference import video as video_mod
backend = VideoBackend()
placement_started = threading.Event()
release_placement = threading.Event()
real_apply = video_mod.apply_memory_plan
def blocking_apply(pipe, plan, **kw):
placement_started.set()
assert release_placement.wait(timeout = 5), "test placement barrier never released"
return real_apply(pipe, plan, **kw)
monkeypatch.setattr(video_mod, "apply_memory_plan", blocking_apply)
load_exc = []
def do_load():
try:
_load_gguf(backend, tmp_path)
except Exception as e: # noqa: BLE001 -- the concurrent unload supersedes this load
load_exc.append(e)
load_thread = threading.Thread(target = do_load)
load_thread.start()
assert placement_started.wait(timeout = 5), "load never reached placement"
# Placement is in flight, holding _generate_lock. unload() must block on its barrier.
unload_done = []
def do_unload():
backend.unload()
unload_done.append(True)
unload_thread = threading.Thread(target = do_unload)
unload_thread.start()
unload_thread.join(timeout = 0.5)
assert not unload_done, "unload() returned while placement still held _generate_lock (the race)"
# Release placement; unload()'s barrier then passes and its teardown runs strictly AFTER
# the load's placement+commit -- never concurrently -- so no two pipelines are ever resident.
release_placement.set()
unload_thread.join(timeout = 5)
load_thread.join(timeout = 5)
assert unload_done, "unload() did not complete after placement released _generate_lock"
assert not load_thread.is_alive() and not load_exc
assert backend._state is None # unload's teardown ran after the load, leaving nothing resident
def test_load_records_engaged_speed_optims(fake_runtime, tmp_path, monkeypatch):
# Regression: the load tail once re-ran the already-filtered speed_optims
# tuple through ``.items()`` as if it were still the raw applied dict, so