Add Qwen30 prefill fidelity experiment
This commit is contained in:
@@ -0,0 +1,400 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Freeze the Qwen30 fixed-shape prefill-only Frontier surface."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import csv
|
||||
import hashlib
|
||||
import importlib.util
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from dataclasses import asdict, dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
|
||||
MODEL = "qwen3-a3b-30b-moe"
|
||||
RATES = (4.0, 8.0, 16.0, 32.0, 64.0)
|
||||
TTFT_SLO_MS = 1256.0
|
||||
TARGET_PASS_RATE = 0.95
|
||||
NUM_BLOCKS = {1: 20080, 2: 76537, 4: 191727}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Config:
|
||||
tp: int
|
||||
mns: int
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return f"tp{self.tp}_mns{self.mns}"
|
||||
|
||||
|
||||
GRID = tuple(Config(tp, mns) for tp in (1, 2, 4) for mns in (8, 16, 32, 64))
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--frontier-source", type=Path, required=True)
|
||||
parser.add_argument("--replayserve-root", type=Path, required=True)
|
||||
parser.add_argument("--profile-root", type=Path, required=True)
|
||||
parser.add_argument("--python-deps", type=Path, required=True)
|
||||
parser.add_argument("--output-root", type=Path, required=True)
|
||||
parser.add_argument("--requests", type=int, default=64)
|
||||
parser.add_argument("--rate", type=float, action="append")
|
||||
parser.add_argument("--config", action="append")
|
||||
parser.add_argument("--timeout-seconds", type=float, default=900.0)
|
||||
parser.add_argument("--resume", action="store_true")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def sha256(path: Path) -> str:
|
||||
digest = hashlib.sha256()
|
||||
with path.open("rb") as source:
|
||||
for chunk in iter(lambda: source.read(1 << 20), b""):
|
||||
digest.update(chunk)
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
def write_json(path: Path, payload: Any) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
temporary = path.with_suffix(path.suffix + ".tmp")
|
||||
temporary.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
|
||||
os.replace(temporary, path)
|
||||
|
||||
|
||||
def load_module(name: str, path: Path):
|
||||
spec = importlib.util.spec_from_file_location(name, path)
|
||||
if spec is None or spec.loader is None:
|
||||
raise ImportError(path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules[name] = module
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
def write_trace(path: Path, *, requests: int, rate: float) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
fields = [
|
||||
"arrived_at",
|
||||
"num_prefill_tokens",
|
||||
"num_decode_tokens",
|
||||
"session_id",
|
||||
"block_hash_ids",
|
||||
]
|
||||
with path.open("w", newline="") as output:
|
||||
writer = csv.DictWriter(output, fieldnames=fields)
|
||||
writer.writeheader()
|
||||
for request_id in range(requests):
|
||||
writer.writerow(
|
||||
{
|
||||
"arrived_at": f"{request_id / rate:.12f}",
|
||||
"num_prefill_tokens": 2048,
|
||||
"num_decode_tokens": 1,
|
||||
"session_id": request_id,
|
||||
"block_hash_ids": "|".join(
|
||||
str(request_id * 128 + block + 1) for block in range(128)
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def profile_paths(root: Path) -> dict[str, Path]:
|
||||
paths = {
|
||||
"linear": root / "linear_op.csv",
|
||||
"attention": root / "attention.csv",
|
||||
"moe": root / "moe.csv",
|
||||
"manifest": root / "manifest.json",
|
||||
}
|
||||
missing = [str(path) for path in paths.values() if not path.is_file()]
|
||||
if missing:
|
||||
raise FileNotFoundError(missing)
|
||||
return paths
|
||||
|
||||
|
||||
def validate_profile(paths: dict[str, Path]) -> dict[str, Any]:
|
||||
manifest = json.loads(paths["manifest"].read_text())
|
||||
expected = manifest["outputs"]
|
||||
for filename in ("linear_op.csv", "attention.csv", "moe.csv"):
|
||||
actual = sha256(paths[{"linear_op.csv": "linear", "attention.csv": "attention", "moe.csv": "moe"}[filename]])
|
||||
if actual != expected[filename]:
|
||||
raise ValueError(f"profile hash mismatch for {filename}")
|
||||
with paths["attention"].open(newline="") as source:
|
||||
rows = list(csv.DictReader(source))
|
||||
coverage = {}
|
||||
for tp in (1, 2, 4):
|
||||
exact = [
|
||||
row
|
||||
for row in rows
|
||||
if int(row["num_tensor_parallel_workers"]) == tp
|
||||
and row["is_prefill"].lower() == "true"
|
||||
and row.get("is_true_mixed_batch", "").lower() != "true"
|
||||
and int(float(row["total_tokens"])) == 2048
|
||||
]
|
||||
if len(exact) != 1:
|
||||
raise ValueError(f"expected one exact TP{tp} 2048-token prefill row, got {len(exact)}")
|
||||
coverage[str(tp)] = {"exact_prefill_2048_rows": len(exact), "profile_batch_size": int(exact[0]["batch_size"])}
|
||||
return {"manifest": manifest, "attention": coverage}
|
||||
|
||||
|
||||
def knobs(config: Config, paths: dict[str, Path], cache: Path) -> dict[str, Any]:
|
||||
return {
|
||||
"simulation_mode": "online",
|
||||
"sys_arch": "co-location",
|
||||
"num_replicas": 1,
|
||||
"cluster_scheduler": "sticky_round_robin",
|
||||
"model_name": MODEL,
|
||||
"device": "h20",
|
||||
"network_device": "h20_dgx",
|
||||
"attn_tensor_parallel_size": config.tp,
|
||||
"attn_data_parallel_size": 1,
|
||||
"moe_tensor_parallel_size": config.tp,
|
||||
"moe_expert_parallel_size": 1,
|
||||
"num_pipeline_stages": 1,
|
||||
"replica_scheduler": "vllm_v1",
|
||||
"decode_cuda_graph_mode": "none",
|
||||
"batch_size_cap": config.mns,
|
||||
"max_tokens_in_batch": 8192,
|
||||
"long_prefill_token_threshold": 0,
|
||||
"block_size": 16,
|
||||
"num_blocks_mode": "explicit",
|
||||
"num_blocks": NUM_BLOCKS[config.tp],
|
||||
"gpu_memory_utilization": 0.92,
|
||||
"non_kv_cache_overhead_bytes": 0,
|
||||
"trace_max_tokens": 40960,
|
||||
"cache_dir": str(cache),
|
||||
"enable_dummy_mode": False,
|
||||
"linear_op_input_file": str(paths["linear"]),
|
||||
"atten_input_file": str(paths["attention"]),
|
||||
"moe_input_file": str(paths["moe"]),
|
||||
"prediction_max_prefill_chunk_size": 18000,
|
||||
"prediction_max_batch_size": 128,
|
||||
"prediction_max_tokens_per_request": 32768,
|
||||
"no_cache": False,
|
||||
"skip_cpu_overhead_modeling": True,
|
||||
"enable_prefix_caching": False,
|
||||
"enable_chunked_prefill": True,
|
||||
}
|
||||
|
||||
|
||||
def find_metrics(run_dir: Path) -> tuple[Path, Path]:
|
||||
systems = list((run_dir / "metrics").rglob("system_metrics.json"))
|
||||
requests = list((run_dir / "metrics").rglob("request_metrics.csv"))
|
||||
if len(systems) != 1 or len(requests) != 1:
|
||||
raise RuntimeError(f"metric pair mismatch: {len(systems)}/{len(requests)}")
|
||||
return systems[0], requests[0]
|
||||
|
||||
|
||||
def score(system_path: Path, request_path: Path, expected_requests: int) -> dict[str, Any]:
|
||||
system = json.loads(system_path.read_text())
|
||||
metadata = system["simulation_metadata"]
|
||||
if int(metadata["completed_requests"]) != expected_requests:
|
||||
raise ValueError("Frontier completion count mismatch")
|
||||
with request_path.open(newline="") as source:
|
||||
rows = list(csv.DictReader(source))
|
||||
if len(rows) != expected_requests:
|
||||
raise ValueError("request metric count mismatch")
|
||||
values = []
|
||||
passed = 0
|
||||
for row in rows:
|
||||
if int(float(row["request_num_prefill_tokens"])) != 2048:
|
||||
raise ValueError("prefill shape drift")
|
||||
if int(float(row["request_num_decode_tokens"])) != 1:
|
||||
raise ValueError("decode shape drift")
|
||||
ttft = float(row["ttft"])
|
||||
if not math.isfinite(ttft) or ttft < 0:
|
||||
raise ValueError("invalid TTFT")
|
||||
values.append(ttft)
|
||||
passed += int(ttft <= TTFT_SLO_MS)
|
||||
ordered = sorted(values)
|
||||
pass_rate = passed / expected_requests
|
||||
return {
|
||||
"ttft_p50_ms": ordered[math.ceil(0.50 * len(ordered)) - 1],
|
||||
"ttft_p95_ms": ordered[math.ceil(0.95 * len(ordered)) - 1],
|
||||
"ttft_max_ms": max(ordered),
|
||||
"passed": passed,
|
||||
"pass_rate": pass_rate,
|
||||
"feasible": pass_rate >= TARGET_PASS_RATE,
|
||||
"throughput_requests_per_second": float(system["throughput_metrics"]["requests_per_second"]),
|
||||
}
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
args.frontier_source = args.frontier_source.resolve()
|
||||
args.replayserve_root = args.replayserve_root.resolve()
|
||||
args.profile_root = args.profile_root.resolve()
|
||||
args.python_deps = args.python_deps.resolve()
|
||||
args.output_root = args.output_root.resolve()
|
||||
rates = tuple(args.rate or RATES)
|
||||
selected = list(GRID)
|
||||
if args.config:
|
||||
wanted = set(args.config)
|
||||
selected = [config for config in GRID if config.name in wanted]
|
||||
if {config.name for config in selected} != wanted:
|
||||
raise ValueError(f"unknown configs: {wanted - {config.name for config in selected}}")
|
||||
paths = profile_paths(args.profile_root)
|
||||
coverage = validate_profile(paths)
|
||||
builder = load_module(
|
||||
"qwen30_prefill_frontier_builder",
|
||||
args.replayserve_root / "tools/run_frontier_sweep.py",
|
||||
)
|
||||
frontier_head = subprocess.run(
|
||||
["git", "-C", str(args.frontier_source), "rev-parse", "HEAD"],
|
||||
check=True,
|
||||
text=True,
|
||||
stdout=subprocess.PIPE,
|
||||
).stdout.strip()
|
||||
traces = {}
|
||||
for rate in rates:
|
||||
trace = args.output_root / "traces" / f"r{rate:g}.csv"
|
||||
write_trace(trace, requests=args.requests, rate=rate)
|
||||
traces[rate] = trace
|
||||
|
||||
config_results = []
|
||||
for config in selected:
|
||||
loads = []
|
||||
config_knobs = knobs(config, paths, args.output_root / "cache")
|
||||
for rate in rates:
|
||||
run_dir = args.output_root / "runs" / config.name / f"r{rate:g}"
|
||||
result_path = run_dir / "result.json"
|
||||
if args.resume and result_path.is_file():
|
||||
loads.append(json.loads(result_path.read_text()))
|
||||
continue
|
||||
run_dir.mkdir(parents=True, exist_ok=True)
|
||||
command = builder.build_frontier_command(
|
||||
python_bin="/usr/bin/python3",
|
||||
trace_file=traces[rate],
|
||||
metrics_root=run_dir / "metrics",
|
||||
run_id=f"qwen30_prefill_{config.name}_r{rate:g}",
|
||||
knobs=config_knobs,
|
||||
)
|
||||
write_json(run_dir / "command.json", command)
|
||||
environment = os.environ.copy()
|
||||
pythonpath = [str(args.python_deps), str(args.frontier_source)]
|
||||
if environment.get("PYTHONPATH"):
|
||||
pythonpath.append(environment["PYTHONPATH"])
|
||||
environment.update(
|
||||
{
|
||||
"PYTHONPATH": ":".join(pythonpath),
|
||||
"CUDA_VISIBLE_DEVICES": "",
|
||||
"NVIDIA_VISIBLE_DEVICES": "void",
|
||||
"WANDB_DISABLED": "true",
|
||||
"VIDUR_DISABLE_WANDB": "1",
|
||||
"FRONTIER_LOG_LEVEL": "WARNING",
|
||||
"PYTHONDONTWRITEBYTECODE": "1",
|
||||
}
|
||||
)
|
||||
started = time.time()
|
||||
with (run_dir / "stdout.log").open("w") as stdout, (
|
||||
run_dir / "stderr.log"
|
||||
).open("w") as stderr:
|
||||
completed = subprocess.run(
|
||||
command,
|
||||
cwd=args.frontier_source,
|
||||
env=environment,
|
||||
stdout=stdout,
|
||||
stderr=stderr,
|
||||
timeout=args.timeout_seconds,
|
||||
check=False,
|
||||
)
|
||||
if completed.returncode != 0:
|
||||
raise RuntimeError(
|
||||
f"Frontier failed for {config.name} rate={rate}: {completed.returncode}"
|
||||
)
|
||||
system_path, request_path = find_metrics(run_dir)
|
||||
result = {
|
||||
"status": "completed",
|
||||
"config": asdict(config) | {"name": config.name},
|
||||
"offered_request_rate": rate,
|
||||
"offered_request_rate_per_gpu": rate / config.tp,
|
||||
"request_count": args.requests,
|
||||
"elapsed_seconds": time.time() - started,
|
||||
"trace_sha256": sha256(traces[rate]),
|
||||
"request_metrics_sha256": sha256(request_path),
|
||||
"score": score(system_path, request_path, args.requests),
|
||||
}
|
||||
write_json(result_path, result)
|
||||
loads.append(result)
|
||||
print(
|
||||
json.dumps(
|
||||
{
|
||||
"config": config.name,
|
||||
"rate": rate,
|
||||
"pass_rate": result["score"]["pass_rate"],
|
||||
"feasible": result["score"]["feasible"],
|
||||
},
|
||||
sort_keys=True,
|
||||
),
|
||||
flush=True,
|
||||
)
|
||||
config_results.append({"config": asdict(config) | {"name": config.name}, "loads": loads})
|
||||
|
||||
capacities = []
|
||||
for item in config_results:
|
||||
feasible = [
|
||||
load["offered_request_rate"]
|
||||
for load in item["loads"]
|
||||
if load["score"]["feasible"]
|
||||
]
|
||||
capacity = max(feasible) if feasible else None
|
||||
capacities.append(
|
||||
{
|
||||
"config": item["config"],
|
||||
"maximum_tested_feasible_request_rate": capacity,
|
||||
"maximum_tested_feasible_request_rate_per_gpu": (
|
||||
capacity / item["config"]["tp"] if capacity is not None else None
|
||||
),
|
||||
"lower_censored": capacity is None,
|
||||
"upper_censored": capacity == max(rates) if capacity is not None else False,
|
||||
}
|
||||
)
|
||||
capacities.sort(
|
||||
key=lambda row: (
|
||||
-(row["maximum_tested_feasible_request_rate_per_gpu"] or -1),
|
||||
row["config"]["name"],
|
||||
)
|
||||
)
|
||||
full = selected == list(GRID) and rates == RATES and args.requests == 64
|
||||
manifest = {
|
||||
"schema": "frontier-qwen30-prefill-surface-v1",
|
||||
"status": "frozen_before_real" if full else "partial_not_decision_bearing",
|
||||
"contract": {
|
||||
"rates": rates,
|
||||
"requests_per_anchor": args.requests,
|
||||
"input_tokens": 2048,
|
||||
"output_tokens": 1,
|
||||
"ttft_slo_ms": TTFT_SLO_MS,
|
||||
"target_pass_rate": TARGET_PASS_RATE,
|
||||
"prefix_caching": False,
|
||||
"arrival": "open_loop_uniform",
|
||||
},
|
||||
"frontier": {
|
||||
"source": str(args.frontier_source),
|
||||
"git_head": frontier_head,
|
||||
"git_status_short": subprocess.run(
|
||||
["git", "-C", str(args.frontier_source), "status", "--short"],
|
||||
check=True,
|
||||
text=True,
|
||||
stdout=subprocess.PIPE,
|
||||
).stdout,
|
||||
},
|
||||
"profiles": {
|
||||
"root": str(args.profile_root),
|
||||
"coverage": coverage,
|
||||
"sha256": {name: sha256(path) for name, path in paths.items()},
|
||||
},
|
||||
"config_results": config_results,
|
||||
"capacity": capacities,
|
||||
}
|
||||
write_json(args.output_root / "frontier_surface_frozen.json", manifest)
|
||||
print(args.output_root / "frontier_surface_frozen.json")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user