GPTQ updates

Summary:

1) reorganized GPTQ
 a) got rid of old GPTQ and renamed GPTQ_MT to GPTQ
 b) moved new GPTQ to prototype
 c) moved quantized linear modules in GPTQ.py to linear_quant_modules.py
2) removed dependence on lm_eval for input_recorder
 a) created new input recorder that doesn't depend on lm_eval
 b) made lm_eval input recorder depend on new generic input_recorder
 c) made TransformerEvalWrapper the base class and made
 d) updated apis generally to work with new input recorder
 LMEvalInputRecorder inherit from it instead of vice-versa
3) reorganized GPTQ tests
 a) moved tests from test_quant_api.py to test_gptq.py
 b) added new test that can be run in CI that doesn't depend on
 lm_eval/llama weights
 c) got rid of test_gptq_mt.py
4) added new documentation for lm_eval
5) GPTQ improvements
 a) reimplemented faster quant
 b) tested compilation of hessian calculation and parts of faster quant,
 generally they were slower.
 c) moved helper functions out of the class. They're largely generic and
 this is less cluttered.
 d) some improvements to the duplication checking and copying to be
 faster when possible
 e) fixed some bugs due to this not being in CI and things changing for
 int4wo tensor subclass.

Test Plan:

1) `python test_gptq.py`

note: the skipped test test_gptq_quantizer_int4_weight_only also ran.

2) I verified that all activation match between old GPTQ and current
   GPTQ

3)

```shell

export CHECKPOINT_PATH=../../../checkpoints # path to checkpoints folder

export MODEL_REPO=meta-llama/Llama-2-7b-chat-hf
python eval.py --checkpoint_path $CHECKPOINT_PATH/$MODEL_REPO/model.pth
python eval.py --checkpoint_path $CHECKPOINT_PATH/$MODEL_REPO/model.pth
--quantization int4wo-64
python eval.py --checkpoint_path $CHECKPOINT_PATH/$MODEL_REPO/model.pth
--quantization int4wo-gptq-64 --calibration_limit 10

export MODEL_REPO=meta-llama/Meta-Llama-3-8B
python eval.py --checkpoint_path $CHECKPOINT_PATH/$MODEL_REPO/model.pth
python eval.py --checkpoint_path $CHECKPOINT_PATH/$MODEL_REPO/model.pth
--quantization int4wo-64
python eval.py --checkpoint_path $CHECKPOINT_PATH/$MODEL_REPO/model.pth
--quantization int4wo-gptq-64 --calibration_limit 10

```
see README.md for results but they show GPTQ is working

Reviewers:

Subscribers:

Tasks:

Tags:

Author: HDCharles

test/quantization/test_gptq.py

