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:
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue