223 lines
7.5 KiB
Python
223 lines
7.5 KiB
Python
#!/usr/bin/env python3
|
|
"""Replay the exact BC-8 simulator cells with the whole-layer decode curve."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
from pathlib import Path
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parent
|
|
REPO = ROOT.parents[1]
|
|
EXPECTED_FRONTIER_COMMIT = "deadc4a321f0baaa534c6ebd17f974123733cdc2"
|
|
CONFIGS = ("tp1_mns16", "tp2_mns16", "tp4_mns16", "tp4_mns32")
|
|
SOURCE_MANIFEST = (
|
|
REPO / "runs/frontier-collective-joint-v0/counterfactual/joint-r2/manifest.json"
|
|
)
|
|
JOINT_INPUTS = REPO / "runs/frontier-knee-sweep-v0/inputs"
|
|
GRID = ROOT / "results/grid.json"
|
|
WRAPPER = ROOT / "run_frontier_with_whole_decode_curve.py"
|
|
CACHE_ROOT = REPO / "runs/frontier-knee-sweep-v0/cache"
|
|
LOCAL_DEPENDENCY_ROOTS = (
|
|
REPO / "runs/frontier-collective-joint-v0/python-deps",
|
|
Path("/home/gahow/.cache/uv/archive-v0/-_kzErLcPO5nASZFX8b9k"),
|
|
Path("/home/gahow/.cache/uv/archive-v0/FbaBs_QJ9QKEbQ9V_4aIR"),
|
|
Path("/home/gahow/.cache/uv/archive-v0/fuHsGXD0Lv_UjFC8yI4-7"),
|
|
Path("/home/gahow/.cache/uv/archive-v0/jFGdqQLpB1eopfm9VxT3j"),
|
|
Path("/home/gahow/.cache/uv/archive-v0/YWW6ExSJuPVvv4-qYQTin"),
|
|
Path("/home/gahow/.cache/uv/archive-v0/3_qxZ5Ll-EpVAGZfbksfe"),
|
|
)
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--frontier-checkout", type=Path, required=True)
|
|
parser.add_argument("--output-root", type=Path, default=ROOT / "replay/bc8")
|
|
parser.add_argument("--jobs", type=int, default=2)
|
|
return parser.parse_args()
|
|
|
|
|
|
def sha256(path: Path) -> str:
|
|
return hashlib.sha256(path.read_bytes()).hexdigest()
|
|
|
|
|
|
def write_json(path: Path, payload) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
|
|
|
|
|
|
def replace_flag(argv: list[str], flag: str, value: str) -> None:
|
|
index = argv.index(flag)
|
|
argv[index + 1] = value
|
|
|
|
|
|
def validate(frontier: Path) -> None:
|
|
commit = subprocess.check_output(
|
|
["git", "-C", str(frontier), "rev-parse", "HEAD"], text=True
|
|
).strip()
|
|
status = subprocess.check_output(
|
|
["git", "-C", str(frontier), "status", "--porcelain"], text=True
|
|
).strip()
|
|
if commit != EXPECTED_FRONTIER_COMMIT or status:
|
|
raise ValueError(
|
|
f"Frontier must be clean at {EXPECTED_FRONTIER_COMMIT}, "
|
|
f"got commit={commit}, dirty={bool(status)}"
|
|
)
|
|
required = (
|
|
SOURCE_MANIFEST,
|
|
GRID,
|
|
WRAPPER,
|
|
CACHE_ROOT,
|
|
*LOCAL_DEPENDENCY_ROOTS,
|
|
)
|
|
missing = [str(path) for path in required if not path.exists()]
|
|
if missing:
|
|
raise ValueError(f"missing replay inputs: {missing}")
|
|
|
|
|
|
def trace_for(config: str) -> Path:
|
|
tp = int(config[2])
|
|
return (
|
|
JOINT_INPUTS
|
|
/ f"traces-per-gpu-low/tp{tp}/w0-short-fixed-uniform-none/"
|
|
"rho0p02/public/frontier.csv"
|
|
)
|
|
|
|
|
|
def run_one(
|
|
config: str,
|
|
*,
|
|
frontier: Path,
|
|
output_root: Path,
|
|
templates: dict,
|
|
) -> dict:
|
|
point = output_root / "raw" / config
|
|
metrics_root = point / "metrics"
|
|
expected = list(metrics_root.glob("**/system_metrics.json"))
|
|
if len(expected) == 1 and (point / "usage.json").is_file():
|
|
return {"config": config, "status": "skipped_complete", "elapsed_s": 0.0}
|
|
|
|
argv = list(templates[config]["argv"])
|
|
argv[0] = sys.executable
|
|
argv[1] = str(WRAPPER.resolve())
|
|
replace_flag(argv, "--trace_request_generator_config_trace_file", str(trace_for(config)))
|
|
replace_flag(argv, "--metrics_config_output_dir", str(metrics_root))
|
|
replace_flag(argv, "--metrics_config_run_id", f"decode_grid_bc8_{config}")
|
|
replace_flag(argv, "--metrics_config_cache_dir", str(CACHE_ROOT / "model"))
|
|
replace_flag(argv, "--vidur_cc_backend_config_cache_dir", str(CACHE_ROOT / "cc"))
|
|
env = os.environ.copy()
|
|
env.update(
|
|
{
|
|
"CUDA_VISIBLE_DEVICES": "",
|
|
"PYTHONDONTWRITEBYTECODE": "1",
|
|
"PYTHONPATH": os.pathsep.join(
|
|
[str(frontier), *(str(path) for path in LOCAL_DEPENDENCY_ROOTS)]
|
|
),
|
|
"FRONTIER_COLLECTIVE_CURVE": str(
|
|
(JOINT_INPUTS / "collective-curve-b4-extrapolated.json").resolve()
|
|
),
|
|
"FRONTIER_COLLECTIVE_CURVE_VARIANT": "drop_mean",
|
|
"FRONTIER_FUSED_MOE_CURVE": str(
|
|
(JOINT_INPUTS / "fused-moe-curve-b4-extrapolated.json").resolve()
|
|
),
|
|
"FRONTIER_WHOLE_DECODE_GRID": str(GRID.resolve()),
|
|
"FRONTIER_CURVE_USAGE": str((point / "usage.json").resolve()),
|
|
}
|
|
)
|
|
point.mkdir(parents=True, exist_ok=True)
|
|
write_json(point / "command.json", argv)
|
|
started = time.monotonic()
|
|
with (point / "run.log").open("w") as output:
|
|
completed = subprocess.run(
|
|
argv,
|
|
cwd=frontier,
|
|
env=env,
|
|
stdout=output,
|
|
stderr=subprocess.STDOUT,
|
|
check=False,
|
|
)
|
|
elapsed = time.monotonic() - started
|
|
metrics = list(metrics_root.glob("**/system_metrics.json"))
|
|
status = (
|
|
"completed"
|
|
if completed.returncode == 0
|
|
and len(metrics) == 1
|
|
and (point / "usage.json").is_file()
|
|
else "failed"
|
|
)
|
|
record = {
|
|
"config": config,
|
|
"status": status,
|
|
"returncode": completed.returncode,
|
|
"elapsed_s": elapsed,
|
|
}
|
|
write_json(point / "run-status.json", record)
|
|
return record
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
if args.jobs < 1:
|
|
raise ValueError("--jobs must be positive")
|
|
frontier = args.frontier_checkout.resolve()
|
|
output_root = args.output_root.resolve()
|
|
validate(frontier)
|
|
source = json.loads(SOURCE_MANIFEST.read_text())
|
|
manifest = {
|
|
"schema": "frontier-decode-grid-bc8-replay.v1",
|
|
"frontier_checkout": str(frontier),
|
|
"frontier_commit": EXPECTED_FRONTIER_COMMIT,
|
|
"configs": list(CONFIGS),
|
|
"rho_per_gpu": 0.02,
|
|
"wrapper": str(WRAPPER.resolve()),
|
|
"wrapper_sha256": sha256(WRAPPER),
|
|
"whole_decode_grid": str(GRID.resolve()),
|
|
"whole_decode_grid_sha256": sha256(GRID),
|
|
"collective_curve_sha256": sha256(
|
|
JOINT_INPUTS / "collective-curve-b4-extrapolated.json"
|
|
),
|
|
"moe_curve_sha256": sha256(
|
|
JOINT_INPUTS / "fused-moe-curve-b4-extrapolated.json"
|
|
),
|
|
"traces": {
|
|
config: {
|
|
"path": str(trace_for(config).resolve()),
|
|
"sha256": sha256(trace_for(config)),
|
|
}
|
|
for config in CONFIGS
|
|
},
|
|
}
|
|
write_json(output_root / "manifest.json", manifest)
|
|
results = []
|
|
with ThreadPoolExecutor(max_workers=args.jobs) as pool:
|
|
futures = {
|
|
pool.submit(
|
|
run_one,
|
|
config,
|
|
frontier=frontier,
|
|
output_root=output_root,
|
|
templates=source["cells"],
|
|
): config
|
|
for config in CONFIGS
|
|
}
|
|
for future in as_completed(futures):
|
|
result = future.result()
|
|
results.append(result)
|
|
print(json.dumps(result, sort_keys=True), flush=True)
|
|
results.sort(key=lambda row: CONFIGS.index(row["config"]))
|
|
write_json(output_root / "run-summary.json", results)
|
|
failures = [row for row in results if row["status"] == "failed"]
|
|
if failures:
|
|
raise SystemExit(f"failed replay cells: {failures}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|