From bb2b14db9795ec135e1c1b3d1744eab9ca1b1317 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 10:21:10 +0000 Subject: [PATCH] ideogram-4: build the FP8 text encoder at target dtype, size it as bf16-resident, optimize both DiTs - The FP8 Qwen3-VL text encoder was constructed at the process fp32 default before the dequantized bf16 weights are copied in. That ~8B-param fp32 scaffold peaks ~2x on host RAM (loading FIRST, before the DiTs), so a 64 GB host can OOM. Build it at the target dtype under set_default_dtype, mirroring the DiT loader; rotary inv_freq is still computed in explicit fp32. - The auto-policy memory table listed the text encoder at 8.8 GB, its FP8 on-disk size, while the DiTs were doubled to their bf16-resident sizes. The loader dequantizes the encoder to bf16 too (~16.3 GB), so the entry understated the resident footprint by ~7.5 GB and could let the planner pick a resident placement that OOMs. Size it as bf16-resident. - Speed (regional compile, QKV fuse) and the attention backend only touched pipe.transformer, so ideogram-4's second denoiser (unconditional_transformer, run every step for dual-branch CFG) stayed eager/native while status reported the optimization as engaged. Iterate every denoiser DiT (mirroring the offload path) so both experts are optimized. Guarded on attr presence, so single-DiT families are unchanged. --- .../core/inference/diffusion_attention.py | 50 +++++++++++------ .../core/inference/diffusion_auto_policy.py | 9 +-- .../core/inference/diffusion_ideogram4.py | 12 +++- .../backend/core/inference/diffusion_speed.py | 55 +++++++++++++++---- .../backend/tests/test_diffusion_attention.py | 11 ++++ studio/backend/tests/test_diffusion_speed.py | 24 ++++++++ 6 files changed, 127 insertions(+), 34 deletions(-) diff --git a/studio/backend/core/inference/diffusion_attention.py b/studio/backend/core/inference/diffusion_attention.py index 7377f8d919..6873026cd1 100644 --- a/studio/backend/core/inference/diffusion_attention.py +++ b/studio/backend/core/inference/diffusion_attention.py @@ -244,13 +244,27 @@ def _ensure_attention_backend_installed(backend: str, logger: Any = None) -> Non ) +def _attention_dits(pipe: Any) -> list: + """Every DiT the denoise loop runs each step: the primary ``transformer`` plus a second + expert some families carry (Ideogram's ``unconditional_transformer`` for its dual-branch + CFG, an MoE ``transformer_2``). The attention backend must be set on ALL of them, else the + second DiT keeps the native default while status reports the requested kernel as engaged.""" + dits: list = [] + for attr in ("transformer", "transformer_2", "unconditional_transformer"): + m = getattr(pipe, attr, None) + if m is not None and m not in dits: + dits.append(m) + return dits + + def apply_attention_backend( pipe: Any, backend: Optional[str], *, logger: Any = None, ) -> Optional[str]: - """Set ``backend`` on ``pipe.transformer`` via the diffusers dispatcher. + """Set ``backend`` on EVERY denoiser DiT (``pipe.transformer`` plus a second expert such as + Ideogram's ``unconditional_transformer``) via the diffusers dispatcher. Returns the backend actually engaged, or None when left at the native default (either because ``backend`` was None or because the requested kernel was unavailable -> graceful @@ -261,28 +275,32 @@ def apply_attention_backend( defaults to None). So a load that wants native must restore it explicitly: otherwise it silently inherits a backend an earlier load pinned (e.g. cuDNN under a speed profile), breaking the bit-identical/``off`` guarantee. Best-effort throughout.""" - transformer = getattr(pipe, "transformer", None) - fn = getattr(transformer, "set_attention_backend", None) - if not callable(fn): + setters = [s for s in (getattr(t, "set_attention_backend", None) for t in _attention_dits(pipe)) + if callable(s)] + if not setters: return None if backend is not None: _ensure_attention_backend_installed(backend, logger) - try: - fn(backend) - # set_attention_backend also pins the backend in diffusers' process-wide - # registry. This transformer's own processors keep it locally (their - # _attention_backend is now explicit), so reset the global default back to - # native -- otherwise a later component whose processors are unconfigured - # (backend None) silently inherits this kernel. + engaged = False + for fn in setters: + try: + fn(backend) + engaged = True + except Exception as exc: # noqa: BLE001 — unavailable kernel -> restore native below + _warn(logger, backend, exc) + if engaged: + # set_attention_backend also pins the backend in diffusers' process-wide registry. + # Each DiT's own processors now keep it locally (their _attention_backend is now + # explicit), so reset the global default back to native ONCE -- otherwise a later + # component whose processors are unconfigured (backend None) inherits this kernel. _reset_global_backend_to_native(logger) if logger is not None: logger.info("diffusion.attention: backend=%s", backend) return backend - except Exception as exc: # noqa: BLE001 — unavailable kernel -> restore native below - _warn(logger, backend, exc) - # No backend requested, or the requested one failed: pin the native default so a stale - # process-wide backend from a previous load can't leak into this one. - _restore_native_backend(fn, logger) + # No backend requested, or every set failed: pin the native default so a stale process-wide + # backend from a previous load can't leak into this one. Fresh DiTs follow the process-wide + # backend, so one reset via any DiT's setter covers them all. + _restore_native_backend(setters[0], logger) return None diff --git a/studio/backend/core/inference/diffusion_auto_policy.py b/studio/backend/core/inference/diffusion_auto_policy.py index 8433e4e2a2..85653364dd 100644 --- a/studio/backend/core/inference/diffusion_auto_policy.py +++ b/studio/backend/core/inference/diffusion_auto_policy.py @@ -61,10 +61,11 @@ _FAMILY_BF16_GB: dict[str, tuple[float, float, float]] = { "krea-2": (26.3, 8.9, 0.5), # Two ~9.3B DiTs (the conditional transformer PLUS the separate # unconditional_transformer driving Ideogram's dual-branch CFG), both resident - # for every generation, and a Qwen3-VL text encoder. The vendor repo stores the - # DiTs as raw float8 (9.29 GB each); these are the bf16-resident sizes after the - # loader's dtype cast, per this table's contract. - "ideogram-4": (37.2, 8.8, 0.2), + # for every generation, and a Qwen3-VL text encoder. The vendor repo stores the DiTs + # AND the text encoder as raw float8 (9.29 GB per DiT, 8.8 GB encoder); these are the + # bf16-resident sizes after the loader's dtype cast, per this table's contract, so the + # encoder doubles to ~16.3 GB just like each DiT (37.2 = 2 x 18.6). + "ideogram-4": (37.2, 16.3, 0.2), } # Base-repo overrides for families whose picker offers multiple sizes under one family diff --git a/studio/backend/core/inference/diffusion_ideogram4.py b/studio/backend/core/inference/diffusion_ideogram4.py index a4ac1dfec2..4e3e321604 100644 --- a/studio/backend/core/inference/diffusion_ideogram4.py +++ b/studio/backend/core/inference/diffusion_ideogram4.py @@ -308,8 +308,16 @@ def load_ideogram4_text_encoder( # Construct normally (so __init__ computes the non-persistent rotary inv_freq # buffers the checkpoint omits) then copy the dequantized weights in with - # assign=False. Host RAM is ample, so the transient dense init is fine. - model = Qwen3VLModel(config).to(dtype) + # assign=False. Build at the target dtype (mirrors the DiT loader below): this ~8B-param + # Qwen3-VL scaffold is ~2x at the process fp32 default (~33 GB vs ~16 GB) and loads FIRST + # on host RAM, so the fp32 transient can OOM a 64 GB host. rotary inv_freq is computed in + # explicit fp32 in __init__, so a bf16 default leaves it correct. + default_dtype = torch.get_default_dtype() + torch.set_default_dtype(dtype) + try: + model = Qwen3VLModel(config).to(dtype) + finally: + torch.set_default_dtype(default_dtype) missing, unexpected = model.load_state_dict(state_dict, strict = False) real_missing = [k for k in missing if not k.endswith("inv_freq")] if real_missing or unexpected: diff --git a/studio/backend/core/inference/diffusion_speed.py b/studio/backend/core/inference/diffusion_speed.py index 8c2524dfda..1f9a736f2f 100644 --- a/studio/backend/core/inference/diffusion_speed.py +++ b/studio/backend/core/inference/diffusion_speed.py @@ -272,6 +272,20 @@ def _vae_channels_last(pipe: Any, logger: Any) -> bool: return False +def _denoiser_dits(pipe: Any) -> list: + """Every DiT the denoise loop runs each step: the primary ``transformer`` plus a second + expert some families carry (Ideogram's ``unconditional_transformer`` for its dual-branch + CFG, an MoE ``transformer_2``). Speed / attention optims must reach ALL of them -- mirroring + the offload path (diffusion_memory streams the same set) -- else the second DiT runs + eager / native for every generation while status over-reports the optimisation as engaged.""" + dits: list = [] + for attr in ("transformer", "transformer_2", "unconditional_transformer"): + m = getattr(pipe, attr, None) + if m is not None and m not in dits: + dits.append(m) + return dits + + def _compile_repeated_blocks( pipe: Any, logger: Any, @@ -280,9 +294,8 @@ def _compile_repeated_blocks( cache_active: bool = False, offload_active: bool = False, ) -> bool: - transformer = getattr(pipe, "transformer", None) - fn = getattr(transformer, "compile_repeated_blocks", None) - if not callable(fn): + dits = [t for t in _denoiser_dits(pipe) if callable(getattr(t, "compile_repeated_blocks", None))] + if not dits: return False # default: mode="default" + dynamic=True -- fast cold start, robust to resolution # changes (no recompile). max: mode="max-autotune-no-cudagraphs" + dynamic=False -- @@ -319,11 +332,19 @@ def _compile_repeated_blocks( for _limit_attr in ("recompile_limit", "cache_size_limit"): # name varies by torch ver if hasattr(dynamo_cfg, _limit_attr): setattr(dynamo_cfg, _limit_attr, max(getattr(dynamo_cfg, _limit_attr) or 0, 64)) - fn(**kwargs) - return True except Exception as exc: # noqa: BLE001 — optimisation only _warn(logger, "compile_repeated_blocks", exc) return False + # Compile every denoiser DiT (a dual-DiT family such as Ideogram runs both each step); a + # per-DiT failure degrades that one to eager without dropping the others. + engaged = False + for transformer in dits: + try: + transformer.compile_repeated_blocks(**kwargs) + engaged = True + except Exception as exc: # noqa: BLE001 — optimisation only + _warn(logger, "compile_repeated_blocks", exc) + return engaged def _enable_cudnn_benchmark(logger: Any) -> bool: @@ -399,16 +420,26 @@ def _enable_fp16_accumulation( def _fuse_qkv(pipe: Any, logger: Any) -> bool: - for owner in (pipe, getattr(pipe, "transformer", None)): - fn = getattr(owner, "fuse_qkv_projections", None) - if callable(fn): + # Prefer the pipe-level fuse (it covers every component the pipe knows about); else fuse each + # denoiser DiT directly so a dual-DiT family (Ideogram) fuses BOTH experts, not just the first. + fn = getattr(pipe, "fuse_qkv_projections", None) + if callable(fn): + try: + fn() + return True + except Exception as exc: # noqa: BLE001 — optimisation only + _warn(logger, "fuse_qkv_projections", exc) + return False + engaged = False + for transformer in _denoiser_dits(pipe): + tfn = getattr(transformer, "fuse_qkv_projections", None) + if callable(tfn): try: - fn() - return True + tfn() + engaged = True except Exception as exc: # noqa: BLE001 — optimisation only _warn(logger, "fuse_qkv_projections", exc) - return False - return False + return engaged def _warn(logger: Any, what: str, exc: Exception) -> None: diff --git a/studio/backend/tests/test_diffusion_attention.py b/studio/backend/tests/test_diffusion_attention.py index 5abaa86fa7..aa8ccd8229 100644 --- a/studio/backend/tests/test_diffusion_attention.py +++ b/studio/backend/tests/test_diffusion_attention.py @@ -177,6 +177,17 @@ def test_apply_sets_backend(): assert engaged == "_native_cudnn" and t.set_to == "_native_cudnn" +def test_apply_sets_backend_on_both_dits(): + # A dual-DiT family (Ideogram) runs transformer + unconditional_transformer each step, so the + # backend must be set on BOTH; otherwise the second DiT keeps the native default while status + # reports the requested kernel as engaged. + t1, t2 = _FakeTransformer(), _FakeTransformer() + pipe = types.SimpleNamespace(transformer = t1, unconditional_transformer = t2) + engaged = apply_attention_backend(pipe, "_native_cudnn") + assert engaged == "_native_cudnn" + assert t1.set_to == "_native_cudnn" and t2.set_to == "_native_cudnn" + + def test_apply_falls_back_on_unavailable_kernel(monkeypatch): # an unavailable kernel must not fail the load -> returns None (diffusers default). monkeypatch.setattr(att, "_active_attention_backend", lambda: "native") diff --git a/studio/backend/tests/test_diffusion_speed.py b/studio/backend/tests/test_diffusion_speed.py index 9a18090e36..28cd4e0416 100644 --- a/studio/backend/tests/test_diffusion_speed.py +++ b/studio/backend/tests/test_diffusion_speed.py @@ -177,6 +177,7 @@ class _Pipe: *, with_compile = False, with_fuse = False, + with_second_dit = False, ) -> None: self.vae = types.SimpleNamespace(mem_format = None, to = self._vae_to) self.transformer = types.SimpleNamespace() @@ -186,6 +187,12 @@ class _Pipe: self.fuse_qkv_projections = self._fuse self.compiled = False self.fused = False + # A dual-DiT family (Ideogram) carries a second denoiser expert that runs every step. + self.second_compiled = False + if with_second_dit: + self.unconditional_transformer = types.SimpleNamespace() + if with_compile: + self.unconditional_transformer.compile_repeated_blocks = self._compile2 def _vae_to(self, *, memory_format): self.vae.mem_format = memory_format @@ -194,6 +201,9 @@ class _Pipe: self.compiled = True self.compile_kwargs = kwargs + def _compile2(self, **kwargs): + self.second_compiled = True + def _fuse(self): self.fused = True @@ -218,6 +228,20 @@ def test_speed_off_applies_nothing(monkeypatch): assert torch.backends.cudnn.benchmark is False +def test_speed_compiles_both_dits_for_dual_dit_family(monkeypatch): + # A dual-DiT family (Ideogram: transformer + unconditional_transformer) runs BOTH DiTs each + # denoise step, so the regional block compile must engage on both, not just the first -- + # otherwise the second DiT runs eager while status reports compile as engaged. + _stub_torch(monkeypatch) + _stub_gguf_accel(monkeypatch) + pipe = _Pipe(with_compile = True, with_second_dit = True) + applied = apply_speed_optims( + pipe, _target(), is_gguf = False, family = _family(), speed_mode = SPEED_DEFAULT + ) + assert applied["compiled"] is True + assert pipe.compiled is True and pipe.second_compiled is True + + def test_speed_default_dense_falls_back_to_regional_compile(monkeypatch): # A DENSE model has no GGUF dequant to compile, so `default` falls back to the # regional block compile (its only compile lever) -- and no GGUF accelerators.