Add module swap quantization API from Quanty

**Summary:** This commit adds a module-swap-based PTQ API from
Quanty, including:

- Quantized linear and embedding modules
- `IntQuantizer` to specify how to quantize weights and activations
- `CodeBookQuantizer` as an alternative to IntQuantizer
- Implementation of K-means to be used for codebook quantization
- Range setting and data getter utility

These new APIs will complement our existing `quantize_` API,
which is primarily used for tensor-subclass-based quantization
today (though it can also support module swaps). All APIs
introduced in this commit are under prototype and highly subject
to change. In particular, we plan to delete `quantize_module_swap`
and `QuantizationRecipe`, and instead integrate this flow with
the `quantize_` API by creating a new `AOBaseConfig`.

All code is migrated from Quanty and written by @TiRune.

**Test Plan:**
python test/quantization/module_swap/test_*

Author: andrewor14

ruff.toml

@@ -7,4 +7,5 @@ lint.ignore = ["E731"]
 # Exclude third-party modules
 exclude = [
     "third_party/*",
+    "prototype/quantization/module_swap/*",
 ]

test/quantization/module_swap/test_kmeans_codebook.py

@@ -0,0 +1,56 @@
+import copy
+import unittest
+
+import torch
+import torch.nn as nn
+
+from torchao.prototype.quantization.module_swap.algorithms import kmeans_codebook
+from torchao.prototype.quantization.module_swap import (
+    CodeBookQuantizer,
+    QuantizedLinear,
+)
+
+
+class SimpleTestNetwork(nn.Module):
+    def __init__(self, weight_group_size: str | int = "per_channel") -> None:
+        super().__init__()
+        if weight_group_size == "per_channel":
+            weight_group_size = 8
+        assert isinstance(weight_group_size, int)
+        weight_quantizer = CodeBookQuantizer(
+            n_bits=2,
+            features=16,
+            codebook_dim=2,
+        )
+
+        self.linear = QuantizedLinear(
+            in_features=16,
+            out_features=8,
+            bias=False,
+            weight_quantizer=weight_quantizer,
+            activation_bits=8,
+            input_quantization=False,
+            output_quantization=False,
+            weight_quantization=True,
+            activation_quantization=False,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.linear(x)
+
+
+class TestKmeansCodebook(unittest.TestCase):
+
+    @unittest.skip("No module named 'faiss'")
+    def test_kmeans_codebook(self) -> None:
+        model = SimpleTestNetwork()
+        codebook_before = copy.deepcopy(model.linear.weight_quantizer.codebook)
+        kmeans_codebook(model)
+        assert not torch.allclose(
+            codebook_before,
+            model.linear.weight_quantizer.codebook,
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/quantization/module_swap/test_llm_ptq_data_getter.py

@@ -0,0 +1,38 @@
+from typing import Tuple
+import unittest
+
+import torch
+
+from torchao.prototype.quantization.module_swap.data_getters import LLMPTQDataGetter
+
+from transformers.models.llama.modeling_llama import LlamaConfig, LlamaForCausalLM
+
+
+test_config = LlamaConfig(
+    vocab_size=10,
+    hidden_size=32,
+    num_hidden_layers=2,
+    num_attention_heads=2,
+    intermediate_size=64,
+)
+
+
+def get_test_llama_model_data() -> Tuple[LlamaForCausalLM, torch.Tensor]:
+    model = LlamaForCausalLM(test_config)
+    input_ids = torch.randint(0, test_config.vocab_size, (1, 10))
+    return model, input_ids
+
+
+class TestPTQDataGetter(unittest.TestCase):
+
+    @unittest.skip("TypeError: cannot unpack non-iterable NoneType object")
+    def test_data_getter(self) -> None:
+        model, data = get_test_llama_model_data()
+        data_getter = LLMPTQDataGetter(model, data, 1)
+        for name, module in model.named_modules():
+            if isinstance(module, torch.nn.Linear):
+                data = data_getter.pop(model, name)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/quantization/module_swap/test_module_swap.py

@@ -0,0 +1,35 @@
+import unittest
+
+import torch
+import torch.nn as nn
+
+from torchao.prototype.quantization.module_swap.module_swap import (
+    QuantizationRecipe,
+    quantize_module_swap,
+)
+
+
+class SimpleEmbeddingTestNetwork(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.embedding = nn.Embedding(10, 64)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.embedding(x)
+
+
+class TestEmbeddingSwap(unittest.TestCase):
+    def test_embedding_swap(self) -> None:
+        model = SimpleEmbeddingTestNetwork()
+        recipe = QuantizationRecipe()
+        recipe.embedding_bits = 4
+        recipe.embedding_quantization = True
+        model = quantize_module_swap(model, recipe)
+        x = torch.randint(0, 10, (10, 64))
+        model(x)
+        assert model.embedding.weight_quantizer.num_bits == 4
+        assert model.embedding.weight_quantizer.group_size == 32
+
+
+if __name__ == "__main__":
+    unittest.main()