3a0bd55bdf
Add a lean, backend-agnostic memory policy that picks a CPU-offload policy and VAE tiling/slicing from measured free device memory vs the model's estimated resident footprint, then applies it to the built pipeline. auto stays resident when the model fits (byte-identical to the prior resident path), and falls to whole-module offload when tight; fast/balanced/low_vram are explicit overrides. Sequential submodule offload is unreliable for GGUF transformers on diffusers 0.38, so it falls back to whole-module offload and status reports the policy actually engaged. Verified on Z-Image-Turbo Q4_K_M (B200): auto reproduces the resident image with no VRAM/latency regression (PSNR inf); balanced/low_vram cut generation peak VRAM 47.9% (15951 -> 8318 MB) with byte-identical output, at the expected latency cost. 73 prior + 35 new CPU tests pass.
369 lines
14 KiB
Python
369 lines
14 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""Unit tests for the diffusion memory planner (``diffusion_memory.py``).
|
|
|
|
Hermetic and CPU-only: no torch, diffusers, GPU, or network. The device target
|
|
and the device-memory snapshot are constructed directly, so the planner's policy
|
|
matrix and the applier's pipeline calls are exercised in isolation.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import types
|
|
|
|
import pytest
|
|
|
|
from core.inference.diffusion_memory import (
|
|
DEFAULT_IMAGE_HEIGHT,
|
|
DEFAULT_IMAGE_WIDTH,
|
|
MEMORY_MODE_BALANCED,
|
|
MEMORY_MODE_FAST,
|
|
MEMORY_MODE_LOW_VRAM,
|
|
OFFLOAD_MODEL,
|
|
OFFLOAD_NONE,
|
|
OFFLOAD_SEQUENTIAL,
|
|
DeviceMemory,
|
|
apply_memory_plan,
|
|
estimate_gguf_dense_mib,
|
|
estimate_image_runtime_mib,
|
|
infer_gguf_quant_label,
|
|
normalize_memory_mode,
|
|
plan_diffusion_memory,
|
|
snapshot_device_memory,
|
|
)
|
|
|
|
|
|
def _target(*, device = "cuda", backend = "cuda", supports_offload = True):
|
|
"""A duck-typed stand-in for DiffusionDeviceTarget (only the fields the
|
|
planner / snapshot read)."""
|
|
return types.SimpleNamespace(
|
|
device = device,
|
|
backend = backend,
|
|
supports_model_cpu_offload = supports_offload,
|
|
)
|
|
|
|
|
|
def _discrete(free_mib, total_mib = None):
|
|
return DeviceMemory("cuda", "cuda", "discrete_vram", free_mib, total_mib or free_mib)
|
|
|
|
|
|
# ── mode normalisation ────────────────────────────────────────────────────────
|
|
|
|
|
|
def test_normalize_memory_mode_accepts_and_rejects():
|
|
assert normalize_memory_mode(None) is None
|
|
assert normalize_memory_mode(" ") is None
|
|
assert normalize_memory_mode("LOW-VRAM") == "low_vram"
|
|
assert normalize_memory_mode("Balanced") == "balanced"
|
|
with pytest.raises(ValueError):
|
|
normalize_memory_mode("ultra")
|
|
|
|
|
|
# ── filename / size estimates ─────────────────────────────────────────────────
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"filename,expected",
|
|
[
|
|
("z-image-turbo-Q4_K_M.gguf", "Q4_K_M"),
|
|
("flux1-dev-Q8_0.gguf", "Q8_0"),
|
|
("model-BF16.gguf", "BF16"),
|
|
("qwen-image-IQ4_XS.gguf", "IQ4_XS"),
|
|
("no-quant-here.gguf", None),
|
|
(None, None),
|
|
],
|
|
)
|
|
def test_infer_gguf_quant_label(filename, expected):
|
|
assert infer_gguf_quant_label(filename) == expected
|
|
|
|
|
|
def test_estimate_gguf_dense_mib_expansion():
|
|
# 4-bit roughly quadruples once dequantised to bf16; F16 is already dense.
|
|
assert estimate_gguf_dense_mib(1000, "Q4_K_M") == 4000
|
|
assert estimate_gguf_dense_mib(1000, "Q8_0") == 2000
|
|
assert estimate_gguf_dense_mib(1000, "BF16") == 1000
|
|
assert estimate_gguf_dense_mib(None, "Q4_K_M") is None
|
|
# Unknown quant falls back to the conservative 4-bit-ish factor.
|
|
assert estimate_gguf_dense_mib(1000, None) == 4000
|
|
|
|
|
|
def test_estimate_image_runtime_scales_with_pixels_and_family():
|
|
base = estimate_image_runtime_mib(width = DEFAULT_IMAGE_WIDTH, height = DEFAULT_IMAGE_HEIGHT)
|
|
bigger = estimate_image_runtime_mib(width = 2048, height = 2048)
|
|
assert bigger > base
|
|
# Distilled / turbo families get a discount.
|
|
turbo = estimate_image_runtime_mib(
|
|
width = DEFAULT_IMAGE_WIDTH, height = DEFAULT_IMAGE_HEIGHT, family = "z-image-turbo"
|
|
)
|
|
assert turbo < base
|
|
|
|
|
|
# ── planner: device classes ───────────────────────────────────────────────────
|
|
|
|
|
|
def test_cpu_target_never_offloads_but_tiles():
|
|
plan = plan_diffusion_memory(
|
|
target = _target(device = "cpu", backend = "cpu", supports_offload = False),
|
|
device_memory = DeviceMemory("cpu", "cpu", "system_memory", 8000, 16000),
|
|
model_dense_mib = 4000,
|
|
runtime_headroom_mib = 2000,
|
|
)
|
|
assert plan.offload_policy == OFFLOAD_NONE
|
|
# CPU/MPS have no separate device pool, so VAE tiling is on to cap the spike.
|
|
assert plan.vae_tiling and plan.vae_slicing
|
|
|
|
|
|
def test_mps_unified_never_auto_offloads():
|
|
plan = plan_diffusion_memory(
|
|
target = _target(device = "mps", backend = "mps", supports_offload = False),
|
|
device_memory = DeviceMemory("mps", "mps", "unified_memory", 4000, 32000),
|
|
model_dense_mib = 20000,
|
|
runtime_headroom_mib = 4000,
|
|
)
|
|
assert plan.offload_policy == OFFLOAD_NONE
|
|
assert any("unified" in r for r in plan.reasons)
|
|
|
|
|
|
def test_unified_cuda_skips_offload_even_if_offload_capable():
|
|
# An integrated CUDA SoC reports unified memory; CPU offload would free nothing.
|
|
plan = plan_diffusion_memory(
|
|
target = _target(device = "cuda", backend = "cuda", supports_offload = True),
|
|
device_memory = DeviceMemory("cuda", "cuda", "unified_memory", 2000, 16000),
|
|
model_dense_mib = 12000,
|
|
runtime_headroom_mib = 4000,
|
|
)
|
|
assert plan.offload_policy == OFFLOAD_NONE
|
|
|
|
|
|
# ── planner: auto budget tiers on a discrete GPU ──────────────────────────────
|
|
|
|
|
|
def test_auto_resident_when_roomy():
|
|
# 80 GB card, ~16 GB model: fits with headroom -> stay resident (bit-identical).
|
|
plan = plan_diffusion_memory(
|
|
target = _target(),
|
|
device_memory = _discrete(80000),
|
|
model_dense_mib = 12000,
|
|
runtime_headroom_mib = 4000,
|
|
)
|
|
assert plan.offload_policy == OFFLOAD_NONE
|
|
assert plan.vae_tiling is False and plan.vae_slicing is False # roomy -> no tiling
|
|
|
|
|
|
def test_auto_model_offload_on_tight_fit():
|
|
# 24 GB free -> reserve max(2048, 2400)=2400 -> budget 21600, 0.85*budget=18360.
|
|
# required = 16000+4000+1000 = 21000: over 0.85*budget but still under budget
|
|
# -> whole-module offload.
|
|
plan = plan_diffusion_memory(
|
|
target = _target(),
|
|
device_memory = _discrete(24000, 24000),
|
|
model_dense_mib = 16000,
|
|
runtime_headroom_mib = 4000,
|
|
base_overhead_mib = 1000,
|
|
)
|
|
assert plan.offload_policy == OFFLOAD_MODEL
|
|
assert plan.vae_tiling is True # offloading -> device is tight -> tile
|
|
|
|
|
|
def test_auto_sequential_when_model_exceeds_budget():
|
|
plan = plan_diffusion_memory(
|
|
target = _target(),
|
|
device_memory = _discrete(8000, 8000),
|
|
model_dense_mib = 40000,
|
|
runtime_headroom_mib = 4000,
|
|
)
|
|
assert plan.offload_policy == OFFLOAD_SEQUENTIAL
|
|
|
|
|
|
def test_auto_stays_resident_when_budget_unknown():
|
|
plan = plan_diffusion_memory(
|
|
target = _target(),
|
|
device_memory = _discrete(None, None),
|
|
model_dense_mib = 40000,
|
|
runtime_headroom_mib = 4000,
|
|
)
|
|
assert plan.offload_policy == OFFLOAD_NONE
|
|
assert any("unknown" in r for r in plan.reasons)
|
|
|
|
|
|
# ── planner: explicit modes + cpu_offload override ────────────────────────────
|
|
|
|
|
|
def test_explicit_modes_force_policy_regardless_of_budget():
|
|
roomy = _discrete(80000)
|
|
assert plan_diffusion_memory(
|
|
target = _target(), device_memory = roomy, model_dense_mib = 1000,
|
|
runtime_headroom_mib = 1000, requested_mode = MEMORY_MODE_FAST,
|
|
).offload_policy == OFFLOAD_NONE
|
|
assert plan_diffusion_memory(
|
|
target = _target(), device_memory = roomy, model_dense_mib = 1000,
|
|
runtime_headroom_mib = 1000, requested_mode = MEMORY_MODE_BALANCED,
|
|
).offload_policy == OFFLOAD_MODEL
|
|
assert plan_diffusion_memory(
|
|
target = _target(), device_memory = roomy, model_dense_mib = 1000,
|
|
runtime_headroom_mib = 1000, requested_mode = MEMORY_MODE_LOW_VRAM,
|
|
).offload_policy == OFFLOAD_SEQUENTIAL
|
|
|
|
|
|
def test_fast_falls_back_to_model_offload_when_it_does_not_fit():
|
|
plan = plan_diffusion_memory(
|
|
target = _target(),
|
|
device_memory = _discrete(8000, 8000),
|
|
model_dense_mib = 40000,
|
|
runtime_headroom_mib = 4000,
|
|
requested_mode = MEMORY_MODE_FAST,
|
|
)
|
|
assert plan.offload_policy == OFFLOAD_MODEL
|
|
|
|
|
|
def test_explicit_cpu_offload_overrides_resident_auto_choice():
|
|
# Roomy GPU -> auto would stay resident, but cpu_offload=True forces offload.
|
|
plan = plan_diffusion_memory(
|
|
target = _target(),
|
|
device_memory = _discrete(80000),
|
|
model_dense_mib = 4000,
|
|
runtime_headroom_mib = 2000,
|
|
explicit_offload = True,
|
|
)
|
|
assert plan.offload_policy == OFFLOAD_MODEL
|
|
assert any("explicit cpu_offload" in r for r in plan.reasons)
|
|
|
|
|
|
def test_explicit_cpu_offload_ignored_on_cpu_target():
|
|
plan = plan_diffusion_memory(
|
|
target = _target(device = "cpu", backend = "cpu", supports_offload = False),
|
|
device_memory = DeviceMemory("cpu", "cpu", "system_memory", 8000, 16000),
|
|
model_dense_mib = 4000,
|
|
runtime_headroom_mib = 2000,
|
|
explicit_offload = True,
|
|
)
|
|
assert plan.offload_policy == OFFLOAD_NONE
|
|
|
|
|
|
# ── snapshot ──────────────────────────────────────────────────────────────────
|
|
|
|
|
|
def test_snapshot_cpu_target_uses_system_memory(monkeypatch):
|
|
import core.inference.diffusion_memory as mem
|
|
|
|
monkeypatch.setattr(mem, "_system_memory_mib", lambda: (16000, 9000))
|
|
snap = snapshot_device_memory(_target(device = "cpu", backend = "cpu"))
|
|
assert snap.memory_kind == "system_memory"
|
|
assert snap.free_mib == 9000 and snap.total_mib == 16000
|
|
|
|
|
|
def test_snapshot_cuda_reads_mem_get_info(monkeypatch):
|
|
import sys
|
|
|
|
fake_torch = types.ModuleType("torch")
|
|
fake_torch.cuda = types.SimpleNamespace(
|
|
mem_get_info = lambda: (10 * 1024 * 1024 * 1024, 24 * 1024 * 1024 * 1024),
|
|
get_device_properties = lambda i: types.SimpleNamespace(integrated = False),
|
|
)
|
|
monkeypatch.setitem(sys.modules, "torch", fake_torch)
|
|
snap = snapshot_device_memory(_target())
|
|
assert snap.memory_kind == "discrete_vram"
|
|
assert snap.free_mib == 10 * 1024 and snap.total_mib == 24 * 1024
|
|
|
|
|
|
def test_snapshot_never_raises_on_probe_failure(monkeypatch):
|
|
import sys
|
|
|
|
fake_torch = types.ModuleType("torch")
|
|
|
|
def _boom():
|
|
raise RuntimeError("no cuda")
|
|
|
|
fake_torch.cuda = types.SimpleNamespace(mem_get_info = _boom)
|
|
monkeypatch.setitem(sys.modules, "torch", fake_torch)
|
|
snap = snapshot_device_memory(_target())
|
|
assert snap.free_mib is None and snap.total_mib is None
|
|
|
|
|
|
# ── applier ───────────────────────────────────────────────────────────────────
|
|
|
|
|
|
class _RecordingPipe:
|
|
def __init__(self) -> None:
|
|
self.calls: list[str] = []
|
|
|
|
def to(self, device):
|
|
self.calls.append(f"to:{device}")
|
|
return self
|
|
|
|
def enable_model_cpu_offload(self):
|
|
self.calls.append("model_offload")
|
|
|
|
def enable_sequential_cpu_offload(self):
|
|
self.calls.append("sequential_offload")
|
|
|
|
def enable_vae_tiling(self):
|
|
self.calls.append("vae_tiling")
|
|
|
|
def enable_vae_slicing(self):
|
|
self.calls.append("vae_slicing")
|
|
|
|
|
|
def _plan(policy, *, tiling):
|
|
return plan_diffusion_memory(
|
|
target = _target(),
|
|
device_memory = _discrete(80000) if policy == OFFLOAD_NONE else _discrete(4000, 8000),
|
|
model_dense_mib = 1000 if policy == OFFLOAD_NONE else 40000,
|
|
runtime_headroom_mib = 1000,
|
|
requested_mode = {
|
|
OFFLOAD_NONE: MEMORY_MODE_FAST,
|
|
OFFLOAD_MODEL: MEMORY_MODE_BALANCED,
|
|
OFFLOAD_SEQUENTIAL: MEMORY_MODE_LOW_VRAM,
|
|
}[policy],
|
|
)
|
|
|
|
|
|
def test_apply_none_places_resident():
|
|
pipe = _RecordingPipe()
|
|
effective = apply_memory_plan(pipe, _plan(OFFLOAD_NONE, tiling = False), device = "cuda")
|
|
assert pipe.calls == ["to:cuda"] # no tiling on a roomy resident run
|
|
assert effective == OFFLOAD_NONE
|
|
|
|
|
|
def test_apply_model_offload_engages_offload_and_tiling():
|
|
pipe = _RecordingPipe()
|
|
effective = apply_memory_plan(pipe, _plan(OFFLOAD_MODEL, tiling = True), device = "cuda")
|
|
assert "model_offload" in pipe.calls
|
|
assert "to:cuda" not in pipe.calls # offload owns placement; never both
|
|
assert "vae_tiling" in pipe.calls and "vae_slicing" in pipe.calls
|
|
assert effective == OFFLOAD_MODEL
|
|
|
|
|
|
def test_apply_sequential_offload():
|
|
pipe = _RecordingPipe()
|
|
effective = apply_memory_plan(pipe, _plan(OFFLOAD_SEQUENTIAL, tiling = True), device = "cuda")
|
|
assert "sequential_offload" in pipe.calls and "to:cuda" not in pipe.calls
|
|
assert effective == OFFLOAD_SEQUENTIAL
|
|
|
|
|
|
def test_apply_sequential_falls_back_to_model_offload_when_unsupported():
|
|
# Sequential offload is unreliable for GGUF on some diffusers versions; the
|
|
# applier must fall back to whole-module offload and report what actually ran.
|
|
class _NoSeqPipe(_RecordingPipe):
|
|
def enable_sequential_cpu_offload(self):
|
|
raise RuntimeError("sequential offload not supported for this transformer")
|
|
|
|
pipe = _NoSeqPipe()
|
|
effective = apply_memory_plan(pipe, _plan(OFFLOAD_SEQUENTIAL, tiling = True), device = "cuda")
|
|
assert effective == OFFLOAD_MODEL
|
|
assert "model_offload" in pipe.calls
|
|
|
|
|
|
def test_apply_tolerates_pipe_without_vae_savers():
|
|
# A pipeline missing enable_vae_* must not crash the applier.
|
|
class _Bare:
|
|
def __init__(self):
|
|
self.moved = None
|
|
|
|
def to(self, device):
|
|
self.moved = device
|
|
|
|
bare = _Bare()
|
|
apply_memory_plan(bare, _plan(OFFLOAD_NONE, tiling = False), device = "cpu")
|
|
assert bare.moved == "cpu"
|