prain: placement and option tables
This commit is contained in:
+184
@@ -57,3 +57,187 @@ def layer_index(name):
|
||||
"""'language_model.model.layers.12.experts' -> 12, or None unless exactly one number."""
|
||||
numbers = [int(part) for part in name.split(".") if part.isdigit()]
|
||||
return numbers[0] if len(numbers) == 1 else None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class LayerPlacement:
|
||||
vram: tuple # pinned experts: slot i holds vram[i] for the whole run
|
||||
ram: tuple # pinned host memory, copied into a cache slot on demand
|
||||
disk: tuple # read from the pack file on demand
|
||||
cache_slots: int # VRAM slots for ram/disk experts, managed by [policy]
|
||||
|
||||
@property
|
||||
def num_slots(self):
|
||||
return len(self.vram) + self.cache_slots
|
||||
|
||||
|
||||
def _opt(default, choices=None, minimum=0, maximum=None):
|
||||
return field(default=default, metadata={"choices": choices, "min": minimum, "max": maximum})
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Policy:
|
||||
"""[policy]: how the VRAM cache slots of every layer are managed."""
|
||||
eviction: str = _opt("lfu", ("lru", "lfu", "gdsf"))
|
||||
heat: str = _opt("weight", ("count", "weight", "recency"))
|
||||
lfu_half_life: float = _opt(64.0, minimum=1) # forwards until an LFU score halves
|
||||
disk_cost: float = _opt(4.0, minimum=1) # gdsf: refetching a disk expert vs. a RAM one
|
||||
recency_half_life: float = _opt(32.0, minimum=1) # heat "recency": tokens per halving
|
||||
scan_tokens: int = _opt(64, minimum=1) # forwards with this many tokens are scans
|
||||
scan_keep: int = _opt(0) # cache slots a scan may not evict
|
||||
promote: int = _opt(2) # experts a scan admits at full priority
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Prefetch:
|
||||
"""[prefetch]: copy experts before their layer runs."""
|
||||
mode: str = _opt("none", ("none", "pregate"))
|
||||
distance: int = _opt(1, minimum=1, maximum=2) # layers ahead
|
||||
budget: int = _opt(2) # experts per layer and forward
|
||||
min_prob: float = _opt(0.05, maximum=1) # ignore predictions below this router probability
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Miss:
|
||||
"""[miss]: what happens to an expert that is not in VRAM."""
|
||||
mode: str = _opt("copy", ("copy", "cpu", "auto"))
|
||||
cpu_max_tokens: int = _opt(1, minimum=1) # auto: CPU runs experts with <= this many tokens
|
||||
cpu_threads: int = _opt(4, minimum=1) # OMP threads of the vLLM worker
|
||||
admit: int = _opt(1) # CPU-run experts copied into VRAM afterwards, hottest first
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Serve:
|
||||
"""[serve]: the vLLM settings prain and its clients both need to know."""
|
||||
max_requests: int = _opt(2, minimum=1) # requests vLLM runs at once (--max-num-seqs)
|
||||
kv_cache_gib: float = _opt(0.75, minimum=0.01) # KV cache, shared by all requests
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Draft:
|
||||
"""[draft]: answers that arrive while other answers already use the GPU."""
|
||||
full_requests: int = _opt(1, minimum=1) # answers that run with all their experts at once
|
||||
min_weight: float = _opt(0.5, maximum=1) # router weight a draft token must still compute
|
||||
refine: str = _opt("regenerate", ("regenerate", "none")) # what happens to a draft later
|
||||
|
||||
|
||||
def _section(cls, raw, where):
|
||||
if not isinstance(raw, dict):
|
||||
raise ValueError(f"{where} must be a table")
|
||||
spec = {f.name: f for f in fields(cls)}
|
||||
_check_keys(raw, set(spec), where)
|
||||
for name, value in raw.items():
|
||||
meta, default = spec[name].metadata, spec[name].default
|
||||
if meta["choices"]:
|
||||
if value not in meta["choices"]:
|
||||
raise ValueError(f"{where} {name} must be one of {meta['choices']}, got {value!r}")
|
||||
continue
|
||||
number = type(value) is int or (type(default) is float and type(value) is float)
|
||||
top = meta["max"]
|
||||
if not number or value < meta["min"] or (top is not None and value > top):
|
||||
raise ValueError(f"{where} {name} must be a number in "
|
||||
f"[{meta['min']}, {top if top is not None else 'inf'}], got {value!r}")
|
||||
return cls(**raw)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Config:
|
||||
pack_dir: str
|
||||
default_tier: str = "ram"
|
||||
cache_slots: int = 16
|
||||
stats_path: str | None = None
|
||||
stats_every: int = 256
|
||||
trace_path: str | None = None
|
||||
policy: Policy = Policy()
|
||||
prefetch: Prefetch = Prefetch()
|
||||
miss: Miss = Miss()
|
||||
serve: Serve = Serve()
|
||||
draft: Draft = Draft()
|
||||
layers: tuple = () # ((selector, section), ...) in file order; later sections win
|
||||
|
||||
def placement(self, layer, num_experts):
|
||||
tier = [self.default_tier] * num_experts
|
||||
cache_slots = self.cache_slots
|
||||
for selector, section in self.layers:
|
||||
if selector != "*" and layer not in parse_ids(selector):
|
||||
continue
|
||||
seen = {}
|
||||
for name in TIERS:
|
||||
for expert in parse_ids(section.get(name, []), num_experts):
|
||||
if expert in seen:
|
||||
raise ValueError(
|
||||
f"[layers.\"{selector}\"] puts expert {expert} in both "
|
||||
f"{seen[expert]} and {name}"
|
||||
)
|
||||
seen[expert] = tier[expert] = name
|
||||
cache_slots = section.get("cache_slots", cache_slots)
|
||||
vram, ram, disk = (
|
||||
tuple(e for e in range(num_experts) if tier[e] == name) for name in TIERS
|
||||
)
|
||||
if (ram or disk) and cache_slots < 1:
|
||||
raise ValueError(
|
||||
f"layer {layer}: cache_slots must be >= 1 while experts live in ram or disk"
|
||||
)
|
||||
if (ram or disk) and self.policy.scan_keep >= cache_slots:
|
||||
raise ValueError(f"layer {layer}: policy.scan_keep must be < cache_slots ({cache_slots})")
|
||||
return LayerPlacement(vram, ram, disk, cache_slots)
|
||||
|
||||
|
||||
def _check_keys(table, allowed, where):
|
||||
unknown = set(table) - allowed
|
||||
if unknown:
|
||||
raise ValueError(f"{where}: unknown keys {sorted(unknown)}, allowed: {sorted(allowed)}")
|
||||
|
||||
|
||||
def _check_count(value, minimum, where):
|
||||
if type(value) is not int or value < minimum:
|
||||
raise ValueError(f"{where} must be an integer >= {minimum}, got {value!r}")
|
||||
|
||||
|
||||
def _path(config_path, value):
|
||||
"""A path from the file at config_path: '~' expands, relative paths start at that file's
|
||||
directory, so a config can point at '../models/...' from anywhere."""
|
||||
if value is None:
|
||||
return None
|
||||
return os.path.join(os.path.dirname(os.path.abspath(config_path)), os.path.expanduser(value))
|
||||
|
||||
|
||||
def load_config(path):
|
||||
with open(path, "rb") as f:
|
||||
raw = tomllib.load(f)
|
||||
_check_keys(raw, _TOP_KEYS, path)
|
||||
if "pack_dir" not in raw:
|
||||
raise ValueError(f"{path}: pack_dir is required")
|
||||
layers = []
|
||||
for selector, section in raw.get("layers", {}).items():
|
||||
where = f'{path} [layers."{selector}"]'
|
||||
if selector != "*":
|
||||
parse_ids(selector)
|
||||
if not isinstance(section, dict):
|
||||
raise ValueError(f"{where} must be a table")
|
||||
_check_keys(section, _LAYER_KEYS, where)
|
||||
if "cache_slots" in section:
|
||||
_check_count(section["cache_slots"], 0, f"{where} cache_slots")
|
||||
layers.append((selector, section))
|
||||
config = Config(
|
||||
pack_dir=_path(path, raw["pack_dir"]),
|
||||
default_tier=raw.get("default_tier", "ram"),
|
||||
cache_slots=raw.get("cache_slots", 16),
|
||||
stats_path=_path(path, raw.get("stats_path")),
|
||||
stats_every=raw.get("stats_every", 256),
|
||||
trace_path=_path(path, raw.get("trace_path")),
|
||||
policy=_section(Policy, raw.get("policy", {}), f"{path} [policy]"),
|
||||
prefetch=_section(Prefetch, raw.get("prefetch", {}), f"{path} [prefetch]"),
|
||||
miss=_section(Miss, raw.get("miss", {}), f"{path} [miss]"),
|
||||
serve=_section(Serve, raw.get("serve", {}), f"{path} [serve]"),
|
||||
draft=_section(Draft, raw.get("draft", {}), f"{path} [draft]"),
|
||||
layers=tuple(layers),
|
||||
)
|
||||
if config.draft.full_requests > config.serve.max_requests:
|
||||
raise ValueError(f"{path}: [draft] full_requests must be <= [serve] max_requests "
|
||||
f"({config.serve.max_requests})")
|
||||
if config.default_tier not in TIERS:
|
||||
raise ValueError(f"{path}: default_tier must be one of {TIERS}")
|
||||
_check_count(config.cache_slots, 0, f"{path}: cache_slots")
|
||||
_check_count(config.stats_every, 1, f"{path}: stats_every")
|
||||
return config
|
||||
|
||||
@@ -23,3 +23,99 @@ class ParseIds(unittest.TestCase):
|
||||
def test_layer_index(self):
|
||||
self.assertEqual(mc.layer_index("language_model.model.layers.12.experts"), 12)
|
||||
self.assertIsNone(mc.layer_index("model.experts"))
|
||||
|
||||
|
||||
class Placement(TmpCase):
|
||||
def test_sections_apply_in_file_order(self):
|
||||
config = mc.load_config(self.write("p.toml", """
|
||||
pack_dir = "/packs"
|
||||
default_tier = "ram"
|
||||
cache_slots = 4
|
||||
|
||||
[layers."*"]
|
||||
vram = "0-1"
|
||||
|
||||
[layers."2-3"]
|
||||
disk = "6-7"
|
||||
vram = "1-2"
|
||||
cache_slots = 2
|
||||
"""))
|
||||
p0 = config.placement(0, 8)
|
||||
self.assertEqual((p0.vram, p0.ram, p0.disk, p0.cache_slots),
|
||||
((0, 1), (2, 3, 4, 5, 6, 7), (), 4))
|
||||
p3 = config.placement(3, 8)
|
||||
self.assertEqual((p3.vram, p3.ram, p3.disk, p3.cache_slots),
|
||||
((0, 1, 2), (3, 4, 5), (6, 7), 2))
|
||||
self.assertEqual(p3.num_slots, 5)
|
||||
|
||||
def test_rejects_typos_conflicts_and_missing_slots(self):
|
||||
with self.assertRaisesRegex(ValueError, "unknown keys"):
|
||||
mc.load_config(self.write("a.toml", 'pack_dir = "/p"\ncache_slot = 3\n'))
|
||||
with self.assertRaisesRegex(ValueError, "default_tier"):
|
||||
mc.load_config(self.write("b.toml", 'pack_dir = "/p"\ndefault_tier = "gpu"\n'))
|
||||
both = mc.load_config(
|
||||
self.write("c.toml", 'pack_dir = "/p"\n[layers."*"]\nvram = "1"\ndisk = "1"\n'))
|
||||
with self.assertRaisesRegex(ValueError, "both vram and disk"):
|
||||
both.placement(0, 4)
|
||||
no_slots = mc.load_config(self.write("d.toml", 'pack_dir = "/p"\ncache_slots = 0\n'))
|
||||
with self.assertRaisesRegex(ValueError, "cache_slots must be >= 1"):
|
||||
no_slots.placement(0, 4)
|
||||
|
||||
def test_all_in_vram_needs_no_cache_slots(self):
|
||||
config = mc.load_config(
|
||||
self.write("v.toml", 'pack_dir = "/p"\ndefault_tier = "vram"\ncache_slots = 0\n'))
|
||||
self.assertEqual(config.placement(5, 4).num_slots, 4)
|
||||
|
||||
def test_policy_prefetch_and_miss_tables(self):
|
||||
config = mc.load_config(self.write("t.toml", """
|
||||
pack_dir = "/p"
|
||||
cache_slots = 3
|
||||
[policy]
|
||||
eviction = "gdsf"
|
||||
lfu_half_life = 16
|
||||
scan_keep = 1
|
||||
[prefetch]
|
||||
mode = "pregate"
|
||||
min_prob = 0.1
|
||||
[miss]
|
||||
mode = "auto"
|
||||
"""))
|
||||
self.assertEqual((config.policy.eviction, config.policy.lfu_half_life, config.policy.heat),
|
||||
("gdsf", 16, "weight"))
|
||||
self.assertEqual((config.prefetch.mode, config.prefetch.distance, config.miss.mode),
|
||||
("pregate", 1, "auto"))
|
||||
self.assertEqual(mc.load_config(self.write("u.toml", 'pack_dir = "/p"\n')).policy, mc.Policy())
|
||||
|
||||
def test_rejects_bad_table_values(self):
|
||||
for body, error in (('[policy]\neviction = "fifo"', "eviction must be one of"),
|
||||
('[policy]\nscan_keep = -1', "scan_keep must be a number"),
|
||||
('[prefetch]\ndistance = 3', "distance must be a number in \\[1, 2\\]"),
|
||||
('[prefetch]\nbudget = 1.5', "budget must be a number"),
|
||||
('[miss]\nmode = "gpu"', "mode must be one of"),
|
||||
('[miss]\nthreads = 4', "unknown keys")):
|
||||
with self.assertRaisesRegex(ValueError, error, msg=body):
|
||||
mc.load_config(self.write("bad.toml", f'pack_dir = "/p"\n{body}\n'))
|
||||
keep_all = mc.load_config(self.write("k.toml", 'pack_dir = "/p"\ncache_slots = 2\n'
|
||||
'[policy]\nscan_keep = 2\n'))
|
||||
with self.assertRaisesRegex(ValueError, "scan_keep must be < cache_slots"):
|
||||
keep_all.placement(0, 4)
|
||||
|
||||
def test_paths_are_relative_to_the_config_file(self):
|
||||
config = mc.load_config(self.write("r.toml", 'pack_dir = "packs/x"\nstats_path = "~/s.json"\n'
|
||||
'trace_path = "/abs/t.gz"\n'))
|
||||
self.assertEqual(config.pack_dir, os.path.join(self.dir, "packs/x"))
|
||||
self.assertEqual(config.stats_path, os.path.expanduser("~/s.json"))
|
||||
self.assertEqual(config.trace_path, "/abs/t.gz")
|
||||
|
||||
def test_serve_and_draft_tables(self):
|
||||
config = mc.load_config(self.write("s.toml", 'pack_dir = "/p"\n[serve]\nmax_requests = 4\n'
|
||||
'kv_cache_gib = 1.5\n[draft]\nfull_requests = 2\n'
|
||||
'refine = "none"\n'))
|
||||
self.assertEqual(config.serve, mc.Serve(max_requests=4, kv_cache_gib=1.5))
|
||||
self.assertEqual(config.draft, mc.Draft(full_requests=2, min_weight=0.5, refine="none"))
|
||||
self.assertEqual(mc.load_config(self.write("d.toml", 'pack_dir = "/p"\n')).serve.max_requests, 2)
|
||||
for body, error in (('[serve]\nmax_requests = 0', "max_requests must be a number"),
|
||||
('[draft]\nmin_weight = 1.5', "min_weight must be a number in \\[0, 1\\]"),
|
||||
('[draft]\nfull_requests = 3', "full_requests must be <= \\[serve\\] max_requests \\(2\\)")):
|
||||
with self.assertRaisesRegex(ValueError, error, msg=body):
|
||||
mc.load_config(self.write("bad.toml", f'pack_dir = "/p"\n{body}\n'))
|
||||
|
||||
Reference in New Issue
Block a user