116 lines
4.1 KiB
Python
116 lines
4.1 KiB
Python
#!/usr/bin/env python3
|
|
"""Convert frozen vLLM collective measurements to Frontier Vidur CC CSV."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import csv
|
|
import hashlib
|
|
import json
|
|
from pathlib import Path
|
|
|
|
|
|
FIELDS = (
|
|
"time_stats.all_reduce.min",
|
|
"time_stats.all_reduce.max",
|
|
"time_stats.all_reduce.mean",
|
|
"time_stats.all_reduce.median",
|
|
"time_stats.all_reduce.std",
|
|
"rank",
|
|
"num_workers",
|
|
"size",
|
|
"collective",
|
|
"devices_per_node",
|
|
"max_devices_per_node",
|
|
)
|
|
|
|
|
|
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 convert(input_path: Path, output_path: Path) -> dict[str, object]:
|
|
payload = json.loads(input_path.read_text())
|
|
if payload.get("schema_version") != "qwen30_vllm020_allreduce_frozen.v1":
|
|
raise ValueError(f"unexpected input schema: {payload.get('schema_version')!r}")
|
|
|
|
rows = []
|
|
seen = set()
|
|
for source in payload["rows"]:
|
|
tp = int(source["tensor_parallel_size"])
|
|
tokens = int(source["num_tokens"])
|
|
key = (tp, tokens)
|
|
if key in seen:
|
|
raise ValueError(f"duplicate collective row: {key}")
|
|
seen.add(key)
|
|
if tp not in (2, 4):
|
|
raise ValueError(f"unsupported TP: {tp}")
|
|
expected_bytes = tokens * int(source["hidden_dim"]) * 2
|
|
if int(source["payload_bytes"]) != expected_bytes:
|
|
raise ValueError(f"payload mismatch for {key}")
|
|
|
|
# Frontier Vidur consumes only the median target. The raw profiler kept
|
|
# per-rank distributions but not aligned per-repeat critical-path
|
|
# samples, so do not invent critical-path min/mean/max/std. Repeating
|
|
# the observed critical-path median in the unused fields keeps the CSV
|
|
# schema explicit without changing the trained target.
|
|
median = float(source["critical_path_median_ms"])
|
|
rows.append(
|
|
{
|
|
"time_stats.all_reduce.min": median,
|
|
"time_stats.all_reduce.max": median,
|
|
"time_stats.all_reduce.mean": median,
|
|
"time_stats.all_reduce.median": median,
|
|
"time_stats.all_reduce.std": 0.0,
|
|
"rank": 0,
|
|
"num_workers": tp,
|
|
"size": expected_bytes,
|
|
"collective": "all_reduce",
|
|
"devices_per_node": tp,
|
|
"max_devices_per_node": 8,
|
|
}
|
|
)
|
|
|
|
expected = {(tp, tokens) for tp in (2, 4) for tokens in (1, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192)}
|
|
if seen != expected:
|
|
raise ValueError(f"collective coverage mismatch: missing={expected - seen}, extra={seen - expected}")
|
|
|
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
|
with output_path.open("w", newline="") as output:
|
|
writer = csv.DictWriter(output, fieldnames=FIELDS, lineterminator="\n")
|
|
writer.writeheader()
|
|
writer.writerows(sorted(rows, key=lambda row: (row["num_workers"], row["size"])))
|
|
|
|
return {
|
|
"schema": "frontier-vidur-allreduce-materialization-v1",
|
|
"source": str(input_path.resolve()),
|
|
"source_sha256": sha256(input_path),
|
|
"output": str(output_path.resolve()),
|
|
"output_sha256": sha256(output_path),
|
|
"rows": len(rows),
|
|
"tp_coverage": [2, 4],
|
|
"target": "time_stats.all_reduce.median",
|
|
"unused_stat_policy": "repeat critical_path_median; std=0",
|
|
"payload_contract": "size=num_tokens*hidden_dim*2_bytes",
|
|
}
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--input", type=Path, required=True)
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
parser.add_argument("--manifest", type=Path, required=True)
|
|
args = parser.parse_args()
|
|
manifest = convert(args.input, args.output)
|
|
args.manifest.parent.mkdir(parents=True, exist_ok=True)
|
|
args.manifest.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n")
|
|
print(json.dumps(manifest, sort_keys=True))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|