Files
aituner/runs/frontier-fidelity-envelope-v1/materialize_qwen235_v020_allreduce.py

214 lines
8.8 KiB
Python

#!/usr/bin/env python3
"""Validate Qwen235 serving collectives and materialize Frontier's CC CSV."""
from __future__ import annotations
import argparse
import csv
import hashlib
import json
import math
from pathlib import Path
from typing import Any
VLLM_COMMIT = "88d34c6409e9fb3c7b8ca0c04756f061d2099eb1"
TOKEN_POINTS = (
1, 2, 4, 8, 16, 24, 32, 40, 48, 56, 63, 64, 65, 72, 80, 88, 96,
104, 112, 120, 128, 136, 144, 152, 160, 168, 176, 184, 192, 200,
208, 216, 224, 232, 240, 248, 255, 256, 257, 512, 1024, 2048, 4096,
8192,
)
TP_SIZES = (4, 8)
HIDDEN_DIM = 4096
ALLOWED_BACKENDS = {
"flashinfer_trtllm_fused_projection",
"pynccl_symmetric_with_copy",
"torch_symmetric_memory",
"pynccl",
}
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 materialize(
input_paths: list[Path], output_path: Path, manifest_path: Path
) -> dict[str, Any]:
if len(input_paths) != 2:
raise ValueError("exactly two raw inputs are required for TP4 and TP8")
rows = []
seen = set()
backends_by_tp: dict[int, set[str]] = {}
environments = []
for path in input_paths:
payload = json.loads(path.read_text())
if payload.get("schema_version") != "vllm020_allreduce_raw.v2":
raise ValueError(f"unexpected raw collective schema in {path}")
environment = payload["environment"]
if environment.get("vllm_version") != "0.20.0":
raise ValueError(f"vLLM version mismatch in {path}")
if environment.get("vllm_source_commit") != VLLM_COMMIT:
raise ValueError(f"vLLM commit mismatch in {path}")
if environment.get("gpu") != "NVIDIA H20":
raise ValueError(f"GPU mismatch in {path}: {environment.get('gpu')!r}")
if Path(environment.get("model", "")).name != "Qwen3-235B-A22B-FP8":
raise ValueError(f"model mismatch in {path}")
if environment.get("collective_contract") != "qwen235-serving-projected":
raise ValueError(f"collective contract mismatch in {path}")
if environment.get("disable_custom_all_reduce") is not True:
raise ValueError(f"custom all-reduce was not disabled in {path}")
backend_env = environment.get("backend_env", {})
if backend_env.get("VLLM_ALLREDUCE_USE_FLASHINFER") != "1":
raise ValueError(f"FlashInfer projection was not enabled in {path}")
if backend_env.get("VLLM_ALLREDUCE_USE_SYMM_MEM") != "1":
raise ValueError(f"symmetric-memory contract mismatch in {path}")
environments.append(environment)
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 TP_SIZES:
raise ValueError(f"unexpected TP: {tp}")
if int(source["hidden_dim"]) != HIDDEN_DIM:
raise ValueError(f"hidden-dim mismatch for {key}")
expected_bytes = tokens * HIDDEN_DIM * 2
if int(source["payload_bytes"]) != expected_bytes:
raise ValueError(f"payload mismatch for {key}")
backend = source["selected_backend"]
if backend not in ALLOWED_BACKENDS:
raise ValueError(f"unexpected backend for {key}: {backend!r}")
fusion_limit = {4: 2 * 1024 * 1024, 8: 512 * 1024}[tp]
if int(source["real_fusion_limit_bytes"]) != fusion_limit:
raise ValueError(f"fusion limit mismatch for {key}")
expected_backend = (
"flashinfer_trtllm_fused_projection"
if expected_bytes <= fusion_limit
else None
)
if expected_backend is not None and backend != expected_backend:
raise ValueError(f"missing fused projection for {key}")
if expected_backend is None and backend == "flashinfer_trtllm_fused_projection":
raise ValueError(f"fused projection exceeds runtime limit for {key}")
backends_by_tp.setdefault(tp, set()).add(backend)
median = float(source["critical_path_median_ms"])
if not math.isfinite(median) or median <= 0:
raise ValueError(f"invalid critical-path median for {key}")
if int(source.get("trials", 0)) != 3 or int(
source.get("repeats_per_trial", 0)
) != 20:
raise ValueError(f"repeat contract mismatch for {key}")
samples = source.get("per_trial_rank_samples_ms", [])
if (
len(samples) != 3
or any(len(trial) != tp for trial in samples)
or any(
len(rank_samples) != 20
or any(
not math.isfinite(float(value)) or float(value) <= 0
for value in rank_samples
)
for trial in samples
for rank_samples in trial
)
):
raise ValueError(f"raw sample coverage mismatch for {key}")
critical_samples = [
max(float(trial[rank][repeat]) for rank in range(tp))
for trial in samples
for repeat in range(20)
]
ordered = sorted(critical_samples)
measured_median = (ordered[29] + ordered[30]) / 2
if not math.isclose(median, measured_median, rel_tol=1e-6, abs_tol=1e-9):
raise ValueError(f"critical-path median/sample mismatch for {key}")
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 TP_SIZES for tokens in TOKEN_POINTS}
if seen != expected:
raise ValueError(
f"collective coverage mismatch: missing={expected - seen}, extra={seen - expected}"
)
if output_path.exists() or manifest_path.exists():
raise FileExistsError("refusing to overwrite immutable collective profile")
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"])))
manifest = {
"schema": "qwen235-v020-serving-allreduce-profile-v1",
"inputs": {str(path.resolve()): sha256(path) for path in input_paths},
"output": str(output_path.resolve()),
"output_sha256": sha256(output_path),
"vllm_source_commit": VLLM_COMMIT,
"hardware": "NVIDIA H20",
"model": "Qwen3-235B-A22B-FP8",
"dtype": "bfloat16",
"hidden_dim": HIDDEN_DIM,
"collective_contract": "qwen235-serving-projected",
"tp_coverage": list(TP_SIZES),
"token_points": list(TOKEN_POINTS),
"rows": len(rows),
"observed_backends_by_tp": {
str(tp): sorted(backends) for tp, backends in sorted(backends_by_tp.items())
},
"frontier_target": "time_stats.all_reduce.median",
"unused_stat_policy": "repeat critical_path_median; std=0",
"raw_environments": environments,
}
manifest_path.parent.mkdir(parents=True, exist_ok=True)
manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n")
return manifest
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--input", type=Path, action="append", required=True)
parser.add_argument("--output", type=Path, required=True)
parser.add_argument("--manifest", type=Path, required=True)
args = parser.parse_args()
manifest = materialize(args.input, args.output, args.manifest)
print(json.dumps(manifest, sort_keys=True))
if __name__ == "__main__":
main()