File size: 2,719 Bytes
714625a | 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 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 | #!/usr/bin/env python3
from __future__ import annotations
import os
import subprocess
import sys
from pathlib import Path
ROOT = Path(__file__).resolve().parent
RUNTIME = ROOT / "runtime"
TEST_SCRIPT = RUNTIME / "lab/scripts/test_checkpoint_capabilities.py"
OUTPUT = ROOT / "gpu-triton-results.json"
def main() -> int:
if not TEST_SCRIPT.is_file():
print(f"ERROR: missing capability runner: {TEST_SCRIPT}", file=sys.stderr)
return 1
if not (ROOT / "lit_model.pth").is_file():
print("ERROR: missing lit_model.pth", file=sys.stderr)
return 1
if not (ROOT / "model_config.yaml").is_file():
print("ERROR: missing model_config.yaml", file=sys.stderr)
return 1
try:
import torch
except Exception as exc:
print(f"ERROR: PyTorch import failed: {exc}", file=sys.stderr)
return 1
if not torch.cuda.is_available():
print("ERROR: CUDA is not available.", file=sys.stderr)
return 1
try:
import triton
except Exception as exc:
print(f"ERROR: Triton import failed: {exc}", file=sys.stderr)
return 1
env = os.environ.copy()
env["CUDA_VISIBLE_DEVICES"] = env.get("CUDA_VISIBLE_DEVICES", "0")
env["MULTISCREEN_BACKEND"] = "triton"
env["TOKENIZERS_PARALLELISM"] = "false"
env["PYTHONUNBUFFERED"] = "1"
old_pythonpath = env.get("PYTHONPATH", "")
env["PYTHONPATH"] = (
str(RUNTIME)
if not old_pythonpath
else str(RUNTIME) + os.pathsep + old_pythonpath
)
command = [
sys.executable,
"-u",
str(TEST_SCRIPT),
"--checkpoint-dir",
str(ROOT),
"--tokenizer-dir",
str(ROOT),
"--dtype",
"bfloat16",
"--device",
"cuda",
"--seed",
"1337",
"--output",
str(OUTPUT),
]
print(f"Checkpoint: {ROOT}")
print(f"GPU: {torch.cuda.get_device_name(0)}")
print("Device: cuda")
print("Backend: triton")
print("Triton imported: True")
print(f"Triton version: {triton.__version__}")
print()
result = subprocess.run(command, cwd=ROOT, env=env)
if result.returncode != 0:
print(
f"GPU Triton smoke: FAIL ({result.returncode})",
file=sys.stderr,
)
return result.returncode
if not OUTPUT.is_file() or OUTPUT.stat().st_size == 0:
print(f"ERROR: result was not created: {OUTPUT}", file=sys.stderr)
return 1
print()
print("GPU Triton smoke: PASS")
print("Backend: triton")
print("Triton imported: True")
print(f"Results: {OUTPUT}")
return 0
if __name__ == "__main__":
raise SystemExit(main())
|