113 lines
4.9 KiB
Python
113 lines
4.9 KiB
Python
"""Measure the prain cost model on this machine and print it as JSON for
|
|
`prain sim --cost` and `prain tune --cost`. Needs the GPU, vLLM and prain.
|
|
|
|
usage: python3 scripts/measure.py PACK_FILE [--quant int4] [--threads N] [--dense-gib 4.45] [--layers 30] > cost.json
|
|
PACK_FILE is any <pack_dir>/layerNNN.w13 file; it is read cold (dropped from the page cache).
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import time
|
|
|
|
import torch
|
|
from vllm.config import VllmConfig, set_current_vllm_config
|
|
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
|
|
from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
|
|
|
|
from prain.config import QUANTS, expert_bytes as size_of
|
|
from prain.quant import cpu_matmul, moe_quant_config, quantize, to_cpu_int4
|
|
from prain.tiered import cpu_experts
|
|
|
|
HIDDEN, INTER, SLOTS = 2816, 704, 5 # Gemma 4 26B-A4B expert, 5 slots per layer
|
|
|
|
|
|
def ms(fn, repeat=20):
|
|
for _ in range(3):
|
|
fn()
|
|
torch.cuda.synchronize()
|
|
start = time.perf_counter()
|
|
for _ in range(repeat):
|
|
fn()
|
|
torch.cuda.synchronize()
|
|
return (time.perf_counter() - start) * 1000 / repeat
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("pack_file")
|
|
parser.add_argument("--threads", type=int, default=4)
|
|
parser.add_argument("--dense-gib", type=float, default=4.45, help="non-expert weights")
|
|
parser.add_argument("--layers", type=int, default=30)
|
|
parser.add_argument("--quant", choices=QUANTS, default="bf16", help="expert weight format")
|
|
args = parser.parse_args()
|
|
torch.set_num_threads(args.threads)
|
|
expert_bytes = size_of(HIDDEN, INTER, args.quant)
|
|
host = torch.empty(expert_bytes, dtype=torch.uint8).pin_memory()
|
|
device = torch.empty_like(host, device="cuda")
|
|
h2d_gbps = expert_bytes / 1e6 / ms(lambda: device.copy_(host, non_blocking=True))
|
|
|
|
fd = os.open(args.pack_file, os.O_RDONLY)
|
|
size = os.fstat(fd).st_size
|
|
os.posix_fadvise(fd, 0, size, os.POSIX_FADV_DONTNEED) # read it cold
|
|
chunk = 2 * HIDDEN * INTER * 2 # one w13 row
|
|
start, reads = time.perf_counter(), 0
|
|
for offset in range(0, min(size, 32 * chunk), chunk):
|
|
reads += len(os.pread(fd, chunk, offset))
|
|
disk_gbps = reads / 1e6 / ((time.perf_counter() - start) * 1000)
|
|
os.close(fd)
|
|
|
|
w1 = torch.randn(SLOTS, 2 * INTER, HIDDEN, dtype=torch.bfloat16, device="cuda") * 0.02
|
|
w2 = torch.randn(SLOTS, HIDDEN, INTER, dtype=torch.bfloat16, device="cuda") * 0.02
|
|
quant_config = None
|
|
if args.quant != "bf16":
|
|
(w1, s1), (w2, s2) = quantize(w1, args.quant), quantize(w2, args.quant)
|
|
quant_config = moe_quant_config(args.quant, s1, s2)
|
|
|
|
def moe(tokens):
|
|
x = torch.randn(tokens, HIDDEN, dtype=torch.bfloat16, device="cuda")
|
|
ids = torch.arange(SLOTS, device="cuda", dtype=torch.int32).repeat(tokens, 1)
|
|
weights = torch.full((tokens, SLOTS), 1 / SLOTS, device="cuda")
|
|
return lambda: fused_experts(x, w1, w2, weights, ids, activation=MoEActivation.GELU_TANH,
|
|
global_num_experts=SLOTS, quant_config=quant_config)
|
|
|
|
with set_current_vllm_config(VllmConfig()):
|
|
round_ms = ms(moe(1))
|
|
route_ms = max(0.0, (ms(moe(256)) - round_ms) / (256 * SLOTS))
|
|
big = torch.empty(256 * 2**20, dtype=torch.uint8, device="cuda")
|
|
copy = torch.empty_like(big)
|
|
device_gbps = 2 * big.numel() / 1e6 / ms(lambda: copy.copy_(big)) # read + write
|
|
other_ms = args.dense_gib * 2**30 / 1e6 / device_gbps / args.layers
|
|
flag = torch.zeros(1, device="cuda")
|
|
sync_ms = ms(lambda: flag.add_(1).item())
|
|
# Serving computes a different expert each time, read from RAM (not from the CPU caches).
|
|
experts = [(torch.randn(2 * INTER, HIDDEN, dtype=torch.bfloat16),
|
|
torch.randn(HIDDEN, INTER, dtype=torch.bfloat16)) for _ in range(23)]
|
|
if args.quant != "bf16": # the CPU kernels on pack-layout weights, as the plugin runs them
|
|
def stored(w):
|
|
q, scale = quantize(w, args.quant)
|
|
return cpu_matmul(args.quant, to_cpu_int4(q) if args.quant == "int4" else q, scale)
|
|
|
|
experts = [(stored(a), stored(b)) for a, b in experts]
|
|
x = torch.randn(1, HIDDEN, dtype=torch.bfloat16)
|
|
one = torch.ones(1, 1)
|
|
ids = torch.zeros(1, 1, dtype=torch.int32)
|
|
|
|
def cpu_route(weights):
|
|
cpu_experts(x, one, ids, {0}, lambda e: weights)
|
|
|
|
for weights in experts[:3]: # warm up like ms(): thread pool and BLAS setup are one-time costs
|
|
cpu_route(weights)
|
|
start = time.perf_counter()
|
|
for weights in experts[3:]:
|
|
cpu_route(weights)
|
|
cpu_ms = (time.perf_counter() - start) * 1000 / len(experts[3:])
|
|
print(json.dumps({"h2d_gbps": round(h2d_gbps, 2), "disk_gbps": round(disk_gbps, 2),
|
|
"round_ms": round(round_ms, 4), "route_ms": round(route_ms, 6),
|
|
"other_ms": round(other_ms, 4), "sync_ms": round(sync_ms, 4),
|
|
"cpu_ms": round(cpu_ms, 4)}, indent=1))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|