Update

[ghstack-poisoned]

torchao/quantization/quant_api.py

@@ -779,6 +779,7 @@ class Int4WeightOnlyConfig(AOBaseConfig):
     use_hqq: bool = False
     zero_point_domain: Optional[ZeroPointDomain] = ZeroPointDomain.NONE
 
+
 # for BC
 # TODO maybe change other callsites
 int4_weight_only = Int4WeightOnlyConfig
@@ -812,7 +813,9 @@ def _int4_weight_only_transform(
     quant_max = 15
     eps = 1e-6
     preserve_zero = LAYOUT_TO_PRESERVE_ZEROS[type(layout)]
-    zero_point_dtype = torch.bfloat16
+    zero_point_dtype = (
+        weight.dtype if isinstance(layout, Int4CPULayout) else torch.bfloat16
+    )
 
     # nonlocal zero_point_domain
     assert (