Update PARQ README

Author: lisjin

test/prototype/test_parq.py

@@ -21,6 +21,7 @@
 from torchao.prototype.parq.quant import (
     Int4UnifTorchaoQuantizer,
     LSBQuantizer,
+    TernaryUnifQuantizer,
     UnifQuantizer,
     UnifTorchaoQuantizer,
 )
@@ -102,14 +103,14 @@ def setUp(self):
     @common_utils.parametrize("unif_quant", [True, False])
     @common_utils.parametrize("hard_prox", [True, False])
     def test_parq_train_loop(self, b: int = 2, unif_quant=True, hard_prox=True):
-        if unif_quant and b == 0:
-            self.skipTest("Ternary uniform quantization not yet supported")
-
         self.model.reset_parameters()
         param_groups = build_param_groups(self.model, b)
         base_optimizer = torch.optim.AdamW(param_groups)
 
-        quantizer = UnifQuantizer() if unif_quant else LSBQuantizer()
+        if unif_quant:
+            quantizer = TernaryUnifQuantizer() if b == 0 else UnifQuantizer()
+        else:
+            quantizer = LSBQuantizer()
         prox_map = (
             ProxHardQuant() if hard_prox else ProxPARQ(anneal_start=0, anneal_end=2)
         )

torchao/prototype/parq/README.md

@@ -48,6 +48,7 @@ optimizer = QuantOptimizer(
 <td valign="top">
 
 ```python
+from torchao.quantization import quantize_
 from torchao.quantization.qat import (
     FakeQuantizeConfig,
     intx_quantization_aware_training,
@@ -67,16 +68,20 @@ quantize_(
 <td valign="top">
 
 ```python
-config = IntXWeightOnlyConfig(
+from torchao.quantization import IntxWeightOnlyConfig, quantize_
+
+config = IntxWeightOnlyConfig(
     weight_dtype=torch.int4, granularity=PerGroup(32)
 )
-optimizer.torchao_quantize_(model, config)
+optimizer.restore_latent_params()
+quantize_(model, config, filter_fn=optimizer._get_filter_fn(model))
 ```
 
 </td>
 <td valign="top">
 
 ```python
+from torchao.quantization import quantize_
 from torchao.quantization.qat import from_intx_quantization_aware_training
 
 quantize_(model, from_intx_quantization_aware_training())

torchao/prototype/parq/__init__.py

@@ -11,4 +11,12 @@
     ProxPARQ,
     QuantOptimizer,
 )
-from .quant import LSBQuantizer, Quantizer, UnifQuantizer  # noqa: F401
+from .quant import (  # noqa: F401
+    Int4UnifTorchaoQuantizer,
+    LSBQuantizer,
+    MaxUnifQuantizer,
+    Quantizer,
+    TernaryUnifQuantizer,
+    UnifQuantizer,
+    UnifTorchaoQuantizer,
+)