127 lines
4.4 KiB
Python
127 lines
4.4 KiB
Python
#!/usr/bin/env python3
|
|
"""Replay one real-trace cell with the structured-attention experiment commit."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import importlib.util
|
|
import json
|
|
import subprocess
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
ROOT = Path(__file__).resolve().parent
|
|
REPO = ROOT.parents[1]
|
|
S3_REAL = REPO / "runs/frontier-s3-real-v0"
|
|
BASE_REFERENCE = (
|
|
REPO
|
|
/ "runs/frontier-collective-joint-v0/counterfactual/joint-r2/manifest.json"
|
|
)
|
|
BASE_COMMIT = "deadc4a321f0baaa534c6ebd17f974123733cdc2"
|
|
EXPERIMENT_COMMIT = "1f8900a4ac64e45754b03d0aa7c1dddab65785cf"
|
|
PATCH = ROOT / "0001-Experiment-with-structured-attention-prefill-predict.patch"
|
|
|
|
|
|
def load_s3_module():
|
|
spec = importlib.util.spec_from_file_location(
|
|
"s3_prefix_replay", S3_REAL / "run_frontier_prefix_replay.py"
|
|
)
|
|
module = importlib.util.module_from_spec(spec)
|
|
sys.path.insert(0, str(S3_REAL))
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
def git(checkout: Path, *args: str) -> str:
|
|
return subprocess.check_output(
|
|
["git", "-C", str(checkout), *args], text=True
|
|
).strip()
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--trace", type=Path, required=True)
|
|
parser.add_argument("--output-root", type=Path, required=True)
|
|
parser.add_argument(
|
|
"--config",
|
|
choices=("tp4_mns16", "tp2_mns16", "tp1_mns16"),
|
|
required=True,
|
|
)
|
|
parser.add_argument("--label", required=True)
|
|
parser.add_argument("--max-tokens", type=int, required=True)
|
|
parser.add_argument("--duration-s", type=float)
|
|
parser.add_argument("--cache-root", type=Path, required=True)
|
|
parser.add_argument(
|
|
"--frontier-checkout",
|
|
type=Path,
|
|
default=Path("/tmp/frontier-attn-structured-v0"),
|
|
)
|
|
parser.add_argument(
|
|
"--attention-profile",
|
|
type=Path,
|
|
default=REPO
|
|
/ "runs/frontier-prefill-kvgrowth-fix-v0/profiles/"
|
|
"profile-v5-kvgrowth/attention.csv",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
frontier = args.frontier_checkout.resolve()
|
|
profile = args.attention_profile.resolve()
|
|
if git(frontier, "rev-parse", "HEAD") != EXPERIMENT_COMMIT:
|
|
raise SystemExit(f"unexpected experiment checkout HEAD: {frontier}")
|
|
if git(frontier, "rev-parse", "HEAD^") != BASE_COMMIT:
|
|
raise SystemExit("experiment commit is not directly based on frozen Frontier")
|
|
if git(frontier, "status", "--porcelain"):
|
|
raise SystemExit("experiment Frontier checkout must be clean")
|
|
if not profile.is_file():
|
|
raise SystemExit(f"attention profile missing: {profile}")
|
|
|
|
reference = json.loads(BASE_REFERENCE.read_text())
|
|
reference["frontier_checkout"] = str(frontier)
|
|
reference["frontier_commit"] = EXPERIMENT_COMMIT
|
|
generated_reference = ROOT / "frontier-reference.json"
|
|
generated_reference.write_text(json.dumps(reference, indent=2))
|
|
|
|
module = load_s3_module()
|
|
module.REFERENCE = generated_reference
|
|
module.EXPECTED_FRONTIER_COMMIT = EXPERIMENT_COMMIT
|
|
original_replace = module.replace_flag
|
|
|
|
def replace_and_override(argv: list[str], flag: str, value: str) -> None:
|
|
original_replace(argv, flag, value)
|
|
if flag.endswith("trace_file"):
|
|
atten_flag = (
|
|
"--random_forrest_execution_time_predictor_config_atten_input_file"
|
|
)
|
|
original_replace(argv, atten_flag, str(profile))
|
|
no_cache = (
|
|
"--random_forrest_execution_time_predictor_config_no_cache"
|
|
)
|
|
if no_cache in argv:
|
|
argv.remove(no_cache)
|
|
|
|
module.replace_flag = replace_and_override
|
|
module.parse_args = lambda: args
|
|
module.main()
|
|
|
|
manifest_path = args.output_root / "manifest.json"
|
|
manifest = json.loads(manifest_path.read_text())
|
|
manifest.update(
|
|
{
|
|
"schema": "frontier-attn-structured-replay-v1",
|
|
"frontier_base_commit": BASE_COMMIT,
|
|
"frontier_experiment_commit": EXPERIMENT_COMMIT,
|
|
"frontier_patch": str(PATCH.resolve()),
|
|
"frontier_patch_sha256": module.sha256(PATCH),
|
|
"attention_profile_override": str(profile),
|
|
"attention_profile_sha256": module.sha256(profile),
|
|
"model_cache_enabled": True,
|
|
}
|
|
)
|
|
manifest_path.write_text(json.dumps(manifest, indent=2))
|
|
print(f"structured replay done: {args.output_root}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|