| 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 | |