Patch RMSNorm for fp8-resident weights; fix crash in prior fp8 commit

8db56cb kept the transformer's weights fp8-resident (instead of upcasting
to bf16 at load) to fix the memory OOM that was forcing this Space onto
the slow enable_model_cpu_offload path. It patched nn.Linear to upcast
its weight just-in-time for the matmul, but this model also uses RMSNorm
extensively (attn.norm_q/norm_k/norm_added_q/norm_added_k in every
transformer block, plus self.txt_norm) with real learnable weights —
confirmed via the checkpoint's safetensors headers, which show 100% of
tensors across all 3 shards as F8_E4M3, including norm gains and biases,
no mixed precision.

diffusers 0.39.0's RMSNorm.forward only upcasts its weight for
float16/bfloat16, not float8. Left unpatched, `hidden_states * self.weight`
multiplies a bf16 activation by an fp8 weight directly, which PyTorch
does not support via type promotion. Verified this reproduces exactly:

RuntimeError: Promotion for Float8 Types is not supported, attempted
to promote Float and Float8_e4m3fn

This adds _fp8_upcast_rmsnorm_forward, a copy of diffusers' RMSNorm.forward
with the same just-in-time upcast idea applied to RMSNorm's weight/bias,
and renames _patch_fp8_linears to _patch_fp8_modules to patch both module
types. Also adds a safety-net scan that logs a warning for any other
fp8-resident parameter neither patcher covers, so a future gap in this
allowlist surfaces at startup instead of mid-inference.

Verified locally (CPU, diffusers 0.39.0 installed in .venv): the
unpatched RMSNorm reproduces the RuntimeError above on an fp8 weight;
the patched version runs and matches a real-bf16-weight reference within
~0.003 max abs diff, consistent with fp8's own quantization noise rather
than a logic error. Still needs a real end-to-end run on the Space's GPU
before merging — this only verifies the patch in isolation, not the full
pipeline.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

app.py

@@ -74,6 +74,7 @@ print("[startup] importing dimensions...", flush=True)
 from dimensions import compute_output_dimensions, max_dim_for_mode
 print("[startup] importing diffusers...", flush=True)
 from diffusers import FlowMatchEulerDiscreteScheduler
+from diffusers.models.normalization import RMSNorm
 print("[startup] importing QwenImageEditPlusPipeline...", flush=True)
 from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
 print("[startup] importing QwenImageTransformer2DModel...", flush=True)
@@ -104,18 +105,62 @@ def _fp8_upcast_linear_forward(self, input):
     return torch.nn.functional.linear(input, weight, bias)
 
 
-def _patch_fp8_linears(model) -> int:
+def _fp8_upcast_rmsnorm_forward(self, hidden_states):
+    # Mirrors diffusers 0.39.0's RMSNorm.forward (CUDA path, models/normalization.py), extended
+    # so an fp8-resident weight/bias gets upcast to the activation's dtype before use instead of
+    # being silently skipped — the stock implementation only special-cases float16/bfloat16, so
+    # an fp8 weight would otherwise reach `hidden_states * self.weight` unconverted and error
+    # (no elementwise op supports bf16 x fp8 operands).
+    input_dtype = hidden_states.dtype
+    variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
+    hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
+
+    if self.weight is not None:
+        weight = self.weight.to(input_dtype) if self.weight.dtype in _FP8_DTYPES else self.weight
+        if weight.dtype in (torch.float16, torch.bfloat16):
+            hidden_states = hidden_states.to(weight.dtype)
+        hidden_states = hidden_states * weight
+        if self.bias is not None:
+            bias = self.bias.to(hidden_states.dtype) if self.bias.dtype in _FP8_DTYPES else self.bias
+            hidden_states = hidden_states + bias
+    else:
+        hidden_states = hidden_states.to(input_dtype)
+
+    return hidden_states
+
+
+def _patch_fp8_modules(model) -> int:
     # This checkpoint ships its weights natively in fp8 (torch_dtype below preserves that
     # instead of upcasting to bf16 at load time, halving resident memory: ~19GB vs ~38GB).
-    # nn.Linear has no fp8 compute kernel on this GPU, so each patched layer upcasts its own
-    # weight to the input's dtype just-in-time for the matmul — mathematically identical to
-    # the old load-time-upcast-everything approach (same values, same target dtype), just
-    # deferred so only one layer's weight is transiently bf16 at a time instead of all of them.
+    # Neither nn.Linear nor RMSNorm (the two module types in this model that own their own
+    # weight/bias, per the checkpoint's safetensors headers — every tensor is fp8, including
+    # norm gains) have an fp8 compute kernel on this GPU, so each patched instance upcasts its
+    # own weight to the input's dtype just-in-time for the op — mathematically identical to the
+    # old load-time-upcast-everything approach (same values, same target dtype), just deferred
+    # so only one layer's weight is transiently bf16 at a time instead of all of them.
     count = 0
     for module in model.modules():
         if isinstance(module, torch.nn.Linear) and module.weight.dtype in _FP8_DTYPES:
             module.forward = types.MethodType(_fp8_upcast_linear_forward, module)
             count += 1
+        elif isinstance(module, RMSNorm) and module.weight is not None and module.weight.dtype in _FP8_DTYPES:
+            module.forward = types.MethodType(_fp8_upcast_rmsnorm_forward, module)
+            count += 1
+
+    # Safety net: flag any other fp8-resident parameter that wasn't patched above, so a gap in
+    # this allowlist surfaces as a startup log line instead of a mid-inference crash — an
+    # unpatched fp8 parameter can't participate in ops with the bf16 activations around it.
+    patched_types = (torch.nn.Linear, RMSNorm)
+    for name, module in model.named_modules():
+        if isinstance(module, patched_types):
+            continue
+        for pname, param in module.named_parameters(recurse=False):
+            if param.dtype in _FP8_DTYPES:
+                print(
+                    f"[startup] WARNING: unpatched fp8 parameter {name}.{pname} "
+                    f"({type(module).__name__}) — will likely error at inference",
+                    flush=True,
+                )
     return count
 
 
@@ -129,8 +174,8 @@ _transformer = QwenImageTransformer2DModel.from_pretrained(
 )
 _hb.set()
 print(f"[startup] transformer loaded in {time.perf_counter()-_t0_load:.1f}s", flush=True)
-_n_fp8_patched = _patch_fp8_linears(_transformer)
-print(f"[startup] patched {_n_fp8_patched} fp8-resident nn.Linear modules for just-in-time upcast", flush=True)
+_n_fp8_patched = _patch_fp8_modules(_transformer)
+print(f"[startup] patched {_n_fp8_patched} fp8-resident nn.Linear/RMSNorm modules for just-in-time upcast", flush=True)
 try:
     print(f"[startup] transformer memory footprint: {_transformer.get_memory_footprint()/1024**3:.2f}GB", flush=True)
 except Exception as e:
@@ -439,7 +484,7 @@ def _log_gpu_properties(cuda_ok):
 # pipeline (~40GB transformer + ~14GB text encoder) peaked at 46.82GB moving
 # onto this Space's 47GB 2g.48gb MIG slice and still OOM'd, wasting ~40s
 # before falling back. The transformer now stays fp8-resident (see
-# _patch_fp8_linears above, ~19GB instead of ~38GB), so the full pipeline
+# _patch_fp8_modules above, ~19GB instead of ~38GB), so the full pipeline
 # should total roughly ~35GB — comfortably under the slice with headroom to
 # spare. Lowered accordingly, but the OOM fallback below stays as a safety
 # net in case that estimate is off.