feat(fp8): enable FP8 storage for Anima

[Pfannkuchensack]

Summary

Stacked on #9414 — depends on the extra_skip_patterns mechanism introduced there. Merge that one first. #9416 and #9478 stack on top of this one.

The fp8_storage toggle was shown for Anima main models but did nothing: AnimaCheckpointModel never called _apply_fp8_layerwise_casting. This wires it in — the state dict is already 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 at all. 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 reach it — they target diffusers' module names (norm, pos_embed, patch_embed, proj_in/out) and this architecture names the equivalent modules differently (x_embedder, final_layer, adaln_modulation_*).

AnimaTransformer now declares _skip_layerwise_casting_patterns, the same attribute diffusers models use, so the loader needs no special-casing.

Note this is the same module as Z-Image's t_embedder, broken through a different mechanism: there diffusers read weight.dtype and cast the input to float8; here the plain precision loss is enough.

Related Issues / Discussions

Follow-up to #8945 (FP8 storage), #9231 (hook-based casting) and the Z-Image PR this is stacked on.

QA Instructions

Needs a CUDA GPU. Model Manager → an Anima main model → Default Settings → enable FP8 Storage → Save.

1. Generate. Expect in the log:

   FP8 layerwise casting enabled for anima-preview (storage=float8_e4m3fn, compute=torch.bfloat16, param_size=2012MB)
   

   and the transformer resident at ~2012MB instead of 3988MB. On main there is no FP8 line at all — that is the dead toggle this PR fixes.

2. Quality vs. bf16. Note a fixed seed, then run once with FP8 on and once off. Two gotchas that will otherwise give you a false result:

   • Disable the invocation cache (PUT /api/v1/app/invocation_cache/disable), or the second run returns the first run's image unchanged and the two look pixel-identical.
   • Anima needs 35 steps / CFG 4.5; at 9 steps / CFG 1.0 both runs look mushy and the comparison says nothing.

   Expect the same composition with slightly coarser fine structure under FP8 — not a different image, and definitely not a dithered mess. A dithered result means the skip patterns are not being applied.

3. Regression: with FP8 off, output must be unchanged from before this PR.

Measured during development, same seed/steps/CFG each run — this is what pins the skip list down:

skip list	param_size	result
none	1994 MB	dithered, no fine detail ❌
t_embedder	2010 MB	clean ✅
+ x_embedder, final_layer	2012 MB	clean ✅
+ adaln_modulation	2180 MB	clean ✅

So t_embedder is necessary and sufficient. The two I/O layers are kept as ~2MB of margin, matching what diffusers skips by default for comparable DiTs. adaln_modulation is deliberately not listed — it costs 168MB and made no difference.

Unit tests:

uv run --extra cuda --extra test pytest tests/backend/model_manager/load/test_load_default_fp8.py tests/backend/anima -q --no-cov

Merge Plan

Merge after the Z-Image PR — _apply_fp8_to_nn_module(..., extra_skip_patterns=...) does not exist without it. No DB schema, no redux slice, no API schema change otherwise.

Checklist

• The PR has a short but descriptive title, suitable for a changelog
• Tests added / updated (if applicable)
• ❗Changes to a redux slice have a corresponding migration — n/a, no redux changes
• Documentation added / updated (if applicable) — n/a
• Updated What's New copy (if doing a release after this PR) — n/a

[lstein]

