perf(image): compile numeric parity, cache-hook compile arming, FBCache toggle crash fix, TE fp8 zero-row guard
Applies the video round-2 accuracy findings to the image diffusion stack and fixes two real image-path bugs found while measuring. All numbers B200, production settings (family default steps/guidance, 1024px, seed 42, 4 fixed prompts), LPIPS (AlexNet) via the new scripts/image_speedmem_bench.py, which drives the production lever functions in the loader's own order. - inductor precision parity: emulate_precision_casts=True on the regional-compile path (fused pointwise kernels keep fp32 intermediates where eager rounds to bf16 between ops). Pairwise LPIPS of the compiled tier vs the same-stack eager tier: Qwen-Image 0.019 to 0.006 at identical speed (72.4 vs 72.5 ms/step), FLUX.1-dev 0.046 to 0.029 at +2% step time (69.8 vs 68.3, reproduced), FLUX.2-klein-4B 0.018 to 0.017 at identical speed. Snapshot/restored with the other process-wide backend flags so an off load never inherits it. - cache x compile composition: re-point each cache hook's fn_ref.original_forward at a torch.compile'd wrapper of the same bound method (armed only where the speed layer compiled the block; restored before every disable_cache and before the partial-hook cleanup). Qwen-Image FBCache computed steps 91.8 to 71.2 ms (back at the uncached compiled rate), 1.21x end to end (7.36 to 6.06 s per 4 images); FLUX.1-dev already traced through its FBCache hook and is measured neutral (same-process armed vs unarmed latents bit-identical). Skip counts within noise (13 vs 11 of 76; pairwise LPIPS 0.005). - FBCache mid-session toggle crash: diffusers 0.39 caches the HookRegistry child list on first cache_context use, so an uncached generation followed by a 20+-step generation (the auto toggle path) enabled hooks the context never reached and crashed with "No context is set" (reproduced live on FLUX.1-dev). Invalidate the stale child cache after every enable_cache. - TE fp8_dynamic zero-row guard: torchao per-row fp8 derives a per-output-channel scale from the row amax, so an all-zero weight row is 0/0 = NaN. SDXL's text_encoder_2 (OpenCLIP bigG) ships exactly such a row, and every explicit fp8_dynamic SDXL render came out black; keep zero-row Linears dense (LPIPS 0.976 black to 0.096 working). Other families' encoders have no such rows and are byte-identical. - No AUTO TE quant exists on the image branch (text_encoder_quant defaults dense, explicit-only), so the video round's auto-dense retune has no image analogue; the explicit lever's cost is now measured (TE fp8_dynamic alone, LPIPS vs bit-exact: Qwen-Image 0.038, FLUX.1-dev 0.084, SDXL 0.096; no speed win, VRAM -6.5 GB on Qwen-Image) for the docs. Tests: 96 passing across the cache/speed/precision suites (11 new arming, 2 child-registry, 2 zero-row, 4 inductor-flag); ruff clean.
This commit is contained in:
Daniel Han
parent
ec90b8658d
commit
de2f22df2b
7 changed files with 997 additions and 2 deletions
|
|
@ -64,6 +64,126 @@ def normalize_transformer_cache(value: Optional[str]) -> Optional[str]:
|
|||
return normalized
|
||||
|
||||
|
||||
def _invalidate_child_registry_cache(transformer: Any) -> None:
|
||||
"""Drop the HookRegistry's cached child-registry list after (un)installing hooks.
|
||||
|
||||
``cache_context`` propagates the state context through ``_get_child_registries``,
|
||||
which diffusers 0.39 caches on first use. An UNCACHED generation already calls
|
||||
``cache_context`` (the pipeline wraps every denoise call), creating the
|
||||
transformer-level registry with an EMPTY cached child list -- so a later
|
||||
``enable_cache`` (the auto step-count toggle engaging FBCache mid-session) installs
|
||||
block hooks that ``_set_context`` never reaches, and the first cached forward dies
|
||||
with "No context is set". Invalidate the stale cache so the next ``cache_context``
|
||||
rebuilds it over the freshly hooked blocks. Best-effort and cheap (one attribute)."""
|
||||
registry = getattr(transformer, "_diffusers_hook", None)
|
||||
if registry is not None and getattr(registry, "_child_registries_cache", None) is not None:
|
||||
try:
|
||||
registry._child_registries_cache = None
|
||||
except Exception: # noqa: BLE001 -- diffusers internals moved; leave as-is
|
||||
pass
|
||||
|
||||
|
||||
# diffusers' cache hook registry names whose compute branch we re-point at a compiled
|
||||
# inner forward (leader = the measuring first block, block = the remaining ones); both
|
||||
# hook families share the fn_ref layout.
|
||||
_CACHE_HOOK_NAMES = (
|
||||
"mag_cache_leader_block_hook",
|
||||
"mag_cache_block_hook",
|
||||
"fbc_leader_block_hook",
|
||||
"fbc_block_hook",
|
||||
)
|
||||
|
||||
|
||||
def _compile_hooked_block_inners(transformer: Any, logger: Any = None) -> int:
|
||||
"""Restore the regional compile on cache-hooked blocks' COMPUTED steps.
|
||||
|
||||
``enable_cache`` replaces each block's ``forward`` with the hook's ``new_forward``
|
||||
(stashing the pre-hook bound method in ``fn_ref.original_forward``), whose skip
|
||||
decision is data-dependent Python: MagCache ``@torch.compiler.disable``s the whole
|
||||
``new_forward`` (recursive -- the compute branch runs EAGER), and even FBCache's
|
||||
traceable ``new_forward`` graph-breaks around its disabled threshold decision,
|
||||
which on some archs (measured: Qwen-Image) drops the compute branch's call into
|
||||
``original_forward`` out of the compiled region -- the block's regional compile
|
||||
artifact (``_compiled_call_impl``) is never reached and the cache forfeits the
|
||||
compile win on every non-skipped step. An explicitly ``torch.compile``d callable
|
||||
re-enables
|
||||
dynamo for its own extent even inside a disabled frame, so re-pointing
|
||||
``fn_ref.original_forward`` at a compiled wrapper of the same bound method restores
|
||||
compiled compute steps while the skip decision stays eager exactly as designed.
|
||||
Measured (B200, scripts/image_speedmem_bench.py): Qwen-Image FBCache computed steps
|
||||
91.8 -> 71.2 ms (= the uncached compiled rate), 1.21x end to end; FLUX.1-dev is
|
||||
neutral (its FBCache ``new_forward`` happens to trace, so computed steps were
|
||||
already compiled -- same-process armed vs unarmed latents bit-identical); on the
|
||||
video DiT balanced MagCache went 39.4 -> 26.9 s at 50 steps.
|
||||
|
||||
Only blocks the speed layer actually compiled are armed (``_compiled_call_impl``
|
||||
guard -- eager tiers stay untouched), and only when ``original_forward`` is a plain
|
||||
bound method (a stacked hook chain, e.g. offload, captures a partial and is
|
||||
skipped). Idempotent via the ``_unsloth_orig_inner`` marker; best-effort. Returns
|
||||
the number of hooks armed."""
|
||||
try:
|
||||
import torch
|
||||
except Exception: # noqa: BLE001 -- no torch, nothing to arm
|
||||
return 0
|
||||
armed = 0
|
||||
try:
|
||||
for module in transformer.modules():
|
||||
registry = getattr(module, "_diffusers_hook", None)
|
||||
if registry is None or getattr(module, "_compiled_call_impl", None) is None:
|
||||
continue
|
||||
hooks = getattr(registry, "hooks", None) or {}
|
||||
for name in _CACHE_HOOK_NAMES:
|
||||
hook = hooks.get(name)
|
||||
fn_ref = getattr(hook, "fn_ref", None) if hook is not None else None
|
||||
orig = getattr(fn_ref, "original_forward", None)
|
||||
if orig is None or getattr(hook, "_unsloth_orig_inner", None) is not None:
|
||||
continue
|
||||
if getattr(orig, "__self__", None) is None:
|
||||
continue # not the plain bound method; arming would miss the block
|
||||
# fullgraph=False / dynamic=True: a cache is active by definition (its
|
||||
# decision points graph-break) and this matches the default tier the
|
||||
# regional compile used. Dynamo caches per code object, so re-arming
|
||||
# after a toggle is effectively free (~0.03 s).
|
||||
fn_ref.original_forward = torch.compile(orig, fullgraph = False, dynamic = True)
|
||||
hook._unsloth_orig_inner = orig
|
||||
armed += 1
|
||||
except Exception as exc: # noqa: BLE001 -- best-effort: the cache still works eager
|
||||
_warn(logger, "cache-hook inner compile", exc)
|
||||
return armed
|
||||
if armed and logger is not None:
|
||||
logger.info(
|
||||
"diffusion.cache: %d cache-hooked block(s) armed with compiled inner forwards",
|
||||
armed,
|
||||
)
|
||||
return armed
|
||||
|
||||
|
||||
def _restore_hooked_block_inners(transformer: Any) -> None:
|
||||
"""Undo ``_compile_hooked_block_inners``: put the plain bound methods back and clear
|
||||
the markers. MUST run before ``disable_cache`` -- ``remove_hook`` splices
|
||||
``fn_ref.original_forward`` back into ``module.forward``, and leaving the compiled
|
||||
wrapper there would pin a stale compiled callable onto the uncached path."""
|
||||
try:
|
||||
modules = list(transformer.modules())
|
||||
except Exception: # noqa: BLE001 -- not a torch module (tests/fakes): nothing armed
|
||||
return
|
||||
for module in modules:
|
||||
registry = getattr(module, "_diffusers_hook", None)
|
||||
if registry is None:
|
||||
continue
|
||||
hooks = getattr(registry, "hooks", None) or {}
|
||||
for name in _CACHE_HOOK_NAMES:
|
||||
hook = hooks.get(name)
|
||||
orig = getattr(hook, "_unsloth_orig_inner", None) if hook is not None else None
|
||||
if orig is None:
|
||||
continue
|
||||
try:
|
||||
hook.fn_ref.original_forward = orig
|
||||
hook._unsloth_orig_inner = None
|
||||
except Exception: # noqa: BLE001 -- per-hook best-effort
|
||||
pass
|
||||
|
||||
|
||||
def _pipeline_opens_cache_context(pipe: Any) -> bool:
|
||||
"""Whether the pipeline enters ``transformer.cache_context(...)`` in its denoise loop.
|
||||
The First-Block-Cache hook requires it at run time, and a CacheMixin transformer alone
|
||||
|
|
@ -137,6 +257,15 @@ def apply_step_cache(
|
|||
|
||||
config = FirstBlockCacheConfig(threshold = thr)
|
||||
enable_cache(config)
|
||||
# enable_cache AFTER the pipe has already run leaves a stale cached child-registry
|
||||
# list on the transformer's HookRegistry; the block hooks just installed would then
|
||||
# never receive the cache context. Must follow every enable_cache.
|
||||
_invalidate_child_registry_cache(transformer)
|
||||
# If the blocks are already regionally compiled (the generation-time toggle
|
||||
# path: compile ran at load), re-point the fresh hooks' compute branch at
|
||||
# compiled inners; the load path (cache before compile) is armed by
|
||||
# _compile_repeated_blocks instead. No-op when nothing is compiled.
|
||||
_compile_hooked_block_inners(transformer, logger)
|
||||
try:
|
||||
transformer._unsloth_step_cache = f"{mode}@{thr}"
|
||||
except Exception: # noqa: BLE001 — marker is best-effort
|
||||
|
|
@ -146,7 +275,10 @@ def apply_step_cache(
|
|||
return mode
|
||||
except Exception as exc: # noqa: BLE001 — incompatible model -> run uncached
|
||||
# enable_cache can fail after hooking some blocks; drop any partial hooks so
|
||||
# the reported-uncached model doesn't actually run half-cached.
|
||||
# the reported-uncached model doesn't actually run half-cached. Any armed
|
||||
# compiled inners must be restored FIRST (remove_hook splices original_forward
|
||||
# back into module.forward).
|
||||
_restore_hooked_block_inners(transformer)
|
||||
try:
|
||||
transformer.disable_cache()
|
||||
except Exception: # noqa: BLE001
|
||||
|
|
@ -225,6 +357,10 @@ def maybe_toggle_step_cache(
|
|||
disable_cache = getattr(transformer, "disable_cache", None)
|
||||
if callable(disable_cache):
|
||||
try:
|
||||
# Before remove_hook splices fn_ref.original_forward back into
|
||||
# module.forward: the compiled inner wrappers must not leak onto the
|
||||
# uncached path.
|
||||
_restore_hooked_block_inners(transformer)
|
||||
disable_cache()
|
||||
transformer._unsloth_step_cache = None
|
||||
if logger is not None:
|
||||
|
|
|
|||
|
|
@ -217,6 +217,24 @@ def _cast_int8_selective(encoder: Any, target: Any, skip_first: int, skip_last:
|
|||
quantize_(encoder, _make_quant_config(TQ_INT8), filter_fn = filter_fn)
|
||||
|
||||
|
||||
def _weight_has_zero_output_row(module: Any) -> bool:
|
||||
"""True when a Linear's weight contains an all-zero OUTPUT row. torchao's per-row
|
||||
fp8 scheme derives a per-output-channel scale from that row's amax, so a dead row
|
||||
yields scale 0 -> 0/0 = NaN through the whole forward. Real checkpoints ship such
|
||||
rows: SDXL's text_encoder_2 (OpenCLIP ViT-bigG) has one in
|
||||
``text_model.encoder.layers.2.self_attn.out_proj`` -- measured on B200: every
|
||||
fp8_dynamic SDXL render came out black (NaN embeddings) until this Linear is left
|
||||
dense. Cheap (one amax per Linear, once per load); False on any error so the
|
||||
caster's own failure handling stays in charge."""
|
||||
try:
|
||||
weight = getattr(module, "weight", None)
|
||||
if weight is None or weight.ndim != 2:
|
||||
return False
|
||||
return bool((weight.abs().amax(dim = -1) == 0).any().item())
|
||||
except Exception: # noqa: BLE001 -- unreadable weight: let quantize_ decide
|
||||
return False
|
||||
|
||||
|
||||
def _cast_fp8_dynamic(encoder: Any, target: Any) -> None:
|
||||
# torchao dynamic fp8 COMPUTE, per-row (per-token activation + per-output-channel weight ->
|
||||
# torch._scaled_mm on the fp8 tensor cores). Unlike the layerwise `fp8` backend this keeps the
|
||||
|
|
@ -232,9 +250,15 @@ def _cast_fp8_dynamic(encoder: Any, target: Any) -> None:
|
|||
|
||||
# require_bf16: scaled_mm asserts a bf16 weight, so skip any stray non-bf16 Linear the encoder
|
||||
# keeps (belt-and-suspenders over the named T5 wo exclusion) rather than aborting the pass.
|
||||
filter_fn = make_filter_fn(
|
||||
base = make_filter_fn(
|
||||
DEFAULT_MIN_LINEAR_FEATURES, _te_exclude_tokens(encoder), require_bf16 = True
|
||||
)
|
||||
|
||||
# A Linear with an all-zero output row NaNs under per-row scaling (scale 0 -> 0/0);
|
||||
# keep exactly those Linears dense so one dead row cannot black out every render.
|
||||
def filter_fn(module: Any, fqn: str = "") -> bool:
|
||||
return base(module, fqn) and not _weight_has_zero_output_row(module)
|
||||
|
||||
quantize_(encoder, _make_quant_config(TQ_FP8), filter_fn = filter_fn)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -77,6 +77,9 @@ def snapshot_backend_flags() -> Optional[dict]:
|
|||
state["cudnn_tf32"] = bool(cudnn.allow_tf32)
|
||||
if hasattr(cudnn, "benchmark"):
|
||||
state["cudnn_benchmark"] = bool(cudnn.benchmark)
|
||||
inductor_cfg = _inductor_config()
|
||||
if inductor_cfg is not None and hasattr(inductor_cfg, "emulate_precision_casts"):
|
||||
state["inductor_emulate_precision_casts"] = bool(inductor_cfg.emulate_precision_casts)
|
||||
return state
|
||||
|
||||
|
||||
|
|
@ -103,6 +106,19 @@ def restore_backend_flags(state: Optional[dict]) -> None:
|
|||
cudnn = getattr(torch.backends, "cudnn", None)
|
||||
_set(cudnn, "allow_tf32", "cudnn_tf32")
|
||||
_set(cudnn, "benchmark", "cudnn_benchmark")
|
||||
_set(_inductor_config(), "emulate_precision_casts", "inductor_emulate_precision_casts")
|
||||
|
||||
|
||||
def _inductor_config() -> Any:
|
||||
"""``torch._inductor.config`` or None. Resolved as attributes off the imported torch
|
||||
module (real torch exposes ``_inductor`` directly after ``import torch``) rather
|
||||
than a submodule import, so a stubbed/partial torch (tests, exotic builds) cleanly
|
||||
reports None instead of picking a stale real module out of ``sys.modules``."""
|
||||
try:
|
||||
import torch
|
||||
return getattr(getattr(torch, "_inductor", None), "config", None)
|
||||
except Exception: # noqa: BLE001 — no inductor -> nothing to snapshot/set
|
||||
return None
|
||||
|
||||
|
||||
def normalize_speed_mode(value: Optional[str]) -> str:
|
||||
|
|
@ -342,6 +358,19 @@ def _compile_repeated_blocks(
|
|||
for _limit_attr in ("recompile_limit", "cache_size_limit"): # name varies by torch ver
|
||||
if hasattr(dynamo_cfg, _limit_attr):
|
||||
setattr(dynamo_cfg, _limit_attr, max(getattr(dynamo_cfg, _limit_attr) or 0, 64))
|
||||
# Match eager's intermediate rounding inside inductor's fused pointwise kernels:
|
||||
# by default they keep chains in fp32 where eager materialises bf16 between ops,
|
||||
# a per-forward rounding delta that a multi-step denoise amplifies chaotically.
|
||||
# Measured (B200, scripts/image_speedmem_bench.py, pairwise LPIPS of the
|
||||
# compiled tier vs the same-stack eager tier): Qwen-Image 0.019 -> 0.006 at
|
||||
# identical speed, FLUX.1-dev 0.046 -> 0.029 at +2% step time, FLUX.2-klein
|
||||
# 0.018 -> 0.017 at identical speed; on the video DiT (HunyuanVideo-1.5-720p)
|
||||
# full-clip LPIPS vs bit-exact drops 0.221 -> 0.052 at zero cost. Process-
|
||||
# global, so snapshot_backend_flags carries it and unload restores the prior
|
||||
# value.
|
||||
inductor_cfg = _inductor_config()
|
||||
if inductor_cfg is not None and hasattr(inductor_cfg, "emulate_precision_casts"):
|
||||
inductor_cfg.emulate_precision_casts = True
|
||||
except Exception as exc: # noqa: BLE001 — optimisation only
|
||||
_warn(logger, "compile_repeated_blocks", exc)
|
||||
return False
|
||||
|
|
@ -354,6 +383,20 @@ def _compile_repeated_blocks(
|
|||
engaged = True
|
||||
except Exception as exc: # noqa: BLE001 — optimisation only
|
||||
_warn(logger, "compile_repeated_blocks", exc)
|
||||
continue
|
||||
# A step cache engaged BEFORE this compile (the production load order) has
|
||||
# already wrapped each block's forward in a @torch.compiler.disable'd hook, so
|
||||
# the compute branch would run eager on every non-skipped step and forfeit the
|
||||
# regional compile entirely. Re-point the hooks' inner forward at compiled
|
||||
# wrappers; no-op when no cache hooks are installed. The toggle path (cache
|
||||
# engaged after load) is armed by apply_step_cache instead. Lazy import:
|
||||
# diffusion_cache imports nothing from this module, but keep the dependency
|
||||
# one-directional at import time.
|
||||
try:
|
||||
from .diffusion_cache import _compile_hooked_block_inners
|
||||
_compile_hooked_block_inners(transformer, logger)
|
||||
except Exception as exc: # noqa: BLE001 — optimisation only
|
||||
_warn(logger, "cache-hook inner compile", exc)
|
||||
return engaged
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -356,3 +356,206 @@ def test_toggle_noop_without_cache_support(monkeypatch):
|
|||
|
||||
def test_toggle_noop_without_transformer():
|
||||
assert maybe_toggle_step_cache(types.SimpleNamespace(), steps = 28) is None
|
||||
|
||||
|
||||
# ── compiled cache-hook inners (regional compile x step cache composition) ──────────
|
||||
import functools # noqa: E402
|
||||
|
||||
from core.inference.diffusion_cache import ( # noqa: E402
|
||||
_compile_hooked_block_inners,
|
||||
_invalidate_child_registry_cache,
|
||||
_restore_hooked_block_inners,
|
||||
)
|
||||
|
||||
|
||||
class _BoundInner:
|
||||
"""Provides a plain bound method for fn_ref.original_forward (__self__ present)."""
|
||||
|
||||
def forward(self, *args, **kwargs):
|
||||
return "eager"
|
||||
|
||||
|
||||
def _hooked_block(
|
||||
*,
|
||||
compiled = True,
|
||||
hook_name = "fbc_block_hook",
|
||||
bound = True,
|
||||
):
|
||||
inner = _BoundInner()
|
||||
orig = inner.forward if bound else functools.partial(_BoundInner.forward, inner)
|
||||
hook = types.SimpleNamespace(fn_ref = types.SimpleNamespace(original_forward = orig))
|
||||
block = types.SimpleNamespace(
|
||||
_diffusers_hook = types.SimpleNamespace(hooks = {hook_name: hook}),
|
||||
_compiled_call_impl = object() if compiled else None,
|
||||
)
|
||||
return block, hook, orig
|
||||
|
||||
|
||||
def _fake_dit(blocks):
|
||||
return types.SimpleNamespace(modules = lambda: [types.SimpleNamespace()] + blocks)
|
||||
|
||||
|
||||
def _stub_torch_compile(monkeypatch):
|
||||
compiled_calls = []
|
||||
|
||||
def _compile(fn, **kwargs):
|
||||
compiled_calls.append((fn, kwargs))
|
||||
wrapper = lambda *a, **k: fn(*a, **k) # noqa: E731
|
||||
wrapper._unsloth_test_compiled_of = fn
|
||||
return wrapper
|
||||
|
||||
torch = types.ModuleType("torch")
|
||||
torch.compile = _compile
|
||||
monkeypatch.setitem(sys.modules, "torch", torch)
|
||||
return compiled_calls
|
||||
|
||||
|
||||
def test_arming_swaps_inner_for_compiled_wrapper(monkeypatch):
|
||||
calls = _stub_torch_compile(monkeypatch)
|
||||
block, hook, orig = _hooked_block()
|
||||
assert _compile_hooked_block_inners(_fake_dit([block])) == 1
|
||||
assert hook.fn_ref.original_forward is not orig
|
||||
assert hook.fn_ref.original_forward._unsloth_test_compiled_of is orig
|
||||
assert hook._unsloth_orig_inner is orig
|
||||
# The inner compile must match the cache-active tier: graph-breakable + dynamic.
|
||||
assert calls[0][1] == {"fullgraph": False, "dynamic": True}
|
||||
|
||||
|
||||
def test_arming_is_idempotent(monkeypatch):
|
||||
_stub_torch_compile(monkeypatch)
|
||||
block, hook, _ = _hooked_block()
|
||||
dit = _fake_dit([block])
|
||||
assert _compile_hooked_block_inners(dit) == 1
|
||||
once = hook.fn_ref.original_forward
|
||||
assert _compile_hooked_block_inners(dit) == 0 # marker short-circuits
|
||||
assert hook.fn_ref.original_forward is once
|
||||
|
||||
|
||||
def test_arming_skips_uncompiled_blocks(monkeypatch):
|
||||
# An eager-tier load has no _compiled_call_impl: the hook must stay untouched
|
||||
# (compiling the inner would ADD compile where the user chose eager).
|
||||
_stub_torch_compile(monkeypatch)
|
||||
block, hook, orig = _hooked_block(compiled = False)
|
||||
assert _compile_hooked_block_inners(_fake_dit([block])) == 0
|
||||
assert hook.fn_ref.original_forward is orig
|
||||
|
||||
|
||||
def test_arming_skips_partial_captured_inner(monkeypatch):
|
||||
# A stacked hook chain (e.g. group offload) captures a functools.partial, not the
|
||||
# plain bound method; arming would compile the wrong layer of the chain.
|
||||
_stub_torch_compile(monkeypatch)
|
||||
block, hook, orig = _hooked_block(bound = False)
|
||||
assert _compile_hooked_block_inners(_fake_dit([block])) == 0
|
||||
assert hook.fn_ref.original_forward is orig
|
||||
|
||||
|
||||
def test_arming_covers_every_cache_hook_family(monkeypatch):
|
||||
# FBCache is the image cache today, but the hook-name table already covers the
|
||||
# MagCache layout too (same fn_ref shape), so a future mode arms for free.
|
||||
_stub_torch_compile(monkeypatch)
|
||||
names = (
|
||||
"mag_cache_leader_block_hook",
|
||||
"mag_cache_block_hook",
|
||||
"fbc_leader_block_hook",
|
||||
"fbc_block_hook",
|
||||
)
|
||||
blocks = [_hooked_block(hook_name = n)[0] for n in names]
|
||||
assert _compile_hooked_block_inners(_fake_dit(blocks)) == len(names)
|
||||
|
||||
|
||||
def test_restore_puts_the_exact_original_back(monkeypatch):
|
||||
_stub_torch_compile(monkeypatch)
|
||||
block, hook, orig = _hooked_block()
|
||||
dit = _fake_dit([block])
|
||||
_compile_hooked_block_inners(dit)
|
||||
_restore_hooked_block_inners(dit)
|
||||
assert hook.fn_ref.original_forward is orig
|
||||
assert hook._unsloth_orig_inner is None
|
||||
|
||||
|
||||
def test_restore_tolerates_fakes_without_modules():
|
||||
_restore_hooked_block_inners(_MixinTransformer()) # no .modules(): no-op
|
||||
|
||||
|
||||
def test_apply_step_cache_arms_compiled_blocks_on_toggle(monkeypatch):
|
||||
# The generation-time toggle engages the cache AFTER the load already compiled the
|
||||
# blocks; apply_step_cache must arm the fresh hooks itself.
|
||||
_stub_diffusers(monkeypatch)
|
||||
_stub_torch_compile(monkeypatch)
|
||||
block, hook, orig = _hooked_block()
|
||||
|
||||
class _T(_MixinTransformer):
|
||||
def modules(self):
|
||||
return [block]
|
||||
|
||||
t = _T()
|
||||
engaged = apply_step_cache(_pipe(t), mode = "fbcache")
|
||||
assert engaged == TC_FBCACHE
|
||||
assert hook.fn_ref.original_forward is not orig
|
||||
assert hook._unsloth_orig_inner is orig
|
||||
|
||||
|
||||
def test_toggle_disable_restores_inners_before_disable(monkeypatch):
|
||||
# remove_hook splices fn_ref.original_forward back into module.forward, so the
|
||||
# compiled wrapper must be swapped out BEFORE disable_cache runs.
|
||||
_stub_diffusers(monkeypatch)
|
||||
order = []
|
||||
|
||||
class _T(_ToggleTransformer):
|
||||
def disable_cache(self):
|
||||
super().disable_cache()
|
||||
order.append("disable")
|
||||
|
||||
def modules(self):
|
||||
order.append("restore-walk")
|
||||
return []
|
||||
|
||||
t = _T()
|
||||
maybe_toggle_step_cache(_pipe(t), steps = 28)
|
||||
mode = maybe_toggle_step_cache(_pipe(t), steps = 8)
|
||||
assert mode is None and t.disables == 1
|
||||
assert order[-2:] == ["restore-walk", "disable"]
|
||||
|
||||
|
||||
def test_enable_failure_restores_inners_before_partial_disable(monkeypatch):
|
||||
# enable_cache can fail after hooking (and arming) some blocks; the partial-hook
|
||||
# cleanup must un-arm them before disable_cache splices original_forward back.
|
||||
_stub_diffusers(monkeypatch)
|
||||
order = []
|
||||
|
||||
class _T(_ToggleTransformer):
|
||||
def enable_cache(self, config):
|
||||
raise RuntimeError("block signature not recognised")
|
||||
|
||||
def disable_cache(self):
|
||||
super().disable_cache()
|
||||
order.append("disable")
|
||||
|
||||
def modules(self):
|
||||
order.append("restore-walk")
|
||||
return []
|
||||
|
||||
t = _T()
|
||||
assert apply_step_cache(_pipe(t), mode = "fbcache") is None
|
||||
assert order == ["restore-walk", "disable"]
|
||||
|
||||
|
||||
# ── stale child-registry cache invalidation (mid-session enable) ────────────────────
|
||||
|
||||
|
||||
def test_enable_invalidates_stale_child_registry_cache(monkeypatch):
|
||||
# diffusers 0.39 caches the child-registry list on first cache_context use; an
|
||||
# UNCACHED generation already populates it (empty), so a later toggle-time
|
||||
# enable_cache would install hooks the context never reaches ("No context is set").
|
||||
_stub_diffusers(monkeypatch)
|
||||
t = _MixinTransformer()
|
||||
t._diffusers_hook = types.SimpleNamespace(_child_registries_cache = ["stale"])
|
||||
assert apply_step_cache(_pipe(t), mode = "fbcache") == TC_FBCACHE
|
||||
assert t._diffusers_hook._child_registries_cache is None
|
||||
|
||||
|
||||
def test_invalidate_child_registry_cache_tolerates_absence():
|
||||
_invalidate_child_registry_cache(types.SimpleNamespace()) # no registry: no-op
|
||||
reg = types.SimpleNamespace(_child_registries_cache = None)
|
||||
_invalidate_child_registry_cache(types.SimpleNamespace(_diffusers_hook = reg))
|
||||
assert reg._child_registries_cache is None
|
||||
|
|
|
|||
|
|
@ -393,3 +393,81 @@ def test_nvfp4_filter_keeps_vision_tower_dense(monkeypatch):
|
|||
assert ff(object(), "lm_head") is False
|
||||
assert ff(object(), "model.decoder.wo") is False
|
||||
assert ff(object(), "model.layers.5.self_attn.q_proj") is True
|
||||
|
||||
|
||||
# ── zero-output-row guard (per-row fp8 NaN protection) ───────────────────────────
|
||||
|
||||
|
||||
class _FakeAmaxVec:
|
||||
def __init__(self, vals):
|
||||
self._vals = vals
|
||||
|
||||
def __eq__(self, other): # noqa: PLW0642 -- tensor-style elementwise compare
|
||||
return _FakeAmaxVec([v == other for v in self._vals])
|
||||
|
||||
def any(self):
|
||||
return _FakeScalar(any(self._vals))
|
||||
|
||||
|
||||
class _FakeScalar:
|
||||
def __init__(self, v):
|
||||
self._v = v
|
||||
|
||||
def item(self):
|
||||
return self._v
|
||||
|
||||
|
||||
class _FakeWeight:
|
||||
"""Tensor-shaped stand-in supporting the exact chain the guard runs:
|
||||
``weight.abs().amax(dim = -1) == 0 -> .any().item()``."""
|
||||
|
||||
ndim = 2
|
||||
|
||||
def __init__(self, rows):
|
||||
self._rows = rows
|
||||
|
||||
def abs(self):
|
||||
return _FakeWeight([[abs(v) for v in r] for r in self._rows])
|
||||
|
||||
def amax(self, dim = -1):
|
||||
return _FakeAmaxVec([max(r) for r in self._rows])
|
||||
|
||||
|
||||
def test_weight_zero_output_row_detection():
|
||||
# A dead output row NaNs torchao's per-row fp8 (scale 0 -> 0/0); SDXL's
|
||||
# text_encoder_2 (OpenCLIP bigG) really ships one in layers.2.self_attn.out_proj --
|
||||
# measured: every fp8_dynamic SDXL render was black until the row is kept dense.
|
||||
zero_row = types.SimpleNamespace(weight = _FakeWeight([[0.1, 0.2], [0.0, 0.0]]))
|
||||
dense = types.SimpleNamespace(weight = _FakeWeight([[0.1, 0.2], [0.3, 0.0]]))
|
||||
assert dp._weight_has_zero_output_row(zero_row) is True
|
||||
assert dp._weight_has_zero_output_row(dense) is False
|
||||
# Non-2D / absent weights are not the per-row scheme's input: never flagged.
|
||||
w3 = _FakeWeight([[1.0]])
|
||||
w3.ndim = 3
|
||||
assert dp._weight_has_zero_output_row(types.SimpleNamespace(weight = w3)) is False
|
||||
assert dp._weight_has_zero_output_row(types.SimpleNamespace()) is False
|
||||
|
||||
# An unreadable weight falls through to quantize_'s own handling.
|
||||
class _Boom:
|
||||
@property
|
||||
def weight(self):
|
||||
raise RuntimeError("meta tensor")
|
||||
|
||||
assert dp._weight_has_zero_output_row(_Boom()) is False
|
||||
|
||||
|
||||
def test_fp8_dynamic_filter_skips_zero_row_linear(monkeypatch):
|
||||
# The fp8_dynamic caster must leave a zero-output-row Linear dense while the rest
|
||||
# of the encoder still quantises (a family-wide deny would forfeit the whole win).
|
||||
_stub_torch(monkeypatch)
|
||||
captured: dict = {}
|
||||
_stub_transformer_quant(monkeypatch, captured)
|
||||
enc = types.SimpleNamespace(_keep_in_fp32_modules = [])
|
||||
|
||||
dp._cast_fp8_dynamic(enc, _target())
|
||||
|
||||
ff = captured["filter_fn"]
|
||||
dead = types.SimpleNamespace(weight = _FakeWeight([[0.5, 0.5], [0.0, 0.0]]))
|
||||
live = types.SimpleNamespace(weight = _FakeWeight([[0.5, 0.5], [0.5, 0.5]]))
|
||||
assert ff(dead, "text_model.encoder.layers.2.self_attn.out_proj") is False
|
||||
assert ff(live, "text_model.encoder.layers.2.mlp.fc1") is True
|
||||
|
|
|
|||
|
|
@ -531,3 +531,85 @@ def test_fp16_accum_allowed_on_fp16_dtype_under_max(monkeypatch):
|
|||
)
|
||||
assert applied["fp16_accum"] is True
|
||||
assert torch.backends.cuda.matmul.allow_fp16_accumulation is True
|
||||
|
||||
|
||||
# ── inductor precision-cast emulation (compile-vs-eager numeric parity) ─────────
|
||||
|
||||
|
||||
def _stub_inductor_config(
|
||||
monkeypatch,
|
||||
torch,
|
||||
*,
|
||||
emulate = False,
|
||||
):
|
||||
"""Attach a fake ``_inductor.config`` to the stubbed torch module (diffusion_speed
|
||||
resolves it as attributes off the imported torch, never via sys.modules -- so the
|
||||
real torch._inductor lingering in sys.modules cannot leak into stubbed tests)."""
|
||||
cfg = types.SimpleNamespace(emulate_precision_casts = emulate)
|
||||
torch._inductor = types.SimpleNamespace(config = cfg)
|
||||
return cfg
|
||||
|
||||
|
||||
def test_regional_compile_enables_emulate_precision_casts(monkeypatch):
|
||||
# Inductor's fused pointwise kernels keep intermediates in fp32 where eager rounds
|
||||
# to bf16 between ops; over a multi-step denoise that compounds to a visible drift.
|
||||
# emulate_precision_casts restores eager's rounding at zero measured speed cost, so
|
||||
# the regional compile path must switch it on.
|
||||
torch = _stub_torch(monkeypatch)
|
||||
_stub_gguf_accel(monkeypatch)
|
||||
cfg = _stub_inductor_config(monkeypatch, torch, emulate = False)
|
||||
pipe = _Pipe(with_compile = True)
|
||||
applied = apply_speed_optims(
|
||||
pipe, _target(), is_gguf = False, family = _family(), speed_mode = SPEED_DEFAULT
|
||||
)
|
||||
assert applied["compiled"] is True
|
||||
assert cfg.emulate_precision_casts is True
|
||||
|
||||
|
||||
def test_snapshot_restores_emulate_precision_casts(monkeypatch):
|
||||
# The flag is process-global, so the unload path must restore the pre-load value
|
||||
# exactly like the TF32 / cudnn.benchmark globals.
|
||||
torch = _stub_torch(monkeypatch)
|
||||
cfg = _stub_inductor_config(monkeypatch, torch, emulate = False)
|
||||
snap = snapshot_backend_flags()
|
||||
assert snap["inductor_emulate_precision_casts"] is False
|
||||
cfg.emulate_precision_casts = True
|
||||
restore_backend_flags(snap)
|
||||
assert cfg.emulate_precision_casts is False
|
||||
|
||||
|
||||
def test_missing_inductor_config_is_tolerated(monkeypatch):
|
||||
# A build without torch._inductor (or with the flag renamed) must neither break the
|
||||
# snapshot nor the compile path.
|
||||
_stub_torch(monkeypatch) # the stub torch has no _inductor attribute
|
||||
_stub_gguf_accel(monkeypatch)
|
||||
snap = snapshot_backend_flags()
|
||||
assert "inductor_emulate_precision_casts" not in snap
|
||||
pipe = _Pipe(with_compile = True)
|
||||
applied = apply_speed_optims(
|
||||
pipe, _target(), is_gguf = False, family = _family(), speed_mode = SPEED_DEFAULT
|
||||
)
|
||||
assert applied["compiled"] is True
|
||||
|
||||
|
||||
def test_regional_compile_arms_cache_hook_inners(monkeypatch):
|
||||
# The production load order engages the step cache BEFORE compile, so the regional
|
||||
# compile pass must re-arm the already-installed cache hooks with compiled inner
|
||||
# forwards (otherwise every computed step runs eager under the hook's
|
||||
# torch.compiler.disable and forfeits the regional compile).
|
||||
_stub_torch(monkeypatch)
|
||||
_stub_gguf_accel(monkeypatch)
|
||||
from core.inference import diffusion_cache as dc_mod
|
||||
|
||||
armed = []
|
||||
monkeypatch.setattr(
|
||||
dc_mod,
|
||||
"_compile_hooked_block_inners",
|
||||
lambda transformer, logger = None: armed.append(transformer) or 1,
|
||||
)
|
||||
pipe = _Pipe(with_compile = True)
|
||||
applied = apply_speed_optims(
|
||||
pipe, _target(), is_gguf = False, family = _family(), speed_mode = SPEED_DEFAULT
|
||||
)
|
||||
assert applied["compiled"] is True
|
||||
assert armed == [pipe.transformer]
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue