mixtral suport

main.py

@@ -181,13 +181,16 @@ def quantize_aq(model: PreTrainedModel, dataloader: Iterable, args: Namespace):
             sequential = get_sequential_groups(model)
         else:
             sequential = [list(find_sublayers(layer).keys())]
-
         for names in sequential:
             if len(args.devices) == 1:
                 assert len(inps) == len(outs) == 1  # number of per-device inputs/outputs
-                aq_handlers = init_aq_engines(layer, names, inps[0], outs[0], **forward_args)
+                aq_handlers = init_aq_engines(
+                    layer, [name for name in names if "gate" not in name], inps[0], outs[0], **forward_args
+                )
             else:
-                aq_handlers = init_aq_engines_parallel(args.devices, layer, names, inps, outs, **forward_args)
+                aq_handlers = init_aq_engines_parallel(
+                    args.devices, layer, [name for name in names if "gate" not in name], inps, outs, **forward_args
+                )
 
             for sublayer_name in aq_handlers.keys():
                 print(f"Quantizing module {sublayer_name} of layer {layer_index}")

src/modelutils.py

@@ -8,7 +8,7 @@
 
 MODEL_ERROR_MSG = "Unsupported model type {} - only 'llama', 'Yi', 'opt' and 'falcon' are supported"
 FALCON_TYPES = ("falcon", "refinedweb", "refinedwebmodel")
-LLAMA_LIKE = ("llama", "Yi", "mistral")
+LLAMA_LIKE = ("llama", "Yi", "mistral", "mixtral")
 
 
 @contextmanager
@@ -119,6 +119,7 @@ def find_sublayers(module, layers=(nn.Conv2d, nn.Linear)):
 
 def get_sequential_groups(model):
     if model.config.model_type in LLAMA_LIKE:
+        assert "mixtral" in model.config.model_type
         return [
             ["self_attn.k_proj", "self_attn.v_proj", "self_attn.q_proj"],
             ["self_attn.o_proj"],