fix mixtral quantization scaler axis when dimension > 2 (#132)

Author: sixiang-google

convert_checkpoints.py

@@ -586,7 +586,7 @@ def main(argv) -> None:
   if FLAGS.quantize_weights:
     quantize_num_bits = 8 if "int8" in FLAGS.quantize_type else 4
     is_blockwise = "blockwise" in FLAGS.quantize_type
-    weight_axis = lambda x: 0 if x in quantize_embedding_weight_map else 1
+    weight_axis = lambda x: 0 if x in quantize_embedding_weight_map else -1
     start = time.perf_counter()
     state_dict = _quantize_state_dict(
         state_dict,