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