Fix (examples/generative): set weight_bit_width in weight_quant (#783)

Xilinx · Dec 13, 2023 · 52daf86 · 52daf86
1 parent 84f4225
commit 52daf86
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 10 deletions.
diff --git a/src/brevitas_examples/common/generative/quantize.py b/src/brevitas_examples/common/generative/quantize.py
@@ -193,6 +193,7 @@ def quantize_model(
     # Modify the weight quantizer based on the arguments passed in
     weight_quant = weight_quant.let(
         **{
+            'bit_width': weight_bit_width,
             'narrow_range': False,
             'block_size': weight_group_size,
             'quantize_zero_point': quantize_weight_zero_point},
@@ -311,15 +312,8 @@ def quantize_model(
                         'group_dim': 1, 'group_size': input_group_size})
 
     quant_linear_kwargs = {
-        'input_quant': linear_2d_input_quant,
-        'weight_quant': weight_quant,
-        'weight_bit_width': weight_bit_width,
-        'dtype': dtype}
-    quant_conv_kwargs = {
-        'input_quant': input_quant,
-        'weight_quant': weight_quant,
-        'weight_bit_width': weight_bit_width,
-        'dtype': dtype}
+        'input_quant': linear_2d_input_quant, 'weight_quant': weight_quant, 'dtype': dtype}
+    quant_conv_kwargs = {'input_quant': input_quant, 'weight_quant': weight_quant, 'dtype': dtype}
 
     quant_mha_kwargs = {
         'in_proj_input_quant': input_quant,

diff --git a/src/brevitas_examples/llm/main.py b/src/brevitas_examples/llm/main.py
@@ -304,7 +304,6 @@ def main():
             seqlen=args.seqlen)
         # Tie back first/last layer weights in case they got untied
         model.tie_weights()
-        print(model)
         print("Model quantization applied.")
 
     if args.act_calibration: