File size: 930 Bytes
36c78b1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 | import os, sys; sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
import torch
from bit_transformer import BitTransformerLM, quantize_dynamic, prepare_qat_fx, convert_qat_fx
from bit_transformer.training import train_loop
def test_qat_matches_dynamic_quant():
data = torch.randint(0, 2, (16, 8), dtype=torch.long)
base = BitTransformerLM(d_model=16, nhead=4, num_layers=1, dim_feedforward=32, max_seq_len=8)
train_loop(base, data, epochs=1, log=False)
dyn = quantize_dynamic(base)
qat_model = BitTransformerLM(d_model=16, nhead=4, num_layers=1, dim_feedforward=32, max_seq_len=8)
qat_model.load_state_dict(base.state_dict())
prepare_qat_fx(qat_model)
convert_qat_fx(qat_model)
inp = torch.randint(0, 2, (10, 8), dtype=torch.long)
out_dyn, _ = dyn(inp)
out_qat, _ = qat_model(inp)
diff = (out_dyn - out_qat).abs().max().item()
assert diff < 0.6
|