Support QAT int4 v1 path for BC (#2888)

**Summary:** `Int4WeightOnlyConfig` supports version 1 (targeting
tinygemm) and version 2 (targeting fbgemm). However, the latter
requires a new dependency (fbgemm_gpu_genai >= 1.2.0), which is
problematic for torchao integrations with other frameworks.
For now, we should continue to support the v1 path for BC.

**Test Plan:**
```
python test/quantization/test_qat.py -k
test_infer_int4_weight_only_config
```

Author: andrewor14

test/quantization/test_qat.py

@@ -69,6 +69,7 @@
 from torchao.quantization.quant_api import (
     Float8DynamicActivationFloat8WeightConfig,
     Float8DynamicActivationInt4WeightConfig,
+    Int4WeightOnlyConfig,
     Int8DynamicActivationInt4WeightConfig,
 )
 from torchao.quantization.quant_primitives import (
@@ -1933,6 +1934,22 @@ def test_quantize_api_fp8_int4(self):
         )
 
     @unittest.skipIf(not _CUDA_IS_AVAILABLE, "skipping when cuda is not available")
+    @unittest.skipIf(
+        not _is_fbgemm_genai_gpu_available(), "Requires fbgemm-gpu-genai >= 1.2.0"
+    )
+    @parametrize("version", [1, 2])
+    def test_quantize_api_int4(self, version: int):
+        """
+        Test the following:
+            quantize_(model, QATConfig(Int4WeightOnlyConfig(), step="prepare"))
+            quantize_(model, QATConfig(Int4WeightOnlyConfig(), step="convert"))
+        """
+        self._test_quantize_api_against_ptq(
+            Int4WeightOnlyConfig(version=version),
+            target_prepare_sqnr=12,
+            target_convert_sqnr=float("inf"),
+        )
+
     def test_infer_fp8_int4_config(self):
         """
         Test that fake quantize configs are correctly inferred from
@@ -1952,6 +1969,29 @@ def test_infer_fp8_int4_config(self):
         self.assertEqual(weight_config.group_size, 128)
         self.assertTrue(weight_config.is_symmetric)
 
+    def test_infer_int4_weight_only_config(self):
+        """
+        Test that fake quantize configs are correctly inferred from `Int4WeightOnlyConfig`.
+        """
+        from torchao.quantization.qat.fake_quantize_config import (
+            _infer_fake_quantize_configs,
+        )
+
+        base_config = Int4WeightOnlyConfig(version=1)
+        (act_config, weight_config) = _infer_fake_quantize_configs(base_config)
+        self.assertIsNone(act_config)
+        self.assertIsInstance(weight_config, IntxFakeQuantizeConfig)
+        self.assertEqual(weight_config.dtype, torch.uint4)
+        self.assertEqual(weight_config.group_size, 128)
+        self.assertFalse(weight_config.is_symmetric)
+
+        base_config = Int4WeightOnlyConfig(version=2)
+        (act_config, weight_config) = _infer_fake_quantize_configs(base_config)
+        self.assertIsNone(act_config)
+        self.assertEqual(weight_config.dtype, torch.int4)
+        self.assertEqual(weight_config.group_size, 128)
+        self.assertTrue(weight_config.is_symmetric)
+
     @unittest.skipIf(not is_sm_at_least_89(), "Need sm89+")
     def test_quantize_api_nvfp4(self):
         """

torchao/quantization/qat/fake_quantize_config.py

@@ -358,14 +358,31 @@ def _infer_fake_quantize_configs(
             is_symmetric=base_config.mapping_type == MappingType.SYMMETRIC,
         )
     elif isinstance(base_config, Int4WeightOnlyConfig):
-        if base_config.version != 2:
-            raise ValueError(f"Only version 2 of {type(base_config)} is supported")
         act_config = None
-        weight_config = IntxFakeQuantizeConfig(
-            dtype=torch.int4,
-            group_size=base_config.group_size,
-            is_symmetric=True,
-        )
+        if base_config.version == 2:
+            weight_config = IntxFakeQuantizeConfig(
+                dtype=torch.int4,
+                group_size=base_config.group_size,
+                is_symmetric=True,
+            )
+        elif base_config.version == 1:
+            # For BC
+            from torchao.quantization.quant_api import (
+                LAYOUT_TO_ZERO_POINT_DOMAIN,
+            )
+
+            if base_config.zero_point_domain == ZeroPointDomain.NONE:
+                zp_domain = LAYOUT_TO_ZERO_POINT_DOMAIN[type(base_config.layout)][0]
+            else:
+                zp_domain = base_config.zero_point_domain
+            weight_config = IntxFakeQuantizeConfig(
+                dtype=torch.uint4,
+                group_size=base_config.group_size,
+                is_symmetric=False,
+                zero_point_domain=zp_domain,
+            )
+        else:
+            raise ValueError(f"Unknown version on base config {type(base_config)}")
     elif isinstance(base_config, Float8DynamicActivationFloat8WeightConfig):
         if base_config.version != 2:
             raise ValueError(f"Only version 2 of {type(base_config)} is supported")