@@ -0,0 +1,202 @@
+import unittest
+from pathlib import Path
+
+import torch
+from torch.testing._internal.common_utils import TestCase
+
+from torchao._models.llama.model import (
+    ModelArgs,
+    Transformer,
+    prepare_inputs_for_model,
+)
+from torchao._models.llama.tokenizer import get_tokenizer
+from torchao.quantization import Int4WeightOnlyConfig, quantize_
+from torchao.quantization.utils import compute_error
+from torchao.utils import (
+    TORCH_VERSION_AT_LEAST_2_4,
+)
+
+torch.manual_seed(0)
+
+
+class TestGPTQ(TestCase):
+    @unittest.skip("skipping until we get checkpoints for gpt-fast")
+    def test_gptq_quantizer_int4_weight_only(self):
+        from torchao._models._eval import (
+            LMEvalInputRecorder,
+            TransformerEvalWrapper,
+        )
+        from torchao.prototype.GPTQ.GPTQ import Int4WeightOnlyGPTQQuantizer
+
+        precision = torch.bfloat16
+        device = "cuda"
+        checkpoint_path = Path(
+            "../../checkpoints/meta-llama/Llama-2-7b-chat-hf/model.pth"
+        )
+        model = Transformer.from_name(checkpoint_path.parent.name)
+        checkpoint = torch.load(str(checkpoint_path), mmap=True, weights_only=True)
+        model.load_state_dict(checkpoint, assign=True)
+        model = model.to(dtype=precision, device="cpu")
+        model.eval()
+
+        tokenizer_path = checkpoint_path.parent / "tokenizer.model"
+        assert tokenizer_path.is_file(), tokenizer_path
+        tokenizer = get_tokenizer(  # pyre-ignore[28]
+            tokenizer_path,
+            "Llama-2-7b-chat-hf",
+        )
+        groupsize = 64
+        blocksize = 128
+        percdamp = 0.01
+        calibration_tasks = ["wikitext"]
+        calibration_limit = 1
+        calibration_seq_length = 100
+        input_prep_func = prepare_inputs_for_model
+        pad_calibration_inputs = False
+        inputs = (
+            LMEvalInputRecorder(
+                tokenizer,
+                calibration_seq_length,
+                input_prep_func,
+                model.config.vocab_size,
+                pad_calibration_inputs,
+                device="cpu",
+            )
+            .record_inputs(
+                calibration_tasks,
+                calibration_limit,
+            )
+            .get_inputs()
+        )
+
+        quantizer = Int4WeightOnlyGPTQQuantizer(
+            groupsize,
+            blocksize,
+            percdamp,
+        )
+        model.setup_caches(max_batch_size=1, max_seq_length=calibration_seq_length)
+
+        model = quantizer.quantize(model, *inputs).cuda()
+
+        model.reset_caches()
+        with torch.device("cuda"):
+            model.setup_caches(max_batch_size=1, max_seq_length=model.config.block_size)
+
+        limit = 1
+        result = TransformerEvalWrapper(
+            model.cuda(),
+            tokenizer,
+            model.config.block_size,
+            prepare_inputs_for_model,
+            device,
+        ).run_eval(
+            ["wikitext"],
+            limit,
+        )
+
+        assert result["results"]["wikitext"]["word_perplexity,none"] < 7.77, (
+            f"accuracy regressed from 7.76 to {result['results']['wikitext']['word_perplexity,none']}"
+        )
+
+
+class TestMultiTensorFlow(TestCase):
+    @unittest.skipIf(not TORCH_VERSION_AT_LEAST_2_4, "Test only enabled for 2.4+")
+    @unittest.skipIf(not torch.cuda.is_available(), "Need CUDA available")
+    def test_multitensor_add_tensors(self):
+        from torchao.prototype.GPTQ.GPTQ import MultiTensor
+
+        tensor1 = torch.randn(3, 3)
+        tensor2 = torch.randn(3, 3)
+        mt = MultiTensor(tensor1)
+        mt.add_tensors(tensor2)
+        self.assertEqual(mt.count, 2)
+        self.assertTrue(torch.equal(mt.values[0], tensor1))
+        self.assertTrue(torch.equal(mt.values[1], tensor2))
+
+    @unittest.skipIf(not TORCH_VERSION_AT_LEAST_2_4, "Test only enabled for 2.4+")
+    @unittest.skipIf(not torch.cuda.is_available(), "Need CUDA available")
+    def test_multitensor_pad_unpad(self):
+        from torchao.prototype.GPTQ.GPTQ import MultiTensor
+
+        tensor1 = torch.randn(3, 3)
+        mt = MultiTensor(tensor1)
+        mt.pad_to_length(3)
+        self.assertEqual(mt.count, 3)
+        mt.unpad()
+        self.assertEqual(mt.count, 1)
+
+    @unittest.skipIf(not TORCH_VERSION_AT_LEAST_2_4, "Test only enabled for 2.4+")
+    @unittest.skipIf(not torch.cuda.is_available(), "Need CUDA available")
+    def test_multitensor_inplace_operation(self):
+        from torchao.prototype.GPTQ.GPTQ import MultiTensor
+
+        tensor1 = torch.ones(3, 3)
+        mt = MultiTensor(tensor1)
+        mt += 1  # In-place addition
+        self.assertTrue(torch.equal(mt.values[0], torch.full((3, 3), 2)))
+
+
+class TestMultiTensorInputRecorder(TestCase):
+    def test_multitensor_input_recorder(self):
+        from torchao.prototype.GPTQ.GPTQ import MultiTensor, MultiTensorInputRecorder
+
+        input_recorder = MultiTensorInputRecorder()
+        in1 = ([1], torch.randn(3, 3), (1, "dog", torch.randn(3, 3)), torch.float)
+        in2 = ([1], torch.randn(3, 3), (1, "dog", torch.randn(3, 3)), torch.float)
+
+        input_recorder(*in1)
+        input_recorder(*in2)
+
+        MT_input = input_recorder.get_recorded_inputs()
+
+        self.assertEqual(MT_input[0], [1])
+        self.assertTrue(isinstance(MT_input[1], MultiTensor))
+        self.assertTrue(isinstance(MT_input[2], tuple))
+        self.assertEqual(MT_input[2][0], 1)
+        self.assertEqual(MT_input[2][1], "dog")
+        self.assertTrue(isinstance(MT_input[2][2], MultiTensor))
+        self.assertEqual(MT_input[3], torch.float)
+
+    def test_gptq_with_input_recorder(self):
+        from torchao.prototype.GPTQ.GPTQ import (
+            Int4WeightOnlyGPTQQuantizer,
+            MultiTensorInputRecorder,
+        )
+
+        torch.set_default_dtype(torch.bfloat16)
+
+        config = ModelArgs(n_layer=2)
+
+        with torch.device("cuda"):
+            model = Transformer(config)
+            model.setup_caches(max_batch_size=2, max_seq_length=100)
+            idx = torch.randint(1, 10000, (10, 2, 50)).to(torch.int32)
+            test_input = prepare_inputs_for_model(idx[0])
+        import copy
+
+        model2 = copy.deepcopy(model)
+        out = model(*test_input)
+        quantize_(model2, Int4WeightOnlyConfig())
+
+        outq = model2(*test_input)
+        del model2
+
+        input_recorder = MultiTensorInputRecorder()
+        for i in range(10):
+            input = prepare_inputs_for_model(idx[i])
+            input_recorder(*input)
+
+        multi_tensor_inputs = input_recorder.get_recorded_inputs()
+
+        quantizer = Int4WeightOnlyGPTQQuantizer()
+
+        quantizer.quantize(model, multi_tensor_inputs)
+
+        outgptq = model(*test_input)
+
+        self.assertGreater(compute_error(outgptq, out), 30)
+        self.assertGreater(compute_error(outgptq, out), compute_error(outq, out))
+
+
+if __name__ == "__main__":
+    unittest.main()