refactor _replace_linear_8da4w

torchao/quantization/GPTQ.py

@@ -893,6 +893,34 @@ def _replace_linear_8da4w(
     linear_class: Type[torch.nn.Module],
     copy_weights: bool = False,
 ):
+
+    from torchao.quantization.quant_api import _replace_with_custom_fn_if_matches_filter
+
+    def filter_fn(child: torch.nn.Module, cur_fqn:str) -> bool:
+        return isinstance(child, nn.Linear) and (_check_linear_int4_k(child.in_features, groupsize) or padding_allowed)
+
+    def replacement_fn(child: torch.nn.Module) -> torch.nn.Module:
+        new_linear = linear_class(
+                    child.in_features,
+                    child.out_features,
+                    bias=False,
+                    device=child.weight.device,
+                    groupsize=groupsize,
+                    precision=precision,
+                    scales_precision=scales_precision,
+                )
+        # In distributed training, the model may be instantiated
+        # on the meta device, in which case there is no need to
+        # copy the weights, and doing so will result in an error
+        if copy_weights and child.weight.device != torch.device("meta"):
+            new_linear.weight = child.weight
+        return new_linear
+        #setattr(module, name, new_linear)
+
+    _replace_with_custom_fn_if_matches_filter(module, replacement_fn, filter_fn)
+
+
+    '''
     for name, child in module.named_children():
         if isinstance(child, nn.Linear):
             if _check_linear_int4_k(child.in_features, groupsize) or padding_allowed:
@@ -921,7 +949,7 @@ def _replace_linear_8da4w(
                 linear_class,
                 copy_weights,
             )
-
+    '''
 def replace_linear_8da4w(
     module: torch.nn.Module,
     groupsize: int,