torchao_int8_weight_only_quantization_quickstart.py

1import torch
2from torchao.quantization import quantize_, int8_weight_only
3
4# 1. Define or load a model
5model = torch.nn.Sequential(
6    torch.nn.Linear(32, 64),
7    torch.nn.ReLU(),
8    torch.nn.Linear(64, 32)
9).cuda().to(torch.bfloat16)
10
11# 2. Apply quantization
12# This transforms the linear layers to use 8-bit weights
13quantize_(model, int8_weight_only())
14
15# 3. Run inference
16input_data = torch.randn(1, 32, device="cuda", dtype=torch.bfloat16)
17output = model(input_data)
18
19print(f"Output shape: {output.shape}")
20print("Quantization applied successfully!")