[Bugfix] Avoid Warnings in SparseML Activation Quantization (vllm-pro…

…ject#5120)
joerunde · Jun 3, 2024 · b48cefe · b48cefe
1 parent 81de9b1
commit b48cefe
Showing 1 changed file with 20 additions and 9 deletions.
diff --git a/...or/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_statictensor.py b/...or/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_statictensor.py
@@ -89,23 +89,34 @@ def create_weights(self, layer: torch.nn.Module,
                            requires_grad=False)
 
         layer.register_parameter("weight", weight)
-        set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0})
-
-        set_weight_attrs(weight, {"weight_loader": weight_loader})
-
+        set_weight_attrs(weight, {
+            "weight_loader": weight_loader,
+            "input_dim": 1,
+            "output_dim": 0,
+        })
         layer.register_parameter("input_scale", input_scale)
-        set_weight_attrs(input_scale, {"weight_loader": weight_loader})
+        set_weight_attrs(input_scale, {
+            "weight_loader": weight_loader,
+            "ignore_warning": True,
+        })
         layer.register_parameter("input_zero_point", input_zero_point)
-        set_weight_attrs(input_zero_point, {"weight_loader": weight_loader})
+        set_weight_attrs(input_zero_point, {
+            "weight_loader": weight_loader,
+            "ignore_warning": True,
+        })
         layer.register_parameter("weight_scale", weight_scale)
-        set_weight_attrs(weight_scale, {"weight_loader": weight_loader})
         set_weight_attrs(
             weight_scale, {
+                "weight_loader": weight_loader,
                 "shard_splitter": self.scales_shard_splitter,
-                "logical_widths": output_partition_sizes
+                "logical_widths": output_partition_sizes,
+                "ignore_warning": True,
             })
         layer.register_parameter("weight_zero_point", weight_zero_point)
-        set_weight_attrs(weight_zero_point, {"weight_loader": weight_loader})
+        set_weight_attrs(weight_zero_point, {
+            "weight_loader": weight_loader,
+            "ignore_warning": True
+        })
 
     def apply_weights(self, layer: torch.nn.Module, x: torch.Tensor):
         weight = layer.weight