File size: 1,903 Bytes
a99edfc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 | """Run with: modal run --profile samuel modal_benchmark.py --gpu L40S"""
import json
from pathlib import Path
import modal
app = modal.App("lfm25-350m-rlcd-benchmark")
image = (modal.Image.debian_slim(python_version="3.11")
.pip_install_from_requirements("requirements.txt")
.add_local_python_source("rlcd")
.add_local_dir("tests", "/root/tests"))
cache = modal.Volume.from_name("lfm25-350m-rlcd-hf-cache", create_if_missing=True)
def execute(suite="diagnostic"):
import os
import subprocess
from rlcd.benchmark import run
os.environ["RLCD_TEST_DEVICE"] = "cuda"
test = subprocess.run(["python", "-m", "pytest", "-q", "-s", "/root/tests"], capture_output=True, text=True)
print(test.stdout, flush=True)
if test.returncode:
raise RuntimeError(test.stdout + test.stderr)
result = run(device="cuda", dtype="float16", warmups=2, repeats=3, suite=suite, max_new_tokens=768 if suite == "stress" else 192)
result["validation"] = {"pytest_stdout": test.stdout, "pytest_stderr": test.stderr, "returncode": test.returncode}
return result
@app.function(image=image, gpu="L40S", timeout=3600, volumes={"/cache": cache}, env={"HF_HOME":"/cache"})
def l40s(suite="diagnostic"):
return execute(suite)
@app.function(image=image, gpu="H100", timeout=3600, volumes={"/cache": cache}, env={"HF_HOME":"/cache"})
def h100(suite="diagnostic"):
return execute(suite)
@app.local_entrypoint()
def main(gpu: str = "L40S", suite: str = "diagnostic"):
if suite not in {"diagnostic", "stress"}:
raise ValueError(suite)
fn = {"L40S":l40s, "H100":h100}[gpu]
result = fn.remote(suite)
path = Path("results") / (gpu.lower() + ("-stress" if suite == "stress" else "") + ".json")
path.parent.mkdir(exist_ok=True)
path.write_text(json.dumps(result, indent=2) + "\n")
print(json.dumps(result["summary"], indent=2))
|