torchao_int8_weight_only_quantization_quickstart.py

1import torch
2import torchao
3from torchao.quantization import quantize_, int8_weight_only
4
5# 1. Create a model
6model = torch.nn.Sequential(
7    torch.nn.Linear(32, 64),
8    torch.nn.ReLU(),
9    torch.nn.Linear(64, 32)
10).cuda().to(torch.bfloat16)
11
12# 2. Apply quantization
13# This converts the linear layer weights to int8
14quantize_(model, int8_weight_only())
15
16# 3. Run inference
17input_tensor = torch.randn(1, 32, device="cuda", dtype=torch.bfloat16)
18output = model(input_tensor)
19
20print(f"Output shape: {output.shape}")
21print(f"Model quantized successfully.")