ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Reserve the fp8 bf16 footprint even when the cache estimate is absent

Browse source 
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
Download patch file Download diff file Expand all files Collapse all files

12
studio/backend/core/inference/diffusion.py
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: