diff --git a/prain/config.py b/prain/config.py index d746288..4150da3 100644 --- a/prain/config.py +++ b/prain/config.py @@ -241,3 +241,64 @@ def load_config(path): _check_count(config.cache_slots, 0, f"{path}: cache_slots") _check_count(config.stats_every, 1, f"{path}: stats_every") return config + + +def moe_dims(hf_config): + """(num_layers, num_experts, top_k, hidden, moe_intermediate) of a HF config.json dict.""" + c = hf_config.get("text_config") or hf_config + + def first(*names): + for name in names: + if c.get(name) is not None: + return c[name] + raise KeyError(f"config.json has none of {names}") + + return ( + first("num_hidden_layers"), + first("num_experts", "num_local_experts", "n_routed_experts"), + first("top_k_experts", "num_experts_per_tok"), + first("hidden_size"), + first("moe_intermediate_size", "intermediate_size"), + ) + + +def expert_bytes(hidden, moe_intermediate, dtype_bytes=2): + """One gated expert: gate+up [2I, H] plus down [H, I].""" + return 3 * hidden * moe_intermediate * dtype_bytes + + +def render_toml(header, top, tables=None, layers=None): + """A placement file: top-level keys, [tables], then [layers."N"] sections (empty ones skipped).""" + lines = [f"# {header}"] + [f"{k} = {json.dumps(v)}" for k, v in top.items() if v is not None] + for name, values in (tables or {}).items(): + lines += ["", f"[{name}]"] + [f"{k} = {json.dumps(v)}" for k, v in values.items()] + for layer, values in sorted((layers or {}).items()): + if values: + lines += ["", f'[layers."{layer}"]'] + [f"{k} = {json.dumps(v)}" for k, v in values.items()] + return "\n".join(lines) + "\n" + + +def plan(stats, num_experts, expert_size, expert_vram_gib, ram_gib, cache_slots, pack_dir, stats_path): + """Hottest experts (by steps that routed to them) pinned in VRAM, next in RAM, rest on disk.""" + uses = {int(layer): s["uses"] for layer, s in stats["layers"].items()} + vram_rows = int(expert_vram_gib * GIB // expert_size) - cache_slots * len(uses) + if vram_rows < 0: + raise ValueError("expert VRAM budget is smaller than the cache slots alone") + ram_rows = int(ram_gib * GIB // expert_size) + ranked = sorted( + ((layer, e) for layer in uses for e in range(num_experts)), + key=lambda le: (-uses[le[0]][le[1]], le), + ) + tier = {le: "vram" if i < vram_rows else "ram" if i < vram_rows + ram_rows else "disk" + for i, le in enumerate(ranked)} + layers = { + layer: {name: format_ids(ids) for name in ("vram", "ram") + if (ids := [e for e in range(num_experts) if tier[(layer, e)] == name])} + for layer in uses + } + return render_toml( + "generated by `prain plan`; experts not listed live on disk", + {"pack_dir": pack_dir, "default_tier": "disk", "cache_slots": cache_slots, + "stats_path": stats_path}, + layers=layers, + ) diff --git a/tests/test_config.py b/tests/test_config.py index 8099645..0a23acc 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -119,3 +119,41 @@ mode = "auto" ('[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')) + + +class Plan(TmpCase): + def test_hottest_to_vram_then_ram_then_disk(self): + stats = {"layers": {"0": {"uses": [5, 0, 9, 1]}, "1": {"uses": [2, 7, 0, 0]}}} + # 1 GiB experts: 4 GiB VRAM = 2 cache slots + 2 pinned, 3 GiB RAM = 3 experts + text = mc.plan(stats, 4, mc.GIB, 4, 3, 1, "/packs", "/s.json") + self.assertEqual(text, """\ +# generated by `prain plan`; experts not listed live on disk +pack_dir = "/packs" +default_tier = "disk" +cache_slots = 1 +stats_path = "/s.json" + +[layers."0"] +vram = "2" +ram = "0,3" + +[layers."1"] +vram = "1" +ram = "0" +""") + p0 = mc.load_config(self.write("plan.toml", text)).placement(0, 4) + self.assertEqual((p0.vram, p0.ram, p0.disk), ((2,), (0, 3), (1,))) + + def test_budget_smaller_than_cache_slots(self): + with self.assertRaisesRegex(ValueError, "smaller than the cache slots"): + mc.plan({"layers": {"0": {"uses": [1]}}}, 1, mc.GIB, 1, 0, 2, "/p", "/s") + + def test_render_toml_round_trips_tables(self): + text = mc.render_toml("h", {"pack_dir": "/p", "stats_path": None}, + {"policy": {"eviction": "lru", "lfu_half_life": 8.0}}, + {3: {"cache_slots": 2}, 4: {}}) + config = mc.load_config(self.write("r.toml", text)) + self.assertEqual((config.policy.eviction, config.policy.lfu_half_life), ("lru", 8.0)) + self.assertEqual(config.placement(3, 4).cache_slots, 2) + self.assertNotIn("stats_path", text) + self.assertNotIn('"4"', text)