Adversarial review at 8dbf0a4b61 (reviewable delta = the 3 Anima files; the other 3 in the diff belong to #9414). I ran the attacks against a real anima-base-v1.0.safetensors on a W7900 — ROCm reports device.type == "cuda", so _device_supports_fp8_storage returns True and this whole path is live there too.

No correctness bug found. Approving.

Attacks attempted, and why each failed

• Compute-dtype poisoning. compute_dtype = next(model.parameters()).dtype resolves to t_embedder.1.linear_1.weight — which is present in the checkpoint and is a skipped module, so it stays bf16. The only strict=False missing keys are the 3 buffers, and init_empty_weights(include_buffers=False) leaves those as real CPU tensors, so no meta/float32 parameter can reach the probe.
• Unanchored regex over-match. re.search("t_embedder", …) would also match text_embedder/context_embedder; this architecture has neither. On a meta-device instantiation the three patterns match exactly 6 modules (t_embedder.1.linear_{1,2}, x_embedder.proj.1, final_layer.{linear,adaln_modulation.1,adaln_modulation.2}) and nothing else.
• Z-Image's failure mode (a weight.dtype read that casts the input). AnimaTransformer has no .dtype attribute, so anima_denoise.py:777's hasattr(transformer, "dtype") falls through to inference_dtype; nothing under invokeai/backend/anima/ touches .weight directly.
• Runtime integrations, all verified with real weights: fully-resident forward, CachedModelWithPartialLoad at 50% VRAM, an LLLite ControlNet bound (adapter demonstrably active — it moves the output by 5.7%), and LoRA patching, where _is_any_part_of_layer_fp8 correctly selects sidecar: zero param dtype changes during or after the patch context, model stays 2012 MB. Storage dtypes are restored to fp8 after every forward in every case.
• Cache staleness on toggling the setting. _LOAD_AFFECTING_SETTINGS eviction is base-agnostic, so Anima gets it for free.

The analysis in the description checks out

All four rows of your size table reproduce to within 0.1 MB (1994.2 / 2010.2 / 2012.0 / 2180.0). TimestepEmbedding.forward with use_adaln_lora=True returns (sample, emb), so t_emb is the raw sinusoidal embedding and those two Linears really do feed only adaln_lora — the comment is exactly right.

The case for t_embedder is in fact stronger than the PR states: 38.2% of its weights flush to zero under unscaled e4m3fn (next worst group is 24.9%), with a round-trip relative error of 0.068 vs 0.028–0.048 everywhere else. It is by a wide margin the most fp8-damaged module in the network. And x_embedder/final_layer are literally diffusers' CosmosTransformer3DModel._skip_layerwise_casting_patterns = ["patch_embed", "final_layer", "norm"], which is apt given Anima is the Cosmos-Predict2 DiT.

Three non-blocking findings

1. The fix has no regression guard. Deleting the model = self._apply_fp8_layerwise_casting(...) line from anima.py outright leaves all 1253 tests in tests/backend/model_manager + tests/backend/anima passing. The dead toggle this PR fixes can come straight back with CI green. (#9414 has the same gap for Z-Image.)

2. The new test doesn't pin the patterns to real module names. test_anima_transformer_declares_t_embedder_skip asserts a string is in a list, then re-tests the loader against a hand-built _Model; x_embedder and final_layer are never asserted at all. Renaming t_embedder in the transformer would silently disable the skip with tests green. Instantiating the real model under accelerate.init_empty_weights() takes ~2 s and pins all three to actual dotted module paths.

3. adaln_modulation "made no difference" isn't supported by measurement. Single-forward velocity error vs bf16, real checkpoint, same inputs each run:

skip list	param_size	rel. L2 vs bf16
none	1994.2 MB	0.15203
t_embedder	2010.2 MB	0.14378
PR list	2012.0 MB	0.13421
PR list + adaln_modulation	2180.0 MB	0.09115

adaln_modulation is the largest remaining error source, and skipping t_embedder alone moves the total only 0.152 → 0.144 — the errors add in quadrature and no single group dominates. The decision is defensible (168 MB for a modest gain), but the comment will be read as a measurement and currently tells the next maintainer the opposite. Suggest softening to "no visible difference in a 35-step A/B". Related nit: final_layer is annotated "output projection", but 1.57 of the 1.70 M params that entry protects are final_layer.adaln_modulation.* — i.e. most of what it shields is the thing the comment two lines down says is deliberately not shielded.

One caveat on my own numbers: I also ran a 25-step CFG-4.5 trajectory to test compounding, and it is chaotic under synthetic conditioning (rel ≈ 0.6 for every config, ordering meaningless). It can't settle the perceptual claim in either direction, so I'm not resting anything on it — your visual A/B remains the evidence for that part.

Adjacent, out of scope

fp8_storage is rendered for ControlNet configs (ControlAdapterModelDefaultSettings.tsx, everything except control_lora), but AnimaControlNetLLLiteModel._load_model never calls the cast — the same dead toggle, still live for Anima LLLite. Those adapters are 16–63 MB, so hiding the toggle is probably a better fix than wiring it.