70 lines
2.2 KiB
Python
70 lines
2.2 KiB
Python
#!/usr/bin/env python3
|
|
|
|
from __future__ import annotations
|
|
|
|
import csv
|
|
import importlib.util
|
|
import json
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
|
|
ROOT = Path(__file__).parent
|
|
|
|
|
|
def load(name: str):
|
|
path = ROOT / name
|
|
spec = importlib.util.spec_from_file_location(path.stem, path)
|
|
assert spec and spec.loader
|
|
module = importlib.util.module_from_spec(spec)
|
|
sys.modules[spec.name] = module
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
class FidelityEnvelopeTest(unittest.TestCase):
|
|
def test_materialize_allreduce(self) -> None:
|
|
module = load("materialize_frontier_allreduce.py")
|
|
rows = []
|
|
for tp in (2, 4):
|
|
for tokens in (1, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192):
|
|
rows.append(
|
|
{
|
|
"tensor_parallel_size": tp,
|
|
"num_tokens": tokens,
|
|
"hidden_dim": 2048,
|
|
"payload_bytes": tokens * 2048 * 2,
|
|
"critical_path_median_ms": tp + tokens / 1000,
|
|
}
|
|
)
|
|
with tempfile.TemporaryDirectory() as temporary:
|
|
root = Path(temporary)
|
|
source = root / "allreduce.json"
|
|
source.write_text(
|
|
json.dumps(
|
|
{
|
|
"schema_version": "qwen30_vllm020_allreduce_frozen.v1",
|
|
"rows": rows,
|
|
}
|
|
)
|
|
)
|
|
output = root / "all_reduce.csv"
|
|
manifest = module.convert(source, output)
|
|
self.assertEqual(manifest["rows"], 24)
|
|
with output.open(newline="") as handle:
|
|
converted = list(csv.DictReader(handle))
|
|
self.assertEqual(converted[0]["num_workers"], "2")
|
|
self.assertEqual(converted[0]["size"], "4096")
|
|
self.assertEqual(converted[-1]["num_workers"], "4")
|
|
self.assertEqual(converted[-1]["size"], str(8192 * 2048 * 2))
|
|
self.assertEqual(
|
|
converted[-1]["time_stats.all_reduce.median"],
|
|
str(4 + 8192 / 1000),
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|