feat(fp8): run scaled and raw fp8 checkpoints on the fp8 tensor cores - #232
Merged
Conversation
Z-Image was excluded from FP8 storage in invoke-ai#8945 because diffusers' enable_layerwise_casting() was called with the global torch dtype (fp16) while Z-Image loads in bf16: skipped modules stayed bf16, hooked ones produced fp16, and attention crashed. That root cause was fixed later in the same PR — the compute dtype now comes from the model's own parameters — so the exclusion is obsolete. Removing it alone is not enough. Our hook-based cast (invoke-ai#9231) dropped one thing diffusers' enable_layerwise_casting() did: honoring the model's declared _skip_layerwise_casting_patterns. Z-Image needs it, and not for quality — TimestepEmbedder.forward reads self.mlp[0].weight.dtype and casts its *input* to it. With an fp8 weight the input becomes float8 before our pre-hook restores the weight, and F.linear dies with: RuntimeError: "addmm_cuda" not implemented for 'Float8_e4m3fn' which is why ZImageTransformer2DModel declares ['t_embedder', 'cap_embedder']. _apply_fp8_to_nn_module now takes extra_skip_patterns and the caller passes the model's list. For other models this is a strict superset of our defaults (FLUX/SD3 pos_embed+norm, UNet norm, CogView4 also proj_out), so it only ever skips more. Also wire the cast into ZImageCheckpointModel: only the diffusers loader called it, so the toggle was a silent no-op for single-file Z-Image models even though both paths build the same ZImageTransformer2DModel. Tested end to end on CUDA: transformer resident VRAM drops from ~11.5GB to 5880MB for both Z-Image-Turbo (diffusers) and Z-Image-Turbo (checkpoint, 14.37GB file), with clean output images in both cases.
The fp8_storage toggle was shown for Anima main models but did nothing: AnimaCheckpointModel never called _apply_fp8_layerwise_casting. Wire it in — the state dict is cast to a single model_dtype before load_state_dict, so the layerwise cast has one unambiguous compute dtype to restore to. Wiring alone renders a heavily dithered image with no fine detail. The cause is t_embedder: it produces the adaln_lora conditioning consumed by every block, so casting it to FP8 corrupts every token everywhere. None of the generic skip patterns match it — they target diffusers' module names (norm, pos_embed, patch_embed, proj_in/out) and this architecture names things differently. AnimaTransformer now declares _skip_layerwise_casting_patterns, the same attribute diffusers models use, so the loader needs no special-casing. Measured on CUDA, same seed/steps/CFG each run: casting nothing = broken at 1994MB; t_embedder alone = clean at 2010MB; adding x_embedder and final_layer changes nothing further (2012MB) and is kept as margin on the I/O layers; adaln_modulation was tested too and is deliberately not listed — it costs 168MB and made no difference. Against a bf16 reference (3988MB) the FP8 result keeps the same composition and loses only a little micro-detail.
Every quantized-format loader reaches _apply_fp8_layerwise_casting, and the cast
there is not a no-op. Verified on real layers:
- GGUF raises "Operation changed the dtype of GGMLTensor unexpectedly" at load.
- bnb NF4 corrupts silently: bnb.nn.LinearNF4 subclasses nn.Linear, so the
isinstance check passes and the packed uint8 payload is cast to float8.
Inference still returns finite numbers and the model just produces garbage
(max abs deviation 50.4 against a reference forward pass).
Both are reachable today by enabling the fp8_storage toggle, which the UI offered
for these models.
Guard on two levels, because a format check alone is not enough — an externally
quantized checkpoint can carry a plain `diffusers` format (e.g. SDNQ):
- _should_use_fp8 rejects gguf_quantized and both bnb formats.
- _apply_fp8_to_nn_module skips any module whose params are non-floating-point
or a torch.Tensor subclass, regardless of the model's declared format.
Frontend hides the toggle for quantized formats, so the control is not shown for
something the backend refuses.
Verified end to end: with fp8_storage forced true in the DB (the legacy case the
UI no longer offers), a GGUF Z-Image model now loads cleanly with no FP8 casting
and no GGMLTensor error, while non-quantized models still show the toggle and
still get cast.
…sor cores InvokeAI dequantized ComfyUI 'scaled fp8' checkpoints to bf16 at load time in three near-identical implementations, discarding both the VRAM saving and the ability to use the fp8 tensor cores. Measured on an RTX 4090: the dequantize round trip makes fp8 *slower* than bf16 (0.90x on FLUX.2 Klein 9B), while torch._scaled_mm reaches 1.29x vs bf16 and 1.63x vs the dequantized path on Krea-2 Turbo, at the same VRAM and with no visible quality loss. Adds a shared invokeai/backend/quantization/fp8_scaled.py that keeps the quantization intact (weight_scale, optional calibrated input_scale, and the per-layer full_precision_matrix_mult hints), and an fp8 branch in CustomLinear._autocast_forward that falls back to the dequantized path whenever any precondition fails, rather than raising mid-generation. Wires up the Krea-2 single-file loader as the first consumer. Remaining loaders (FLUX.2, Z-Image, Qwen-Image) still dequantize eagerly. Gated behind `fp8_compute` (default off): the fp8 matmul quantizes activations too, so images change at a fixed seed. The same flag also decides whether the weights stay quantized, since keeping them fp8 without the fp8 matmul would halve VRAM but run slower.
…em on fp8 tensor cores The Krea-2 single-file loader dequantized ComfyUI 'scaled fp8' checkpoints to bf16 at load, discarding both the VRAM saving and any chance of using the fp8 tensor cores. It now keeps the quantization and hands the scales to CustomLinear, which multiplies via torch._scaled_mm where the checkpoint allows it. Measured on an RTX 4090 with krea2TurboOfficialComfy_krea2TurboFp8 at 1024x1024: 884 vs 1107 ms/step (1.25x) and 12.24 GiB resident instead of ~25 GB in bf16, so the model is fully resident rather than streamed. Images are indistinguishable from the dequantized path (PSNR 29.95 dB). The checkpoint's _quantization_metadata marks 96 of 256 layers with full_precision_matrix_mult; those are honored and stay in bf16, which is what keeps fidelity in line with ComfyUI (and costs ~23% of the speedup). Two subtleties the measurements exposed, both covered by tests now: - apply_custom_layers_to_model leaves device autocasting disabled for fully resident models, so a check living only in _autocast_forward never runs. The fp8 branch is consulted before the autocasting split. - _quantization_metadata names layers in the native scheme while the scales are extracted after the native -> diffusers rename, so the per-layer flags matched nothing. The metadata paths are pushed through the same converter. Gated behind `fp8_compute` (default off): activations are quantized too, so images change at a fixed seed.
…data read fp8 weights are force-routed to sidecar patching, and the sidecar wrapper dispatches through _autocast_forward, so the fp8 branch has to survive that route with the LoRA residual added on top. Verified on Krea-2 Turbo fp8 with a 256-layer LoRA: 1008 vs 1241 ms/step (1.231x, against 1.251x without the LoRA), all 256 fp8 modules routed to sidecar, and images equivalent between both paths (PSNR 26.83 dB). Reading the safetensors header metadata no longer fails the load. It only enriches fp8 handling with the per-layer full_precision_matrix_mult hints, so an unreadable header now warns and continues rather than raising — but it does warn, because without the hints layers the quantizer marked unsafe would silently be multiplied in fp8.
Adds a settings matrix (only `fp8_compute` is needed; `fp8_storage` is bypassed on that path) and logs when a redundant fp8_storage setting was skipped, so the case is not silent.
The Qwen-Image i2l node hardcoded vae.disable_tiling(), so a full-frame encode was the only option. At 2560x1440 that peaks at 9.26 GiB — on top of a resident multi-GB transformer, which is what makes an upscale round-trip run out of headroom exactly at this node while every other node fits. Adds `tiled` / `tile_size` input fields following the SD/SDXL i2l node, OR'd with the global force_tiled_decode setting. Off by default, so behaviour is unchanged unless enabled. estimate_vae_working_memory_qwen_image gains a matching tile_size parameter. Without it the change would be inert: the cache would keep reserving the full-frame figure (10.99 GiB at 2560x1440) and evict models to honour it, no matter what the VAE actually does. Tiled, it budgets one tile plus 25% overlap plus the resident RGB image, mirroring estimate_vae_working_memory_wan. Measured through the node at 2560x1440: 10.99 -> 0.26 GiB reserved, 9.26 -> 0.17 GiB actual peak, identical latent shape. Tiled latents differ by ~1.4% relative L2 on noise input (worst case for tile blending; real images blend far better), which is why this stays opt-in.
Both nodes reserve working memory for a full-frame operation, which at high resolutions exceeds a 24 GB card, so the model cache evicts everything else to honour it. On CUDA at 2560x1440: 19.91 GiB for the decode and 10.99 GiB for the encode. Tiling is the intended escape hatch, but it did not work on either node: - qwen_image_i2l hardcoded vae.disable_tiling(), so it could not be enabled. - qwen_image_l2i honoured the global force_tiled_decode, but computed its working-memory estimate before and independently of that flag. Tiling bounded the VAE while the cache still reserved the full-frame figure, so the memory was never freed for anything else — effectively inert. Adds `tiled` / `tile_size` input fields to both nodes following the SD/SDXL i2l/l2i nodes, OR'd with force_tiled_decode. Off by default; behaviour is unchanged unless enabled. estimate_vae_working_memory_qwen_image gains a matching tile_size parameter, and both nodes resolve tile_size=0 to the VAE default (256px) before estimating. Tiled it budgets one tile plus 25% overlap plus the resident RGB image, mirroring estimate_vae_working_memory_wan. Without this the change would be cosmetic on i2l and remain inert on l2i. Measured through the i2l node at 2560x1440: 10.99 -> 0.26 GiB reserved, 9.26 -> 0.17 GiB actual peak, identical latent shape. Verified across eight resolutions that tiled and untiled encodes produce the same latent dimensions. Tiled latents differ by ~1.4% relative L2 on noise input (worst case for tile blending), which is why this stays opt-in. Also fixes a crash in qwen_image_i2l: `width`/`height` are `int | None`, but the workflow UI sends 0 for an unset number input, and `0 is not None` reached `image.resize((0, 0))` -> "height and width must be > 0". Non-positive values are now treated as unset, matching how tile_size uses 0.
Switching the encoder to fp8 compute left everything the checkpoint does not quantize in bf16, growing it from 4236MB to 4999MB. On a GPU already holding a ~12GB transformer that was enough to push the transformer out of a full VRAM load, which is far more expensive than the encoder change ever saved. Cast the remainder to fp8 storage, with two exclusions. Layers carrying a weight_scale keep it and go through _scaled_mm -- the cast hooks would upcast them without applying the scale. nn.Embedding is skipped because the token embedding table is the encoder's input representation: quantizing it doubles the error against bf16 (relative L2 0.0079 -> 0.0163) to save 371MiB, a bad trade for a model whose whole job is text fidelity. The old fp8_storage path did cast it, so this is strictly more accurate than what shipped before. fp8_storage 116.8 ms 4.14 GiB rel L2 0.0351 (original) fp8 compute 61.9 ms 4.88 GiB rel L2 0.0079 (previous commit) this 65.3 ms 4.50 GiB rel L2 0.0079
…llings A scaled-fp8 checkpoint may ship an input_scale of exactly 1.0, meaning the producer wrote the field without calibrating it. Taking it at face value replaces the per-forward amax scale with no scaling at all, so activations above the fp8 maximum saturate. Measured relative error against a bf16 reference, dynamic vs a 1.0 scale: |x|max 368 0.0274 0.0262 |x|max 1928 0.0277 0.4356 |x|max 30976 0.0253 0.9639 Inside +/-448 the two are equivalent -- fp8_e4m3 is a floating-point format, so a scale factor buys no relative precision the way it would for int8. Above it the unscaled path collapses. Non-finite and non-positive scales are rejected for the same reason: they cannot be a valid divisor. Also accept `.scale_input` as an alias for `.input_scale`, mirroring the `.scale_weight`/`.weight_scale` pair we already handle. Previously such a key was left in the state dict, and the Qwen3-VL loader deleted it outright, so a calibrated activation scale was silently discarded and every forward paid the amax reduction. That delete is now redundant and removed.
FP8 Compute had no user-facing documentation, and the existing FP8 Storage page actively claimed a compute path might arrive "later" — it is already here. The part worth writing down is the reproducibility constraint. torch._scaled_mm needs both operands on the same device, so a layer whose weights are still in RAM silently falls back to the dequantized BF16 path. Which layers that hits depends on how much of the model happened to fit, and that shifts between runs, so the same seed stops reproducing. Measured on a 24GB card with a ~12GB transformer at 88-95% residency: two runs with identical seed and settings differed in 98.7% of pixels; fully resident, repeated runs were bit-identical. Also corrects "FP8 + partial loading: fully supported" — true for Storage, but for Compute it costs 47% per step on top of the reproducibility loss. Regenerates settings.json, which predated both fp8 settings.
A checkpoint can ship fp8 weights with no weight_scale. The runtime already handles them — scaled_mm_linear treats weight_scale as optional — but the loaders never let them through: FLUX, FLUX.2 and Z-Image cast the whole state dict to bf16, discarding both the VRAM saving and the tensor cores. Krea-2 kept them by accident and said nothing about it. Only nn.Linear.weight is preserved. That restriction is not cosmetic: a Z-Image checkpoint quantized everything, 243 of its 453 fp8 tensors being 1-D biases, norm weights and a learned pad token. Keeping those fp8 saves nothing usable and breaks inference — the value reaches the activations and the next Linear gets an fp8 input, which dies in x.abs() with "abs_cuda" not implemented. A model's own _skip_layerwise_casting_patterns is honored on top, for Linears whose forward casts activations to their weight's dtype. Also stops fp8 storage from silently defeating fp8 compute: layerwise casting restores the compute dtype before every forward, so on an already-fp8 checkpoint the matmul would quietly fall back and the VRAM toggle would make the model slower with no indication why. Verified end-to-end on Z-Image unstableRevolution_V2Fp8, 1024x1024, 30 steps: transformer 11739MB -> 5881MB (both 100% resident), 1.60 -> 1.27 s/it. 1.297 -> 1.023 s/it (3 warm runs each), transformer 11740MB -> 5881MB, residency 95.5-100% -> 100%.
Conflicts in the fp8 loaders: - load_default.py: keep both guards. main added the idempotence early-return (the FP8_COMPUTE_DTYPE_ATTR marker), this branch added the "already fp8, leave it on the tensor cores" early-return. They gate different things, so both stay; the marker check runs first. - krea2.py: the transformer loader no longer calls _dequantize_scaled_fp8 — extract_fp8_scaled_layers/dequantize_fp8_scaled replaces it and honours the fp8_compute setting. main's RAM-spike fix to that helper is still relevant for its remaining callers, so the helper keeps main's dtype-aware version. target_device/model_dtype now come from main's hoisted position. - krea2.py Qwen3-VL encoder: same, plus main's .comfy_quant / scale_input key cleanup is already covered — extract_fp8_scaled_layers pops both.
…evice Review follow-ups for invoke-ai#9478. 1. A scaled fp8 layer that matched a skip pattern (or was not an nn.Linear weight) went through cast_state_dict's plain .to(dtype), which drops the weight_scale; attach_fp8_scales then skipped it for no longer being fp8, so the scale was lost and the weight ended up off by 1/weight_scale. Krea-2 hits this on ordinary ComfyUI exports — time_embed.linear_1/linear_2 are quantized like any other Linear and match the model's `time_embed` pattern. split_fp8_scaled_layers() now folds exactly those layers first, with the scale applied, and drops them from the mapping. The predicate the three callers share lives in can_stay_quantized(). attach_fp8_scales returning fewer than len(layers) now warns instead of reading as success. 2. device_supports_fp8_matmul tested compute capability >= (8, 9). On ROCm that reports the gfx arch, so RDNA3 (gfx1100 -> (11, 0)) passed and every forward then raised. It now probes once per device with a real 16x16 _scaled_mm; the capability compare is only a pre-filter. Without this the fallback that exists to avoid a mid-generation crash was gated on the same wrong answer. 3. Flux2CheckpointModel._dequantize_fp8_weights converted every float8 tensor unconditionally, before cast_state_dict ever saw the state dict — so the FLUX.2 half of the raw-fp8 path never executed and its log line could not fire. It now takes the same keep_fp8 gate. 4. Z-Image deleted .scale_weight without applying it and could not tell a scaled checkpoint from a raw one. It goes through extract_fp8_scaled_layers + attach_fp8_scales like Krea-2 now, reads the full_precision_matrix_mult hints from both the header and the .comfy_quant markers, and accepts both scale spellings. The fused-QKV split carries the quantization side-channel with it: a scale left on `...attention.qkv` keys onto a module the diffusers model does not have, so all three split weights would stay quantized but unscaled. 5. make_room charged 1 byte/element for every fp8 tensor, but only 2-D Linear weights outside the skip patterns actually stay quantized. On a checkpoint that quantized all 453 of its tensors that is most of the reservation. predict_cast_state_dict_size() answers with the same predicate the cast uses.
ComfyUI-style scaled fp8 checkpoints spell the scale either `.weight_scale` or `.scale_weight`; fp8_scaled.py has always known both, but four loaders carried their own `.weight_scale`-only literals. flux.py (FLUX.2, Qwen-Image) and the z_image text-encoder path applied one spelling and stripped both, so a `.scale_weight` checkpoint lost its scales without a word and every quantized weight came out off by 1/weight_scale. mistral_encoder.py neither applied nor stripped it, so the leftover key tripped load_state_dict(..., strict=True). Adds iter_weight_scale_pairs() and is_scale_metadata_key() and routes all four loaders through them. Validated against five real checkpoints (Krea-2 x2, FLUX.2 Klein 9B, Qwen-Image Edit 2511, Qwen2.5-VL 7B), folding real quantized layers through the real loader code: results are bit-identical for the spelling each file already used, and identical again once its scale keys are renamed to the other spelling. On the previous code the renamed variants come out with a relative error of 530-1682.
zImageTurboFP8Kijai_fp8ScaledE4m3fn is ComfyUI scaled fp8 with the `.scale_weight` spelling and scales between 1.5 and 7.6. Loading it on main drops the scales in the transformer's keys_to_remove filter, leaving each weight at 1/scale of its true magnitude -- 13% to 65%, layer by layer, silently. Captures the key layout (76 keys, shapes and dtypes) as a fixture next to the existing bf16 one, with tests asserting every `.scale_weight` is recognized, the keys are consumed rather than left for a strict load, and the scale value is carried through. One test guards the fixture itself, since a bf16 recapture would make the rest vacuous.
OCP Microscaling stores one E8M0 exponent per 32-element block, written as `uint8` because safetensors has no E8M0 dtype. The previous commit's block-wise expansion made such checkpoints *loadable* - and they load wrong: end to end they generate a pure-noise image with nothing in the log. Decoding the byte as `2**(v-127)` is not the fix. Verified against a real pair: the MXFP8 and scaled-fp8 builds of `krea2TurboOfficialComfy` share all 174 bf16 tensors bit-for-bit, so the scaled build is an exact reference for the same weights. Against it the decoded weights correlate only 0.60 - worse than the unscaled raw codes at 0.75 - while the block axis itself checks out (4% spread within a 32-block along dim 1, against 38% for the alternative). The measured per-block scale has no monotonic relation to the byte: 112 and 116 yield the same true scale. That rules out a wrong exponent bias and points at a swizzled scale layout, which is a de-swizzle to implement, not a constant to correct. So refuse the format at extraction, naming the layer and saying what would happen otherwise. A loud refusal beats both the shape error this used to raise and the silent garbage it would produce now.
…be caching M1. The Qwen3-VL encoder was the one caller that never went through `split_fp8_scaled_layers`. Its per-key `if dtype is not FP8_DTYPE` cast looks equivalent to `cast_state_dict` and is not: it keeps *every* fp8 tensor, so a 1-D fp8 norm from a "quantizes everything" checkpoint kept its scale as an unused buffer on a non-Linear and the forward computed on raw fp8 codes, off by 1/weight_scale with nothing logged. Two more facets of the same hole: an e5m2 scaled weight was cast without its scale, and a block-wise scale reached `scaled_mm_linear` unchecked and raised mid-generation. Now routed through extract -> split -> cast like every other loader; the `skip=` callback for the storage pass is unchanged. m1. A multi-element `input_scale` aborted the load with a raw shape error from `reshape(())`. `scaled_mm_linear` scales activations by one value, so there is nothing to do with a per-channel activation scale — it is dropped and the dynamic amax path takes over, matching how every other malformed side-channel in the module behaves. m2. `reattach_layer_sidechannel` tested for `<destination>.weight`, which assumes every quantizable module stores its parameter as `weight`. A quantized norm writes `<path>.scale`, so the guard rejected a destination it could have placed and dropped the scale. It now asks whether any tensor sits under the destination. The orphan report at the krea2 call site moves from DEBUG to INFO: a dropped scale has no other symptom. m3. Mistral strips `text_encoder.`/`language_model.` on top of the generic prefixes, but its hints only went through the generic tuple — so a prefixed redistribution shipping `_quantization_metadata` had every `full_precision_matrix_mult` silently ignored. The prefix list is now a named constant used by both the sd strip and the hint rename. m4. The two legacy folds multiplied `(out, in) * (out,)`, which broadcasts on the last axis and therefore scales input channels: wrong on a square weight, a shape error otherwise. Both now use `expand_weight_scale`, as the shared dequantizers already did. m5. `make_room` moved back ahead of `split_fp8_scaled_layers` at all five call sites. The split dequantizes its unusable subset through fp32, so reserving afterwards let that transient peak land on an unreserved cache. The prediction is unaffected — it applies the same `can_stay_quantized` predicate. m6. `_probe_fp8_matmul` cached every unrecognized `RuntimeError` as "this GPU cannot do fp8", which is the failure mode the OOM branch exists to avoid, just reached through a different wording. It now caches only errors that name the support constraint and treats the open set as inconclusive; an inconclusive answer costs one 16x16 matmul on the next load. m7. Test gaps closed: the already-fp8 storage guard and the `skip` callback now have tests that fail if reverted; `set_fp8_matmul_enabled` is restored to `None` rather than left pinned to `False` after `TestCustomLinearIntegration`; the 1-D wrong-length branch of `is_matmul_usable_scale` is pinned; and hint-rename plumbing is pinned per loader, including FLUX.1, which had no test at all. The Mistral case pins the failure mode as well as the fix. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Picks up invoke-ai#9538, which lands in the same code this branch rewrites. - krea2.py: keep both. main added `sd.clear()` after `load_state_dict` so the `assign=True` aliases free before the FP8 cast; it now runs ahead of this branch's fp8 logging block, which returns early on an fp8 checkpoint and would otherwise skip it. - z_image.py: main's block was a comment recording the scaled-fp8 mis-load as a known caveat ("the raw fp8 codes are cast to model_dtype unscaled and the model loads with wrong weights"). That is the bug this branch fixes, so the caveat is obsolete rather than merged. - test_load_default_fp8.py: additive on both sides, both kept. Semantic follow-up, not a textual conflict: invoke-ai#9538 introduced `_model_declared_skip_patterns`, which unions `_skip_layerwise_casting_patterns` with `_keep_in_fp32_modules`. The five loaders in this branch still read the first attribute directly, so the checkpoint load path would have honoured a narrower skip list than the layerwise cast does. All five now call the shared helper. test_z_image_fp8_wiring.py builds its loader with `object.__new__` and stubs only `_ram_cache`; the scaled-fp8 path also reads the safetensors header and asks the device about the fp8 matmul, so the stub gained `_logger` and `_torch_device`.
…olds Round-4 review follow-ups for invoke-ai#9478. R1. `make_room` runs before `split_fp8_scaled_layers`, but the prediction applied only `can_stay_quantized` while the split applies `is_matmul_usable_scale` on top. A block-wise-scaled 2-D Linear weight therefore predicted at 1 byte per element and arrived at 2 — half the reservation, and on a checkpoint using that layout the shortfall is the whole quantized-Linear set. Both now go through one predicate, `survives_split_and_cast`, and `predict_cast_state_dict_size` takes the `scaled_layers` mapping the caller is about to hand the split. All five call sites pass it; the docstrings and the per-call-site comment that claimed the prediction was split-independent are corrected. Also found while verifying: a 1-D weight_scale whose length matches neither layout reaches the fold and dies in torch's broadcast, naming neither the layer nor the file. `expand_weight_scale` now reports it as the malformed quantization metadata it is. r1. The `expand_weight_scale` call sites in the Mistral encoder and the Z-Image Qwen3 encoder had no test, so reverting either to its local multiply passed CI — the exact gap this fix exists to close, since a wrong scale axis is silent on a square weight. Both are pinned now. The Z-Image fold was a loop inside `_load_from_singlefile`, unreachable without a real checkpoint, so it is extracted as `_fold_comfy_scaled_weights`. r2. The Qwen3-VL split fix is pinned at helper level, no transformers needed: a scaled 1-D norm is folded with its scale applied, a block-wise Linear drops out of the returned mapping, and an e5m2 weight keeps its scale on the way down — the three things the old per-key loop got wrong. r3/r4. `_load_text_encoder` resolved `should_keep_fp8_weights` twice, which can disagree now that an inconclusive probe is deliberately not cached; it is resolved once. And `keep_fp8` now covers the storage path as well (`keep_matmul_fp8 or use_fp8_storage`), so the fp8_compute-off path no longer casts raw fp8 Linears to bf16 only for `_apply_fp8_to_nn_module` to re-quantize them — roughly 4.4 -> 8.9 GiB of reservation on the 4B encoder for a value-exact round trip.
Round-5 review follow-up for invoke-ai#9478. Blocker. Extracting `_fold_comfy_scaled_weights` in 9cd24e0 inserted it between `@ModelLoaderRegistry.register(...)` and the class it decorates, so the registry bound the helper. `register` does no type check and returns its argument unchanged, so nothing complained: the helper kept working at its call sites, the module imported, ruff was clean and the suite passed. The only symptom was that `ModelLoadService.load_model` instantiates the registered implementation, so every Z-Image Qwen3 single-file encoder load died with `TypeError: _fold_comfy_scaled_weights() got an unexpected keyword argument 'app_config'`. The helper now sits above the decorator. Pinned by a registry-wide guard rather than a case for this one loader: every entry must be a class and a `ModelLoaderBase` subclass. That covers the whole failure class, including the loaders with no end-to-end test. n1. `_apply_fp8_to_nn_module` skipped a module and left whatever dtype it arrived with. That is fine for a full-precision weight and silently fatal for one a loader kept in fp8: no cast back, no upcast pre-hook, and the forward runs on raw fp8 codes. Until now it depended on the loader's own skip list and `_FP8_DEFAULT_SKIP_PATTERNS` never naming the same 2-D Linear — true for Qwen3-VL by coincidence, not by construction. A pattern skip now restores the compute dtype. The `skip=` callback deliberately keeps the old behaviour: its one caller excludes scaled-fp8 layers, which are meant to stay quantized and carry their weight_scale into `_scaled_mm`, so upcasting them would drop the scale.
…sack/InvokeAI into feat/fp8_compute_raw
…d raw checkpoints Ports invoke-ai#9478 ahead of its upstream merge. Its three stacked prerequisites (invoke-ai#9414, invoke-ai#9415, invoke-ai#9416) are already on main, and the fp8 files were byte-identical to the PR's base, so only two conflicts arose: krea2's import block (both sides kept) and the regenerated openapi.json.
Pfannkuchensack
requested review from
JPPhoto,
blessedcoolant and
lstein
as code owners
September 9, 2026 00:07
The generator emits raw json.dumps output; both openapi-checks and the frontend prettier check compare against a prettier-formatted file.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Ports invoke-ai/InvokeAI#9478 into v7 ahead of its upstream merge, so the work stacked behind it can start here instead of waiting.
Adds FP8 Compute: checkpoints that already ship FP8 weights run directly on the FP8 tensor cores via
torch._scaled_mminstead of being unpacked to BF16 at load. Distinct from the existing FP8 Storage toggle, which quantizes a full-precision model yourself and still does the math in BF16. New global settingfp8_computeininvokeai.yaml, defaultfalse, plusfp8_compute_full_precision_hints(defaulttrue).Covers both kinds of checkpoint in the wild — scaled FP8 (FP8 weight plus
weight_scale, the ComfyUI convention) and raw FP8 (no scale at all) — wired into FLUX.1, FLUX.2, Anima, Krea-2, Z-Image and the Qwen3-VL and Mistral encoders. The upstream description carries the full rationale and the measurements; this PR does not restate them.Not in scope, unchanged from upstream: Qwen-Image and SDXL. The v7-only loaders (MiniMax H3, Wan, CogView4, ErnieImage, Ideogram 4) are likewise untouched here.
Related Issues / Discussions
main, so nothing else had to come along.QA Instructions
This is a merge of the upstream PR branch rather than a copied diff, so that git already knows these commits when invoke-ai#9478 lands upstream and
mainis synced.Net effect against
mainis exactly the upstream PR: 33 files, +5103/-271, no unrelated content.Two conflicts, both resolved:
krea2.py— import block only: this repo'snormalize_qwen3vl_rope_configagainst the PR'sfp8_scaledimports. Both kept; both are used (5 and 12 references respectively).invokeai/frontend/web/openapi.json— regenerated rather than hand-merged, withpnpm typegenforschema.ts. Both new settings reach the generated schema.config_default.pyauto-merged cleanly with this repo's own options intact.Checks run locally (Windows, RTX 4090, diffusers 0.40.0):
pytest tests/backend/quantization/ tests/backend/model_manager/load/ tests/test_config.py tests/backend/util/test_fp8.py— 965 passed, 150 skipped, 0 failedruff checkandruff format --check— cleanNot run: the full test suite, left to CI. Not run: the runtime QA from the upstream PR — no FP8 checkpoint was generated with on this machine, so the speed and VRAM claims are inherited from upstream, not reproduced here.
Review
Agent-assisted. Independent review subagents were not run; this had a self-review pass only, which is a weaker guarantee and is disclosed here rather than implied. The conflict resolutions and the regenerated artifacts are the parts worth a human eye.
Compatibility / Rollout
Both new settings default to the previous behaviour (
fp8_compute: false), so nothing changes until a user opts in.openapi.jsonandschema.tsare regenerated and additive.Checklist
What's Newcopy (if doing a release after this PR)