244 lines
9.6 KiB
Python
244 lines
9.6 KiB
Python
#!/usr/bin/env python3
|
|
"""Run Frontier on the four-cell Qwen235 vLLM 0.20 latency surface."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import importlib.util
|
|
import json
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
from dataclasses import asdict, dataclass
|
|
from pathlib import Path
|
|
|
|
|
|
HERE = Path(__file__).resolve().parent
|
|
Q30_RUNNER = HERE / "run_frontier_qwen30_exact_trace_surface.py"
|
|
|
|
|
|
def load_q30():
|
|
spec = importlib.util.spec_from_file_location("q30_exact_surface", Q30_RUNNER)
|
|
if spec is None or spec.loader is None:
|
|
raise ImportError(Q30_RUNNER)
|
|
module = importlib.util.module_from_spec(spec)
|
|
sys.modules[spec.name] = module
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
Q30 = load_q30()
|
|
BASE = Q30.BASE
|
|
MODEL = "Qwen3-235B-A22B"
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Config:
|
|
tp: int
|
|
mns: int
|
|
moe_tp: int
|
|
moe_ep: int
|
|
|
|
@property
|
|
def name(self) -> str:
|
|
return f"tp{self.tp}_ep{self.moe_ep}_mns{self.mns}"
|
|
|
|
|
|
GRID = tuple(
|
|
Config(tp, mns, 4 if tp == 4 else 1, 1 if tp == 4 else 8)
|
|
for tp in (4, 8)
|
|
for mns in (64, 128)
|
|
)
|
|
|
|
|
|
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("--runtime-contract", type=Path, required=True)
|
|
parser.add_argument("--trace-tp", action="append", required=True, help="TP=PATH")
|
|
parser.add_argument("--config", action="append")
|
|
parser.add_argument("--prefix-caching", action=argparse.BooleanOptionalAction, default=True)
|
|
parser.add_argument("--allreduce-csv", type=Path, required=True)
|
|
parser.add_argument("--timeout-seconds", type=float, default=7200)
|
|
parser.add_argument("--resume", action="store_true")
|
|
return parser.parse_args()
|
|
|
|
|
|
def profile_paths(root: Path) -> dict[str, Path]:
|
|
paths = {
|
|
"linear": root / "linear_op.csv",
|
|
"attention": root / "attention.csv",
|
|
"moe": root / "moe.csv",
|
|
"linear_kernel": root / "linear_op_kernel_only.csv",
|
|
"attention_kernel": root / "attention_kernel_only.csv",
|
|
"moe_kernel": root / "moe_kernel_only.csv",
|
|
"manifest": root / "manifest.json",
|
|
}
|
|
missing = [str(path) for path in paths.values() if not path.is_file()]
|
|
if missing:
|
|
raise FileNotFoundError(missing)
|
|
manifest = json.loads(paths["manifest"].read_text())
|
|
for name, path in paths.items():
|
|
if name == "manifest":
|
|
continue
|
|
expected = manifest.get("outputs", {}).get(path.name)
|
|
if expected != BASE.sha256(path):
|
|
raise ValueError(f"profile hash mismatch: {path}")
|
|
return paths
|
|
|
|
|
|
def knobs(config: Config, paths: dict[str, Path], contract: dict, cache: Path, prefix: bool):
|
|
resolved = contract[config.name]
|
|
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.moe_tp,
|
|
"moe_expert_parallel_size": config.moe_ep,
|
|
"num_pipeline_stages": 1,
|
|
"replica_scheduler": "vllm_v1",
|
|
"decode_cuda_graph_mode": "piecewise",
|
|
"batch_size_cap": config.mns,
|
|
"max_tokens_in_batch": 8192,
|
|
"long_prefill_token_threshold": 0,
|
|
"enable_chunked_prefill": True,
|
|
"block_size": 16,
|
|
"num_blocks_mode": "explicit",
|
|
"num_blocks": int(resolved["num_gpu_blocks"]),
|
|
"gpu_memory_utilization": 0.80,
|
|
"non_kv_cache_overhead_bytes": 0,
|
|
"trace_max_tokens": 40960,
|
|
"cache_dir": str(cache / config.name),
|
|
"enable_prefix_caching": prefix,
|
|
"enable_dummy_mode": False,
|
|
"linear_op_input_file": str(paths["linear"]),
|
|
"atten_input_file": str(paths["attention"]),
|
|
"moe_input_file": str(paths["moe"]),
|
|
"linear_op_kernel_only_input_file": str(paths["linear_kernel"]),
|
|
"atten_kernel_only_input_file": str(paths["attention_kernel"]),
|
|
"moe_kernel_only_input_file": str(paths["moe_kernel"]),
|
|
"prediction_max_prefill_chunk_size": 8192,
|
|
"prediction_max_tokens_per_request": 40960,
|
|
"prediction_max_batch_size": max(int(v) for v in resolved["capture_sizes"]),
|
|
"no_cache": True,
|
|
}
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
for name in ("frontier_source", "replayserve_root", "profile_root", "python_deps", "output_root", "runtime_contract", "allreduce_csv"):
|
|
setattr(args, name, getattr(args, name).resolve())
|
|
paths = profile_paths(args.profile_root)
|
|
contract = json.loads(args.runtime_contract.read_text())["configs"]
|
|
trace_by_tp = {}
|
|
for specification in args.trace_tp:
|
|
raw_tp, separator, path = specification.partition("=")
|
|
if not separator:
|
|
raise ValueError(f"trace-tp must be TP=PATH: {specification}")
|
|
tp = int(raw_tp)
|
|
trace_by_tp[tp] = Q30.parse_trace(
|
|
f"eval={path}", rate_contract="trace-window", prefix_caching=args.prefix_caching
|
|
)
|
|
if set(trace_by_tp) != {4, 8}:
|
|
raise ValueError("trace-tp must provide exactly TP4 and TP8")
|
|
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}}")
|
|
|
|
builder = BASE.load_module(
|
|
"qwen235_frontier_builder", args.replayserve_root / "tools/run_frontier_sweep.py"
|
|
)
|
|
environment = os.environ.copy()
|
|
environment.update(
|
|
{
|
|
"PYTHONPATH": ":".join([str(args.python_deps), str(args.frontier_source)]),
|
|
"CUDA_VISIBLE_DEVICES": "",
|
|
"NVIDIA_VISIBLE_DEVICES": "void",
|
|
"WANDB_DISABLED": "true",
|
|
"VIDUR_DISABLE_WANDB": "1",
|
|
"FRONTIER_LOG_LEVEL": "WARNING",
|
|
"PYTHONDONTWRITEBYTECODE": "1",
|
|
}
|
|
)
|
|
results = []
|
|
for config in selected:
|
|
config_knobs = knobs(config, paths, contract, args.output_root / "cache", args.prefix_caching)
|
|
for trace in (trace_by_tp[config.tp],):
|
|
run_dir = args.output_root / "runs" / config.name / trace["label"]
|
|
result_path = run_dir / "result.json"
|
|
if args.resume and result_path.is_file():
|
|
previous = json.loads(result_path.read_text())
|
|
if previous.get("status") == "completed":
|
|
results.append(previous)
|
|
continue
|
|
run_dir.mkdir(parents=True, exist_ok=True)
|
|
command = builder.build_frontier_command(
|
|
python_bin="/usr/bin/python3",
|
|
trace_file=trace["path"],
|
|
metrics_root=run_dir / "metrics",
|
|
run_id=f"qwen235_v020_{config.name}_{trace['label']}",
|
|
knobs=config_knobs,
|
|
)
|
|
command.extend(["--cudagraph_capture_sizes", *(str(v) for v in contract[config.name]["capture_sizes"])])
|
|
command = BASE.configure_cc_command(
|
|
command,
|
|
backend="vidur",
|
|
allreduce_csv=args.allreduce_csv,
|
|
cache=args.output_root / "cc-cache",
|
|
)
|
|
BASE.write_json(run_dir / "command.json", command)
|
|
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,
|
|
)
|
|
result = {
|
|
"config": asdict(config) | {"name": config.name},
|
|
"trace": Q30.trace_manifest_entry(trace),
|
|
"elapsed_seconds": time.time() - started,
|
|
"returncode": completed.returncode,
|
|
}
|
|
if completed.returncode == 0:
|
|
metrics = Q30.find_request_metrics(run_dir)
|
|
result.update(status="completed", metrics=Q30.score(metrics, trace["shapes"]), request_metrics_sha256=BASE.sha256(metrics))
|
|
else:
|
|
stderr_text = (run_dir / "stderr.log").read_text(errors="replace")
|
|
result.update(status="failed", failure_class=Q30.classify_frontier_failure(stderr_text))
|
|
BASE.write_json(result_path, result)
|
|
results.append(result)
|
|
print(json.dumps({"config": config.name, "trace": trace["label"], "status": result["status"]}, sort_keys=True), flush=True)
|
|
|
|
manifest = {
|
|
"schema": "qwen235-v020-frontier-latency-surface-v1",
|
|
"frontier_commit": subprocess.check_output(["git", "-C", str(args.frontier_source), "rev-parse", "HEAD"], text=True).strip(),
|
|
"profiles": {name: BASE.sha256(path) for name, path in paths.items()},
|
|
"runtime_contract_sha256": BASE.sha256(args.runtime_contract),
|
|
"prefix_caching": args.prefix_caching,
|
|
"results": results,
|
|
}
|
|
BASE.write_json(args.output_root / "frontier_surface.json", manifest)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|