File size: 5,735 Bytes
95200a7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 | """Asking a graphics card for more than it has is not an error on Windows.
The driver backs the excess with system memory, the allocation succeeds, the
server reports itself healthy, and then every token drags weights across the
PCIe bus. On the benchmark laptop -- a 4 GB card with 2.3 GB free and a 4.5 GB
model -- switching offloading on without a ceiling took twenty workloads that
had been completing to twenty that timed out on their first generation. Nothing
in the logs said "too big". It just went quiet.
So the ceiling is computed before anything is asked for, and these are the
properties of that computation.
"""
from __future__ import annotations
import struct
from distinct_agent.gguf import ModelShape, read_shape
from distinct_agent.server_runner import (
VRAM_BUDGET_SHARE,
VRAM_FIXED_OVERHEAD,
plan_offload,
)
GB = 1024**3
MB = 1024**2
#: OLMo-2 7B: 32 layers, 4096-wide, 32 heads with no grouped-query attention.
OLMO2_7B = ModelShape(layers=32, embedding_length=4096, heads=32, kv_heads=32)
def _plan(free_bytes, model_bytes=4_500_000_000, shape=OLMO2_7B, context=4096):
return plan_offload(
model_bytes=model_bytes,
shape=shape,
free_vram_bytes=free_bytes,
context_tokens=context,
)
# -- the ceiling -------------------------------------------------------------
def test_the_plan_never_asks_for_more_than_the_card_has_free() -> None:
"""The whole point. Every rung must fit in the budget it was given."""
for free in (1 * GB, 2 * GB, 4 * GB, 8 * GB, 24 * GB):
budget = int(free * VRAM_BUDGET_SHARE) - VRAM_FIXED_OVERHEAD
per_layer = 4_500_000_000 / (OLMO2_7B.layers + 1) + OLMO2_7B.kv_bytes_per_layer(4096)
for layers in _plan(free):
assert layers * per_layer <= max(0, budget) + per_layer
def test_a_laptop_card_with_a_seven_billion_model_offloads_some_of_it() -> None:
"""2.3 GB free is not nothing, and it is nowhere near the whole model."""
plan = _plan(2378 * MB)
assert 0 < plan[0] < OLMO2_7B.layers
def test_a_card_that_fits_the_whole_model_takes_the_whole_model() -> None:
assert _plan(24 * GB)[0] == OLMO2_7B.layers
def test_a_card_with_nothing_free_offloads_nothing() -> None:
assert _plan(200 * MB) == (0,)
def test_no_card_at_all_offloads_nothing() -> None:
assert _plan(0) == (0,)
def test_the_key_value_cache_is_counted_and_not_just_the_weights() -> None:
"""At 4096 tokens it can cost more per layer than the layer's weights."""
assert OLMO2_7B.kv_bytes_per_layer(4096) > 0
assert _plan(3 * GB, context=4096)[0] < _plan(3 * GB, context=512)[0]
def test_a_model_whose_header_cannot_be_read_offloads_nothing() -> None:
"""Unreadable is not an excuse to guess: it is the CPU, as it always was."""
assert _plan(8 * GB, shape=ModelShape()) == (0,)
# -- the ladder under the ceiling -------------------------------------------
def test_the_ladder_descends_and_ends_at_the_cpu() -> None:
plan = _plan(6 * GB)
assert list(plan) == sorted(plan, reverse=True)
assert plan[-1] == 0
def test_the_ladder_has_no_repeated_rungs() -> None:
for free in (400 * MB, 800 * MB, 1 * GB, 2 * GB, 3 * GB, 6 * GB, 24 * GB):
plan = _plan(free)
assert len(plan) == len(set(plan))
# -- reading the header ------------------------------------------------------
def _gguf(tmp_path, pairs, architecture="olmo2"):
"""A minimal but valid GGUF header carrying ``pairs`` as uint32 values."""
out = bytearray(b"GGUF")
out += struct.pack("<I", 3)
out += struct.pack("<Q", 0)
out += struct.pack("<Q", len(pairs))
for suffix, value in pairs.items():
key = f"{architecture}{suffix}".encode()
out += struct.pack("<Q", len(key)) + key
out += struct.pack("<I", 4) + struct.pack("<I", value)
path = tmp_path / "model.gguf"
path.write_bytes(bytes(out))
return path
def test_the_shape_is_read_out_of_the_file(tmp_path) -> None:
path = _gguf(
tmp_path,
{
".block_count": 32,
".embedding_length": 4096,
".attention.head_count": 32,
".attention.head_count_kv": 8,
},
)
shape = read_shape(path)
assert (shape.layers, shape.embedding_length, shape.heads, shape.kv_heads) == (32, 4096, 32, 8)
assert shape.head_dimension == 128
assert shape.known
def test_multi_head_attention_omits_the_kv_head_count_and_that_is_fine(tmp_path) -> None:
path = _gguf(
tmp_path,
{".block_count": 32, ".embedding_length": 4096, ".attention.head_count": 32},
)
assert read_shape(path).kv_heads == 32
def test_grouped_query_attention_needs_far_less_cache(tmp_path) -> None:
"""Eight key/value heads instead of thirty-two is four times less per layer."""
many = ModelShape(layers=32, embedding_length=4096, heads=32, kv_heads=32)
few = ModelShape(layers=32, embedding_length=4096, heads=32, kv_heads=8)
assert few.kv_bytes_per_layer(4096) * 4 == many.kv_bytes_per_layer(4096)
assert _plan(3 * GB, shape=few)[0] > _plan(3 * GB, shape=many)[0]
def test_a_file_that_is_not_a_gguf_reads_as_unknown_rather_than_raising(tmp_path) -> None:
path = tmp_path / "not-a-model.gguf"
path.write_bytes(b"this is not a model file at all")
assert read_shape(path).known is False
def test_a_truncated_header_reads_as_unknown_rather_than_raising(tmp_path) -> None:
path = tmp_path / "cut.gguf"
path.write_bytes(b"GGUF" + struct.pack("<I", 3) + b"\x00\x00")
assert read_shape(path).known is False
def test_a_missing_file_reads_as_unknown_rather_than_raising(tmp_path) -> None:
assert read_shape(tmp_path / "absent.gguf").known is False
|