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).
This commit is contained in:
Daniel Han 2026-07-06 11:48:11 +00:00
commit 4f00222c2e
2 changed files with 22 additions and 2 deletions

View file

@ -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:

View file

@ -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)