fix tp bug (#66)

ModelTC · Sep 5, 2024 · 8b66f07 · 8b66f07
1 parent 5d8ef6f
commit 8b66f07
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/llmc/compression/quantization/base_blockwise_quantization.py b/llmc/compression/quantization/base_blockwise_quantization.py
@@ -147,13 +147,15 @@ def set_quant_config(self):
 
         self.quant_out = self.quant_config.get('quant_out', False)
         self.tp = self.quant_config.get('tp', 1)
+        self.quant_config['weight']['tp'] = self.tp
 
         # set weight quant config
         self.wquantizer = Quantizer(**self.quant_config['weight'])
 
         # set act quant config
         if 'act' in self.quant_config:
             self.w_only = False
+            self.quant_config['act']['tp'] = self.tp
             self.aquantizer = Quantizer(**self.quant_config['act'])
         else:
             self.w_only = True