207 lines
11 KiB
Python
207 lines
11 KiB
Python
"""prain command line: check, plan or tune a placement, simulate it, serve with vLLM, and
|
|
chat in a TUI or the browser."""
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import sys
|
|
from dataclasses import asdict
|
|
|
|
from prain.config import (GIB, QUANTS, Miss, Policy, Prefetch, expert_bytes, format_ids, load_config,
|
|
moe_dims, plan, render_toml)
|
|
|
|
|
|
def _model_dims(hf_config_path, quant="bf16"):
|
|
with open(hf_config_path) as f:
|
|
num_layers, num_experts, top_k, hidden, inter = moe_dims(json.load(f))
|
|
return num_layers, num_experts, top_k, expert_bytes(hidden, inter, quant)
|
|
|
|
|
|
def check(config_path, hf_config_path):
|
|
config = load_config(config_path)
|
|
num_layers, num_experts, top_k, size = _model_dims(hf_config_path, config.quant)
|
|
print("layer vram ram disk cache_slots")
|
|
vram_rows = ram_rows = few_slots = 0
|
|
for layer in range(num_layers):
|
|
p = config.placement(layer, num_experts)
|
|
vram_rows += p.num_slots
|
|
ram_rows += len(p.ram)
|
|
few_slots += bool((p.ram or p.disk) and p.cache_slots < top_k)
|
|
print(f"{layer:5} {len(p.vram):4} {len(p.ram):4} {len(p.disk):4} {p.cache_slots:11}")
|
|
if few_slots:
|
|
print(f"warning: {few_slots} layers have fewer cache slots than top_k={top_k}; "
|
|
"their decode steps run in several rounds", file=sys.stderr)
|
|
print(f"expert size {size / 2**20:.2f} MiB")
|
|
print(f"VRAM {vram_rows * size / GIB:.2f} GiB (pinned experts + cache slots)")
|
|
print(f"KV {config.serve.kv_cache_gib:.2f} GiB (shared by up to "
|
|
f"{config.serve.max_requests} parallel requests)")
|
|
print(f"RAM {ram_rows * size / GIB:.2f} GiB (pinned host memory)")
|
|
print(f"disk {num_layers * num_experts * size / GIB:.2f} GiB (pack in {config.pack_dir})")
|
|
|
|
|
|
def _variant_text(variant):
|
|
policy, prefetch, miss = variant
|
|
fetch = "none" if prefetch.mode == "none" else f"pregate d{prefetch.distance} b{prefetch.budget}"
|
|
cpu = miss.mode if miss.mode == "copy" else f"{miss.mode} admit{miss.admit}"
|
|
if miss.mode == "auto":
|
|
cpu += f" <={miss.cpu_max_tokens}tok"
|
|
return (f"{policy.eviction:5} {policy.heat:8} keep{policy.scan_keep} "
|
|
f"promote{policy.promote} {fetch:15} {cpu}")
|
|
|
|
|
|
def sim(trace, config_path, hf_config_path, cost_path, expert_mib):
|
|
from prain.sim import describe, load_cost, read_forwards, simulate
|
|
|
|
config = load_config(config_path)
|
|
_, num_experts, _, size = _model_dims(hf_config_path, config.quant)
|
|
forwards = read_forwards(trace)
|
|
cost = load_cost(cost_path, expert_mib or size / 2**20)
|
|
layers = sorted({layer for fwd in forwards for layer in fwd})
|
|
placements = {layer: config.placement(layer, num_experts) for layer in layers}
|
|
variant = (config.policy, config.prefetch, config.miss)
|
|
result = simulate(forwards, placements, *variant, cost)
|
|
bound = simulate(forwards, placements, Policy(scan_tokens=config.policy.scan_tokens),
|
|
Prefetch(), Miss(), cost, belady=True)
|
|
print("\n".join(describe(result, bound, cost)))
|
|
|
|
|
|
def tune_command(args):
|
|
from prain.config import Miss, Policy, Prefetch
|
|
from prain.sim import describe, load_cost, read_forwards, tune
|
|
|
|
_, num_experts, _, size = _model_dims(args.hf_config, args.quant)
|
|
cost = load_cost(args.cost, args.expert_mib or size / 2**20)
|
|
forwards = read_forwards(args.trace)
|
|
best, ram, per, final, bound, rows = tune(forwards, num_experts, args.expert_vram_gib,
|
|
args.ram_gib, cost,
|
|
base=(Policy(), Prefetch(), Miss(cpu_threads=args.cpu_threads)))
|
|
print(" rank total_s decode_ms hit% eviction heat keep promote prefetch miss",
|
|
file=sys.stderr)
|
|
for rank, (variant, r) in enumerate(rows[: args.top], 1):
|
|
print(f"{rank:5} {r.total_ms / 1000:7.2f} {r.ms_per_decode:9.1f} {100 * r.hit_rate:4.1f}"
|
|
f" {_variant_text(variant)}", file=sys.stderr)
|
|
print("best, with per-layer cache slots:", file=sys.stderr)
|
|
print("\n".join(" " + line for line in describe(final, bound, cost)), file=sys.stderr)
|
|
policy, prefetch, miss = best
|
|
layers = {layer: {k: v for k, v in (("ram", format_ids(ram[layer])),
|
|
("cache_slots", per[layer])) if v != ""}
|
|
for layer in per}
|
|
print(render_toml(
|
|
f"generated by `prain tune` from {os.path.basename(args.trace)}",
|
|
{"pack_dir": os.path.abspath(args.pack_dir), "quant": None if args.quant == "bf16" else args.quant, "default_tier": "disk",
|
|
"cache_slots": min(per.values()), "trace_path": os.path.abspath(args.trace)},
|
|
{"policy": asdict(policy), "prefetch": asdict(prefetch), "miss": asdict(miss)},
|
|
layers,
|
|
), end="")
|
|
|
|
|
|
SERVE_OWNED = ("--enforce-eager", "--moe-backend", "--scheduling-policy", "--max-num-seqs",
|
|
"--kv-cache-memory-bytes") # set by `prain serve` from the placement file
|
|
|
|
|
|
def serve(args):
|
|
config = load_config(args.config) # fail now, not after vLLM spent minutes loading
|
|
extra = args.vllm_args[1:] if args.vllm_args[:1] == ["--"] else args.vllm_args
|
|
owned = [a for a in extra if a.split("=")[0] in SERVE_OWNED]
|
|
if owned:
|
|
raise SystemExit(f"prain serve: {' '.join(owned)} comes from the placement file "
|
|
"([serve] max_requests, kv_cache_gib); remove it after --")
|
|
os.environ["PRAIN_CONFIG"] = os.path.abspath(args.config)
|
|
if config.miss.mode != "copy": # vLLM pins torch to 1 CPU thread unless this is set
|
|
os.environ.setdefault("OMP_NUM_THREADS", str(config.miss.cpu_threads))
|
|
os.execvp("vllm", ["vllm", "serve", args.model, "--enforce-eager", "--moe-backend", "triton",
|
|
"--scheduling-policy", "priority",
|
|
"--max-num-seqs", str(config.serve.max_requests),
|
|
"--kv-cache-memory-bytes", str(int(config.serve.kv_cache_gib * GIB)),
|
|
*extra])
|
|
|
|
|
|
def main(argv=None):
|
|
parser = argparse.ArgumentParser(prog="prain")
|
|
sub = parser.add_subparsers(dest="command", required=True)
|
|
c = sub.add_parser("check", help="validate a placement file and print what it costs")
|
|
c.add_argument("config")
|
|
c.add_argument("hf_config", help="the model's config.json")
|
|
p = sub.add_parser("plan", help="print a placement file built from routing stats")
|
|
p.add_argument("stats", help="stats_path written by a prain run")
|
|
p.add_argument("hf_config", help="the model's config.json")
|
|
p.add_argument("--expert-vram-gib", type=float, required=True,
|
|
help="VRAM for experts: pinned ones plus cache slots")
|
|
p.add_argument("--ram-gib", type=float, required=True, help="pinned host memory for experts")
|
|
p.add_argument("--cache-slots", type=int, default=16)
|
|
p.add_argument("--pack-dir", required=True)
|
|
p.add_argument("--quant", choices=QUANTS, default="bf16", help="expert weight format")
|
|
for name, text in (("sim", "simulate a placement file on a routing trace"),
|
|
("tune", "search policies and per-layer slots on a trace; prints a placement file")):
|
|
s = sub.add_parser(name, help=text)
|
|
s.add_argument("trace", help="trace_path written by a prain run, or `prain synth`")
|
|
if name == "sim":
|
|
s.add_argument("config")
|
|
s.add_argument("hf_config", help="the model's config.json")
|
|
s.add_argument("--cost", help="JSON from scripts/measure.py (default: RX 6600 XT guesses)")
|
|
s.add_argument("--expert-mib", type=float, help="what-if: expert size, e.g. 2.9 for int4")
|
|
if name == "tune":
|
|
s.add_argument("--expert-vram-gib", type=float, required=True)
|
|
s.add_argument("--ram-gib", type=float, required=True)
|
|
s.add_argument("--pack-dir", required=True)
|
|
s.add_argument("--quant", choices=QUANTS, default="bf16", help="expert weight format")
|
|
s.add_argument("--cpu-threads", type=int, default=4,
|
|
help="[miss] cpu_threads to write: the --threads scripts/measure.py used")
|
|
s.add_argument("--top", type=int, default=10, help="variants to list")
|
|
y = sub.add_parser("synth", help="write a synthetic routing trace, to try sim and tune")
|
|
y.add_argument("out")
|
|
for flag, default in (("forwards", 300), ("layers", 30), ("experts", 128), ("top-k", 8),
|
|
("seed", 0)):
|
|
y.add_argument(f"--{flag}", type=int, default=default)
|
|
s = sub.add_parser("serve", help="exec `vllm serve` with the placement active")
|
|
s.add_argument("config")
|
|
s.add_argument("model")
|
|
s.add_argument("vllm_args", nargs=argparse.REMAINDER, help="extra vllm serve args after --")
|
|
for name, what in (("tui", "the terminal"), ("web", "the browser")):
|
|
ui = sub.add_parser(name, help=f"chat in {what}; experts warm up while you type")
|
|
ui.add_argument("--url", default="http://127.0.0.1:8000", help="the vLLM server")
|
|
ui.add_argument("--warmup", choices=("idle", "always", "off"), default="idle",
|
|
help="idle: not while an answer streams (default)")
|
|
ui.add_argument("--warmup-priority", type=int, default=1,
|
|
help="vLLM priority of warm-ups; 0 for a server without "
|
|
"--scheduling-policy priority")
|
|
if name == "web":
|
|
ui.add_argument("--port", type=int, default=8080)
|
|
ui.add_argument("--config", help="placement file: [serve] max_requests and the "
|
|
"[draft] table (default: their defaults)")
|
|
args = parser.parse_args(argv)
|
|
|
|
if args.command == "check":
|
|
check(args.config, args.hf_config)
|
|
elif args.command == "plan":
|
|
_, num_experts, _, size = _model_dims(args.hf_config, args.quant)
|
|
with open(args.stats) as f:
|
|
stats = json.load(f)
|
|
print(plan(stats, num_experts, size, args.expert_vram_gib, args.ram_gib,
|
|
args.cache_slots, os.path.abspath(args.pack_dir), os.path.abspath(args.stats), args.quant), end="")
|
|
elif args.command == "sim":
|
|
sim(args.trace, args.config, args.hf_config, args.cost, args.expert_mib)
|
|
elif args.command == "tune":
|
|
tune_command(args)
|
|
elif args.command == "synth":
|
|
from prain.sim import synthetic_trace, write_trace
|
|
|
|
write_trace(args.out, synthetic_trace(args.layers, args.experts, args.top_k,
|
|
args.forwards, seed=args.seed))
|
|
elif args.command == "serve":
|
|
serve(args)
|
|
elif args.command == "tui":
|
|
from prain.tui import main as tui # curses only loads when needed
|
|
|
|
tui(args.url, args.warmup, args.warmup_priority)
|
|
else:
|
|
from prain.config import Config
|
|
from prain.web import serve as web
|
|
|
|
config = load_config(args.config) if args.config else Config(pack_dir="") # defaults
|
|
web(args.url, args.port, args.warmup, args.warmup_priority, config.serve, config.draft)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|