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:
parent
ec90b8658d
commit
de2f22df2b
7 changed files with 997 additions and 2 deletions
|
|
@ -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)
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue