diff --git a/scripts/measure.py b/scripts/measure.py new file mode 100644 index 0000000..2fc4e6f --- /dev/null +++ b/scripts/measure.py @@ -0,0 +1,92 @@ +"""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 [--threads N] [--dense-gib 4.45] [--layers 30] > cost.json +PACK_FILE is any /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.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) + args = parser.parse_args() + torch.set_num_threads(args.threads) + expert_bytes = 3 * HIDDEN * INTER * 2 + host = torch.empty(expert_bytes // 2, dtype=torch.bfloat16).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 + + 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) + + 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()) + w13_cpu = torch.randn(2 * INTER, HIDDEN, dtype=torch.bfloat16).pin_memory() + w2_cpu = torch.randn(HIDDEN, INTER, dtype=torch.bfloat16).pin_memory() + x = torch.randn(1, HIDDEN, dtype=torch.bfloat16) + one = torch.ones(1, 1) + ids = torch.zeros(1, 1, dtype=torch.int32) + start = time.perf_counter() + for _ in range(20): + cpu_experts(x, one, ids, {0}, lambda e: (w13_cpu, w2_cpu)) + cpu_ms = (time.perf_counter() - start) * 1000 / 20 + 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()