Handle weight-only embeddings in torchao_convert

Author: lisjin

test/prototype/test_parq.py

@@ -161,6 +161,7 @@ def build_param_groups(
     model,
     b: int = 2,
     group_size: Optional[int] = None,
+    embed_b: int = 4,
 ):
     params_quant, params_embed, params_no_quant = split_param_groups(model)
     quant_kwargs = {}
@@ -171,14 +172,27 @@ def build_param_groups(
         {"params": params_no_quant},
     ]
     if params_embed:
-        param_groups.append({"params": params_embed, "quant_bits": 4})
+        param_groups.append({"params": params_embed, "quant_bits": embed_b})
     return param_groups
 
 
-def get_optim_kwargs(base_optimizer, embedding=True, quant_cls=UnifTorchaoQuantizer):
+def get_optim_kwargs(
+    model, base_optimizer, embedding=True, quant_cls=UnifTorchaoQuantizer
+):
     optim_kwargs = {}
     if embedding:
-        group_idx = len(base_optimizer.param_groups) - 2
+        embed_data_ptrs = set(
+            (
+                m.weight.data_ptr()
+                for m in model.modules()
+                if isinstance(m, nn.Embedding)
+            )
+        )
+        group_idx = -1
+        for i, group in enumerate(base_optimizer.param_groups):
+            if all(p.data_ptr() in embed_data_ptrs for p in group["params"]):
+                group_idx = i
+                break
         assert group_idx > -1
         optim_kwargs["group_quantizer_map"] = {group_idx: quant_cls()}
     return optim_kwargs
@@ -221,7 +235,7 @@ def compare_parq_convert(
     orig_model = copy.deepcopy(model)  # save copy of PARQ quantized model
 
     # equivalent to torchao's convert step
-    optimizer.torchao_convert(model, weight_only=weight_only)
+    optimizer.torchao_convert(model, weight_only=weight_only, embed_weight_only=True)
 
     inputs = model.example_inputs(device=_DEVICE)
     torch.testing.assert_close(model(inputs), orig_model(inputs))
@@ -289,13 +303,15 @@ def test_parq_train_loop(
             quantizer = TernaryUnifQuantizer() if b == 0 else UnifQuantizer()
         else:
             quantizer = LSBQuantizer()
-        param_groups = build_param_groups(model, b)
+        param_groups = build_param_groups(model, b, embed_b=b)
         base_optimizer = torch.optim.AdamW(param_groups)
 
         prox_map = (
             ProxHardQuant() if hard_prox else ProxPARQ(anneal_start=0, anneal_end=2)
         )
-        optim_kwargs = get_optim_kwargs(base_optimizer)
+        optim_kwargs = get_optim_kwargs(
+            model, base_optimizer, quant_cls=type(quantizer), embedding=False
+        )
         optimizer = QuantOptimizer(base_optimizer, quantizer, prox_map, **optim_kwargs)
         for _ in range(3):
             x = model.example_inputs(device=_DEVICE)
@@ -365,7 +381,7 @@ def test_int4_weight_only_e2e(self, group_size: int = 32):
 
         b = 4
         base_optimizer = torch.optim.AdamW(build_param_groups(model, b, group_size))
-        optim_kwargs = get_optim_kwargs(base_optimizer, embedding=False)
+        optim_kwargs = get_optim_kwargs(model, base_optimizer, embedding=False)
         optimizer = QuantOptimizer(
             base_optimizer,
             Int4UnifTorchaoQuantizer(),
@@ -387,7 +403,7 @@ def test_intx_weight_only_e2e(self, b: int = 2, group_size: int = 32):
         quantize_(m_ref, config)
 
         base_optimizer = torch.optim.AdamW(build_param_groups(model, b, group_size))
-        optim_kwargs = get_optim_kwargs(base_optimizer, embedding=False)
+        optim_kwargs = get_optim_kwargs(model, base_optimizer, embedding=False)
         optimizer = QuantOptimizer(
             base_optimizer,
             UnifTorchaoQuantizer(),
@@ -464,7 +480,7 @@ def test_intx_weight_only_e2e(self, b: int = 2, group_size: int = 32):
         quantize_(m_ref, config, filter_fn=_is_linear)
 
         base_optimizer = torch.optim.AdamW(build_param_groups(model, b, group_size))
-        optim_kwargs = get_optim_kwargs(base_optimizer, embedding=False)
+        optim_kwargs = get_optim_kwargs(model, base_optimizer, embedding=False)
         optimizer = QuantOptimizer(
             base_optimizer,
             quantizer,
@@ -486,7 +502,7 @@ def test_intx_weight_only_tied_embed_linear(
 
         quantizer = StretchedUnifTorchaoQuantizer(b)
         base_optimizer = torch.optim.SGD(build_param_groups(model, b))
-        optim_kwargs = get_optim_kwargs(base_optimizer)
+        optim_kwargs = get_optim_kwargs(model, base_optimizer)
         optimizer = QuantOptimizer(
             base_optimizer,
             quantizer,
@@ -498,7 +514,7 @@ def test_intx_weight_only_tied_embed_linear(
         optimizer.step()
 
         apply_activation_quantization(model, optimizer, model_dtype)
-        optimizer.torchao_convert(model)
+        optimizer.torchao_convert(model, embed_weight_only=True)
         check_torchao_tensor_subclass(self, model)
         self.assertTrue(
             torch.equal(model.embed_tokens.weight.qdata, model.linear2.weight.qdata)
@@ -540,7 +556,7 @@ def test_int8_dynamic_activation_intx_e2e(
 
         # quantize weights with PARQ
         base_optimizer = torch.optim.SGD(build_param_groups(model, b, group_size))
-        optim_kwargs = get_optim_kwargs(base_optimizer, embedding=False)
+        optim_kwargs = get_optim_kwargs(model, base_optimizer, embedding=False)
         optimizer = QuantOptimizer(
             base_optimizer,
             quantizer,

torchao/prototype/parq/optim/quantopt.py

@@ -160,9 +160,15 @@ def _get_quantizer(self, group_idx: int) -> Optional[Quantizer]:
             return self.group_quantizer_map[group_idx]
         return self.quantizer
 
-    def torchao_convert(self, model: nn.Module, weight_only: bool = False) -> None:
+    def torchao_convert(
+        self,
+        model: nn.Module,
+        weight_only: bool = False,
+        embed_weight_only: bool = False,
+    ) -> None:
         """Converts model parameters to torchao quantized tensor subclasses."""
         model.eval()
+        self.restore_latent_params()
 
         # TODO(lvj): find more robust way to identify embedding layers
         embed_data_ptrs = set()
@@ -175,9 +181,10 @@ def torchao_convert(self, model: nn.Module, weight_only: bool = False) -> None:
             elif _is_linear(module) and module.weight.data_ptr() not in embed_data_ptrs:
                 linear_data_ptrs.add(module.weight.data_ptr())
 
-        tied_embeddings = getattr(model, "_tied_weights_keys", None) is not None
-        if tied_embeddings:
+        tied_embeddings = False
+        if not embed_weight_only and getattr(model, "_tied_weights_keys", None):
             # Workaround for dynamic activations on tied embeddings
+            tied_embeddings = True
             for module in embed_modules:
                 setattr(module, "bias", None)