From 025375dfabc08aea492063049d9133d646f696e3 Mon Sep 17 00:00:00 2001 From: andrewor14 Date: Wed, 16 Oct 2024 07:52:14 -0700 Subject: [PATCH] Add missing QAT linear classes in BC path Summary: https://github.com/pytorch/ao/pull/1037 moved the module API and attempted to maintain BC, but it seems it missed the linear classes. This unintentionally broke some workflows. This commit adds the missing linear classes back until we are ready to fully remove this path. --- torchao/quantization/prototype/qat/_module_swap_api.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/torchao/quantization/prototype/qat/_module_swap_api.py b/torchao/quantization/prototype/qat/_module_swap_api.py index 0b44974f2..532a3e1e5 100644 --- a/torchao/quantization/prototype/qat/_module_swap_api.py +++ b/torchao/quantization/prototype/qat/_module_swap_api.py @@ -2,7 +2,9 @@ # These will be removed in the future from .linear import ( + Int8DynActInt4WeightQATLinear, Int8DynActInt4WeightQATQuantizer as Int8DynActInt4WeightQATQuantizerModuleSwap, + Int4WeightOnlyQATLinear, Int4WeightOnlyQATQuantizer as Int4WeightOnlyQATQuantizerModuleSwap, enable_8da4w_fake_quant as enable_8da4w_fake_quant_module_swap, disable_8da4w_fake_quant as disable_8da4w_fake_quant_module_swap,