torchao_int4_weight_only_quantization_benchmark_with_tinygemm.py

1import torch
2from torchao.quantization import quantize_, int4_weight_only
3import torch.utils.benchmark as benchmark
4
5# 1. Create a model
6model = torch.nn.Sequential(
7    torch.nn.Linear(1024, 1024),
8    torch.nn.ReLU(),
9    torch.nn.Linear(1024, 1024)
10).cuda().to(torch.bfloat16)
11
12# 2. Quantize the model
13# apply int4 weight-only quantization using the tinygemm kernel
14quantize_(model, int4_weight_only())
15
16# 3. Compile the model to fuse kernels and speed up execution
17model = torch.compile(model, mode="max-autotune")
18
19# 4. Benchmark the model
20input_tensor = torch.randn(1024, 1024, device="cuda", dtype=torch.bfloat16)
21
22def benchmark_model(model, input_tensor):
23    with torch.no_grad():
24        model(input_tensor)
25
26t0 = benchmark.Timer(
27    stmt='benchmark_model(model, input_tensor)',
28    setup='from __main__ import benchmark_model',
29    globals={'model': model, 'input_tensor': input_tensor}
30)
31
32print(f"Int4 Weight-only quantization execution time: {t0.timeit(100).mean * 1000:.3f} ms")