File size: 3,804 Bytes
8db7cbc | 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 | """Turn on the two speed knobs that actually help Chronos-2 on a GPU.
Weight-only INT8 is smaller, not faster, on an RTX 3090. What *does* help:
1. TF32 — use the GPU's fast float math (off by default in PyTorch).
2. torch.compile — fuse kernels so a tiny 120M model is not launch-bound.
Usage after snapshot_download of this repo:
from load import load
from fast_infer import speedup
pipe = speedup(load(".", device="cuda"))
Or from the command line:
python fast_infer.py --repo oxfrug/chronos-2-int8-torchao
python fast_infer.py --fp32 amazon/chronos-2 # original model, same knobs
"""
from __future__ import annotations
import argparse
import time
import numpy as np
import torch
from chronos import Chronos2Pipeline
def enable_tf32() -> None:
if not torch.cuda.is_available():
return
torch.set_float32_matmul_precision("high")
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
def speedup(pipe: Chronos2Pipeline, compile: bool = True, tf32: bool = True) -> Chronos2Pipeline:
"""Return the same pipeline with TF32 and (optionally) torch.compile."""
if tf32:
enable_tf32()
if compile and torch.cuda.is_available():
pipe.model = torch.compile(pipe.model, mode="reduce-overhead")
return pipe
def _predict_once(pipe: Chronos2Pipeline, y: np.ndarray, h: int) -> None:
pipe.predict_quantiles(inputs=[y], prediction_length=h, quantile_levels=[0.5])
def bench(pipe: Chronos2Pipeline, n_warm: int = 3, n_run: int = 8, h: int = 24) -> float:
y = np.sin(np.linspace(0, 40, 512)).astype(np.float32)
if torch.cuda.is_available():
torch.cuda.synchronize()
for _ in range(n_warm):
_predict_once(pipe, y, h)
if torch.cuda.is_available():
torch.cuda.synchronize()
t0 = time.perf_counter()
for _ in range(n_run):
_predict_once(pipe, y, h)
if torch.cuda.is_available():
torch.cuda.synchronize()
return 1000.0 * (time.perf_counter() - t0) / n_run
def main() -> None:
ap = argparse.ArgumentParser(description="Load Chronos-2 with TF32 + torch.compile and time one call.")
ap.add_argument("--repo", default="oxfrug/chronos-2-int8-torchao")
ap.add_argument("--fp32", default=None, help="Load this HF id with Chronos2Pipeline.from_pretrained instead.")
ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
ap.add_argument("--no-compile", action="store_true")
ap.add_argument("--no-tf32", action="store_true")
args = ap.parse_args()
if args.fp32:
pipe = Chronos2Pipeline.from_pretrained(
args.fp32,
device_map=args.device,
dtype=torch.float32,
)
else:
from pathlib import Path
import importlib.util
here = Path(__file__).resolve().parent
load_py = here / "load.py"
if load_py.exists():
spec = importlib.util.spec_from_file_location("c2int8", load_py)
mod = importlib.util.module_from_spec(spec)
spec.loader.exec_module(mod)
pipe = mod.load(str(here if (here / "model.safetensors").exists() else args.repo), device=args.device)
else:
spec = importlib.util.spec_from_file_location("c2int8", "load.py")
raise SystemExit("need load.py next to this script, or pass --fp32 amazon/chronos-2")
pipe = speedup(pipe, compile=not args.no_compile, tf32=not args.no_tf32)
print("first call compiles; later calls are the ones that count", flush=True)
ms = bench(pipe)
print(f"{ms:.2f} ms / call (512 context, horizon 24, one series)")
print("device", args.device, "compile", not args.no_compile, "tf32", not args.no_tf32)
if __name__ == "__main__":
main()
|