Merge main into ja: apply fp8-resident transformer fix to Japanese space

app.py

@@ -3,6 +3,7 @@ import gc
 import time
 import threading
 import traceback
+import types
 
 # cudaMallocAsync bypasses NVML memory queries that fail on MIG GPU instances
 os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "backend:cudaMallocAsync")
@@ -73,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)
@@ -94,16 +96,90 @@ def _start_heartbeat(label: str) -> threading.Event:
     return done
 
 
+_FP8_DTYPES = (torch.float8_e4m3fn, torch.float8_e5m2)
+
+
+def _fp8_upcast_linear_forward(self, input):
+    weight = self.weight.to(input.dtype) if self.weight.dtype in _FP8_DTYPES else self.weight
+    bias = self.bias.to(input.dtype) if (self.bias is not None and self.bias.dtype in _FP8_DTYPES) else self.bias
+    return torch.nn.functional.linear(input, weight, bias)
+
+
+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).
+    # 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
+
+
 _t0_load = time.perf_counter()
 print("[startup] loading transformer from_pretrained (prithivMLmods/Qwen-Image-Edit-Rapid-AIO-V23)...", flush=True)
 _hb = _start_heartbeat("transformer")
 _transformer = QwenImageTransformer2DModel.from_pretrained(
     "prithivMLmods/Qwen-Image-Edit-Rapid-AIO-V23",
-    torch_dtype=dtype,
+    torch_dtype=torch.float8_e4m3fn,
     device_map="cpu",
 )
 _hb.set()
 print(f"[startup] transformer loaded in {time.perf_counter()-_t0_load:.1f}s", 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:
+    print(f"[startup] transformer memory footprint: unavailable ({e})", flush=True)
 
 _t1_load = time.perf_counter()
 print("[startup] loading pipeline from_pretrained (FireRedTeam/FireRed-Image-Edit-1.1)...", flush=True)
@@ -387,45 +463,113 @@ def infer(images_b64_json, prompt, seed, randomize_seed, guidance_scale, steps,
         raise
 
 
-@spaces.GPU(duration=lambda *a, **kw: int(a[8]) if len(a) > 8 else 60)
-def _infer_gpu(pil_images, prompt, seed, guidance_scale, steps, width, height, mode, gpu_duration=20):
-    _cuda_ok = torch.cuda.is_available()
-    timer = _InferTimer(_cuda_ok)
-    t0 = time.perf_counter()
-
+def _log_infer_start(prompt, steps, guidance_scale, seed, gpu_duration, mode):
     print(f"[infer] ===== START =====")
     print(f"[infer] steps={steps}, guidance={guidance_scale}, seed={seed}, gpu_duration={gpu_duration}s, mode={mode}")
     print(f"[infer] prompt={repr(prompt[:120])}")
 
-    if _cuda_ok:
-        p = torch.cuda.get_device_properties(0)
-        print(f"[infer] GPU: {p.name}, total={p.total_memory/1024**3:.1f}GB, cap={p.major}.{p.minor}")
-        torch.cuda.reset_peak_memory_stats()
-
-    # Sequential offload: only the component currently doing work (text_encoder,
-    # then transformer, then vae — see model_cpu_offload_seq) sits on GPU at a
-    # time, instead of the full ~37GB pipeline resident simultaneously. Needed
-    # to fit on smaller MIG slices (e.g. the 47GB 2g.48gb partition) without OOM.
-    if getattr(pipe.transformer, "_hf_hook", None) is None:
-        pipe.enable_model_cpu_offload(device=device)
-        print(f"[infer] enabled sequential cpu offload on {device}")
-    print(f"[infer] {_gpu_mem_str(_cuda_ok)} — t={time.perf_counter()-t0:.1f}s")
-
-    print(f"[infer] {len(pil_images)} image(s) pre-decoded, output={width}x{height}, seed={seed}")
-
-    generator = torch.Generator(device=device).manual_seed(seed)
 
-    _step_t = []
+def _log_gpu_properties(cuda_ok):
+    if not cuda_ok:
+        return None
+    p = torch.cuda.get_device_properties(0)
+    print(f"[infer] GPU: {p.name}, total={p.total_memory/1024**3:.1f}GB, cap={p.major}.{p.minor}")
+    torch.cuda.reset_peak_memory_stats()
+    return p
+
+
+# Each ZeroGPU call runs in a fresh worker (hooks are always unset here), so
+# cpu_offload buys no cross-call reuse — it only trades one bulk to(device)
+# transfer for several slower hook-managed ones. The previous bf16-everywhere
+# 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_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.
+_FAST_PATH_MIN_GB = 40
+
+
+def _place_pipe_on_device(cuda_ok, gpu_props, t0):
+    if getattr(pipe.transformer, "_hf_hook", None) is not None:
+        return  # already placed by an earlier call sharing this worker
+    if cuda_ok and gpu_props.total_memory / 1024**3 >= _FAST_PATH_MIN_GB:
+        try:
+            pipe.to(device)
+            print(f"[infer] moved full pipe to {device} — t={time.perf_counter()-t0:.1f}s")
+            return
+        except torch.cuda.OutOfMemoryError:
+            print(f"[infer] OOM moving full pipe to {device}, falling back to cpu offload")
+            pipe.to("cpu")
+            torch.cuda.empty_cache()
+            pipe.enable_model_cpu_offload(device=device)
+            print(f"[infer] enabled cpu offload on {device} (fallback)")
+            return
+    pipe.enable_model_cpu_offload(device=device)
+    print(f"[infer] enabled cpu offload on {device} (slice too small for fast path)")
+
+
+def _instrument_first_touch(modules_with_names, t0):
+    """Install self-removing forward-pre-hooks that log the moment each module is first entered."""
+    def _make_hook(name, handle_box):
+        def _hook(mod, inputs):
+            print(f"[infer] first call into {name} — {_gpu_mem_str(True, sync=True)} | t={time.perf_counter()-t0:.1f}s")
+            handle_box["h"].remove()
+        return _hook
+    for module, name in modules_with_names:
+        handle_box = {}
+        handle_box["h"] = module.register_forward_pre_hook(_make_hook(name, handle_box))
+
+
+def _make_step_callback(steps, timer, t0):
+    """Build the diffusers step callback that logs per-step timing and marks timer checkpoints."""
+    step_times = []
     def _step_cb(pipeline, step_idx, timestep, cb_kwargs):
         now = time.perf_counter()
-        _step_t.append(now)
+        step_times.append(now)
         if step_idx == 0:
             timer.mark("first_step")
         timer.mark("last_step")  # overwritten each step; final value = end of last step
-        delta_ms = (now - (_step_t[-2] if len(_step_t) > 1 else t0)) * 1000
-        tag = " ← includes compile" if step_idx == 0 else ""
+        delta_ms = (now - (step_times[-2] if len(step_times) > 1 else t0)) * 1000
+        tag = " ← includes cold-start (offload hook install + first weight transfer)" if step_idx == 0 else ""
         print(f"[infer] step {step_idx+1}/{steps} done — {delta_ms:.0f}ms{tag} | t={now-t0:.1f}s")
         return cb_kwargs
+    return _step_cb
+
+
+def _log_infer_error(e, t0, timer):
+    print(f"[infer] ERROR: {type(e).__name__}: {e} | t={time.perf_counter()-t0:.1f}s")
+    print(traceback.format_exc())
+    try:
+        torch.cuda.synchronize()
+    except Exception as cuda_err:
+        print(f"[infer] CUDA synchronize after error: {cuda_err}")
+    timer.print_timings()
+
+
+@spaces.GPU(duration=lambda *a, **kw: int(a[8]) if len(a) > 8 else 60)
+def _infer_gpu(pil_images, prompt, seed, guidance_scale, steps, width, height, mode, gpu_duration=20):
+    _cuda_ok = torch.cuda.is_available()
+    timer = _InferTimer(_cuda_ok)
+    t0 = time.perf_counter()
+
+    _log_infer_start(prompt, steps, guidance_scale, seed, gpu_duration, mode)
+    gpu_props = _log_gpu_properties(_cuda_ok)
+
+    _place_pipe_on_device(_cuda_ok, gpu_props, t0)
+    print(f"[infer] {_gpu_mem_str(_cuda_ok)} — t={time.perf_counter()-t0:.1f}s")
+
+    if _cuda_ok:
+        _instrument_first_touch(
+            [(pipe.text_encoder, "text_encoder"), (pipe.transformer, "transformer"), (pipe.vae, "vae")],
+            t0,
+        )
+
+    print(f"[infer] {len(pil_images)} image(s) pre-decoded, output={width}x{height}, seed={seed}")
+
+    generator = torch.Generator(device=device).manual_seed(seed)
+    step_cb = _make_step_callback(steps, timer, t0)
 
     timer.mark("pipe_start")
     print(f"[infer] calling pipe... t={time.perf_counter()-t0:.1f}s")
@@ -434,7 +578,7 @@ def _infer_gpu(pil_images, prompt, seed, guidance_scale, steps, width, height, m
             image=pil_images, prompt=prompt, negative_prompt=negative_prompt,
             height=height, width=width, num_inference_steps=steps,
             generator=generator, true_cfg_scale=guidance_scale,
-            callback_on_step_end=_step_cb,
+            callback_on_step_end=step_cb,
             callback_on_step_end_tensor_inputs=["latents"],
         ).images[0]
         timer.mark("pipe_end")
@@ -443,18 +587,14 @@ def _infer_gpu(pil_images, prompt, seed, guidance_scale, steps, width, height, m
         duration = timer.elapsed_ms("pipe_start", "pipe_end") / 1000.0
         return result_image, seed, duration
     except Exception as e:
-        print(f"[infer] ERROR: {type(e).__name__}: {e} | t={time.perf_counter()-t0:.1f}s")
-        print(traceback.format_exc())
-        try:
-            torch.cuda.synchronize()
-        except Exception as cuda_err:
-            print(f"[infer] CUDA synchronize after error: {cuda_err}")
-        timer.print_timings()
+        _log_infer_error(e, t0, timer)
         raise
     finally:
-        # No manual pipe.to("cpu") — that fights the offload hooks' own device
-        # bookkeeping. Hooks already return each component to CPU after its
-        # forward; empty_cache() just reclaims the freed GPU blocks.
+        # No manual pipe.to("cpu"): in the offload-fallback case that fights the
+        # hooks' own device bookkeeping (they return each component to CPU after
+        # its forward), and in the normal fast-path case the worker's GPU access
+        # is reclaimed by ZeroGPU when this call returns regardless — paying for
+        # a D2H transfer here would just be wasted GPU-billed time.
         gc.collect()
         torch.cuda.empty_cache()
         print(f"[infer] ===== END t={time.perf_counter()-t0:.1f}s =====")

requirements.txt

@@ -1,6 +1,6 @@
-git+https://github.com/huggingface/accelerate.git
-git+https://github.com/huggingface/diffusers.git
-git+https://github.com/huggingface/peft.git
+accelerate==1.14.0
+diffusers==0.39.0
+peft==0.19.1
 transformers==4.57.1
 huggingface_hub
 pyarrow