refuse a quantized model the card cannot hold, instead of proving it slowly

_reject_offloaded_quantized already catches a 4-bit model whose weights landed
on the CPU -- that is what stopped Qwen3.5-9B answering every request with "".
But it runs at the BOTTOM of the device_map ladder. So a full-GPU attempt that
OOMs still walks every rung: 427 weight shards across GPU and host RAM, many
minutes of it, before raising an error that was knowable when the ladder
started. The line printed immediately above the fallback says as much --
"quantized kernels cannot read offloaded weights" -- and then falls through to
the ladder anyway.

The other half is that the attempt which failed was never retried honestly. It
ran against whatever happened to be free on a shared card, which is a lottery;
another model's residency is a reason to evict, not a reason to mis-place this
one. Reproduced on the live server today: 7.2 GB free because an embedding model
and a whisper-server held the rest, and nothing asked them to move.

So a quantized model that fails full-GPU now evicts and tries once more, and
refuses immediately if that still does not fit, naming the three things that
help. Unquantized models keep the ladder: offload is legitimate for dense
weights, only slower.

This makes a failure faster and louder rather than slower and quieter, which is
a behaviour change on a live path.
Co-Authored-By: 's avatarClaude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014T3rZnLhydX9QqUcfBbqLu
parent 998d6d97
......@@ -813,6 +813,45 @@ class NvidiaBackend(ModelBackend):
return self._get_vram_percentages_for_strategy(strategy, is_moe, total_vram_gb)
def _evict_for_quantized(self, model_name: str, need_gb: float) -> bool:
"""Free VRAM for a model that has to be fully resident. True if it helped.
A quantized model has no offload path worth taking, so its first full-GPU
attempt running against whatever happened to be free on a shared card is
a lottery. Losing it is a reason to evict and try again, not a reason to
mis-place the model. Never raises: the caller reports the real failure.
"""
import gc
import torch
def _free_gb() -> float:
try:
return torch.cuda.mem_get_info(0)[0] / 1e9
except Exception:
return 0.0
before = _free_gb()
try:
from codai.models.manager import multi_model_manager
# Same call ensure_vram_for makes, and it takes no lock of its own,
# so it is safe from inside a load the manager started. It skips a
# model that is busy serving rather than crashing its CUDA context.
multi_model_manager._evict_models_for_vram(float(need_gb or 0.0))
except Exception as exc:
print(f" (could not evict for {model_name}: {exc})")
return False
try:
gc.collect()
torch.cuda.empty_cache()
except Exception:
pass
after = _free_gb()
if after > before + 0.1:
print(f" freed {after - before:.1f} GB — {after:.1f} GB now free")
return True
return False
def _derive_cuda_device(self) -> str:
"""Derive the CUDA device string from global args.
......@@ -1169,6 +1208,45 @@ class NvidiaBackend(ModelBackend):
model = self._try_load_model(model_name, _fg_kwargs, self.device)
if model is not None:
print(f" ✓ Model loaded full-GPU on {cuda_device}")
elif _quantized:
# The ladder below is exactly what the message above says
# cannot work: quantized kernels cannot read offloaded
# weights, so a 4-bit model placed by device_map loads
# happily and then answers every request with "" and
# completion_tokens: 0. Falling through to it turns a
# transient shortage of VRAM into a model that is broken
# for as long as it stays resident, and says nothing.
#
# So: make room properly and try once more. The first
# attempt ran against whatever happened to be free, which
# on a shared card is a lottery — another model's
# residency is a reason to evict, not to mis-place this
# one.
print(" ✗ Full-GPU load OOMed — evicting other models "
"and retrying (a quantized model cannot be served "
"from CPU offload)")
if self._evict_for_quantized(model_name, expected_vram_gb):
model = self._try_load_model(model_name, _fg_kwargs,
self.device)
if model is not None:
print(f" ✓ Model loaded full-GPU on {cuda_device} "
f"after eviction")
else:
_free_gb = 0.0
try:
_free_gb = torch.cuda.mem_get_info(0)[0] / 1e9
except Exception:
pass
raise RuntimeError(
f"{model_name} is quantized and needs about "
f"{expected_vram_gb:.1f} GB of VRAM to itself; "
f"{_free_gb:.1f} GB is free after evicting what "
f"could be evicted. Serving it from CPU offload "
f"is not an option — quantized kernels cannot "
f"read offloaded weights, and the model would "
f"answer every request with an empty completion. "
f"Free the card, lower n_ctx, or configure this "
f"model unquantized.")
else:
print(" ✗ Full-GPU load OOMed — falling back to "
"device_map offload ladder")
......
......@@ -213,3 +213,94 @@ def test_an_unquantized_model_may_still_offload():
hf_device_map={"layers.0": 0, "layers.1": "cpu"},
config=types.SimpleNamespace(quantization_config=None))
b._reject_offloaded_quantized(model, "m", {})
# ============================== A5, second half: refuse before the long walk
# _reject_offloaded_quantized catches a quantized model that ended up on the
# CPU -- but it runs at the BOTTOM of the device_map ladder, after every rung
# has been tried. For a 4-bit 9B that means loading 427 weight shards onto
# GPU+CPU over many minutes before raising, with the host RAM and the card
# churned for a result that was known to be useless when the ladder started:
# the line above it says so in as many words ("quantized kernels cannot read
# offloaded weights"). A shortage of VRAM is also not a reason to give up
# before asking the other models on the card to move.
def test_a_quantized_oom_evicts_and_retries_instead_of_walking_the_ladder():
"""Asserted on the source, like the ladder guarantee above, so it survives
a refactor of the branch."""
import inspect
from codai.backends import cuda
src = inspect.getsource(cuda.NvidiaBackend.load_model)
head = src[src.index("Full-GPU load OOMed"):]
# The quantized branch comes first and never reaches the ladder message.
quant = head[:head.index("falling back to")]
assert "_evict_for_quantized" in quant
assert "raise RuntimeError" in quant
# And an unquantized model still has the ladder, which is real recovery.
assert "falling back to" in head
def test_the_refusal_says_what_to_do_about_it():
import inspect
from codai.backends import cuda
src = inspect.getsource(cuda.NvidiaBackend.load_model)
msg = src[src.index("is quantized and needs about"):]
msg = msg[:msg.index('")') + 2]
for hint in ("empty completion", "n_ctx", "unquantized"):
assert hint in msg, hint
class _FakeMem:
"""Stands in for torch.cuda.mem_get_info: a list of successive free bytes."""
def __init__(self, *free_gb):
self.vals = [int(g * 1e9) for g in free_gb]
def __call__(self, _i=0):
v = self.vals[0] if len(self.vals) == 1 else self.vals.pop(0)
return (v, int(24 * 1e9))
def _backend():
from codai.backends.cuda import NvidiaBackend
return object.__new__(NvidiaBackend)
def test_eviction_that_freed_room_reports_success(monkeypatch):
import sys
import types
import torch
freed = []
mgr = types.SimpleNamespace(
_evict_models_for_vram=lambda gb: freed.append(gb))
mod = types.ModuleType("codai.models.manager")
mod.multi_model_manager = mgr
monkeypatch.setitem(sys.modules, "codai.models.manager", mod)
monkeypatch.setattr(torch.cuda, "mem_get_info", _FakeMem(7.2, 19.0))
monkeypatch.setattr(torch.cuda, "empty_cache", lambda: None)
assert _backend()._evict_for_quantized("m", 12.9) is True
assert freed == [12.9]
def test_eviction_that_freed_nothing_reports_failure(monkeypatch):
import sys
import types
import torch
mod = types.ModuleType("codai.models.manager")
mod.multi_model_manager = types.SimpleNamespace(
_evict_models_for_vram=lambda gb: None)
monkeypatch.setitem(sys.modules, "codai.models.manager", mod)
monkeypatch.setattr(torch.cuda, "mem_get_info", _FakeMem(7.2, 7.2))
monkeypatch.setattr(torch.cuda, "empty_cache", lambda: None)
assert _backend()._evict_for_quantized("m", 12.9) is False
def test_eviction_never_raises_when_the_manager_is_not_there(monkeypatch):
import sys
import types
import torch
broken = types.ModuleType("codai.models.manager") # no multi_model_manager
monkeypatch.setitem(sys.modules, "codai.models.manager", broken)
monkeypatch.setattr(torch.cuda, "mem_get_info", _FakeMem(7.2))
monkeypatch.setattr(torch.cuda, "empty_cache", lambda: None)
assert _backend()._evict_for_quantized("m", 12.9) is False
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment