From 4f00222c2ed27db0ffae2e977a672e1c86efeda2 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 11:48:11 +0000 Subject: [PATCH] Reserve the fp8 bf16 footprint even when the cache estimate is absent For a narrow (fp8) Ideogram-4 base the planner reserves the family's known bf16 component total, but the reservation was gated on model_dense_mib being non-None. On a first-time load an empty blob cache (or a best-effort download probe that swallowed a transient HF error) leaves model_dense_mib None, so the guard skipped the reservation exactly when it was needed: the planner then read 'size unknown -> stay resident' and the ~54 GB pipeline OOMed a card that offload would have fit. family_bf16_components_gb is a network-free constant, so reserve it whenever the cache signal is absent (use it directly when None, else take the max). --- studio/backend/core/inference/diffusion.py | 12 ++++++++++-- studio/backend/tests/test_diffusion_more_families.py | 12 ++++++++++++ 2 files changed, 22 insertions(+), 2 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 91a2acb699..25dc36bee2 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -1714,9 +1714,17 @@ class DiffusionBackend: is_narrow_base = ideogram4_repo_is_fp8(repo_id) if is_narrow_base: table = family_bf16_components_gb(fam, fam.base_repo) - if table is not None and model_dense_mib is not None: + if table is not None: + # family_bf16_components_gb is a network-free constant, so reserve the bf16 + # footprint even when the cache-derived estimate is absent (empty blob cache, + # or a best-effort download probe that swallowed a transient HF error and + # returned nothing). Otherwise model_dense_mib stays None and the planner + # reads "size unknown -> stay resident", so the ~54 GB fp8 pipeline plans a + # resident placement and OOMs a card that offload would have fit. table_mib = int(sum(table) * (1000.0**3) / (1024.0 * 1024.0)) - model_dense_mib = max(model_dense_mib, table_mib) + model_dense_mib = ( + table_mib if model_dense_mib is None else max(model_dense_mib, table_mib) + ) companion_mib = None else: if transformer_resident_override_mib is not None: diff --git a/studio/backend/tests/test_diffusion_more_families.py b/studio/backend/tests/test_diffusion_more_families.py index 9d89eb5ab3..b2ffcd093e 100644 --- a/studio/backend/tests/test_diffusion_more_families.py +++ b/studio/backend/tests/test_diffusion_more_families.py @@ -59,6 +59,18 @@ def test_ideogram4_generation_defaults(): assert default_generation_params("ideogram-ai/ideogram-4-fp8") == (48, 7.0) +def test_ideogram4_bf16_reservation_table_present(): + # The memory planner reserves this bf16 footprint for a narrow (fp8) ideogram-4 base even + # when the blob-cache estimate is absent (empty cache / a best-effort download probe that + # swallowed a transient HF error), so the ~54 GB pipeline never plans a resident placement + # it cannot fit. If this constant table ever went None, that fp8 OOM safeguard would + # silently disable, so pin that it is present and sums to the expected ~54 GB. + fam = detect_family("ideogram-ai/ideogram-4-fp8") + table = family_bf16_components_gb(fam, fam.base_repo) + assert table is not None + assert sum(table) > 50.0 # transformer (37.2) + bf16 text encoder (16.3) + VAE (0.2) + + def test_ideogram4_memory_table_counts_both_dits(): fam = detect_family("ideogram-ai/ideogram-4-fp8") components = family_bf16_components_gb(fam)