[Minor] modify no_split_module_classes

bef5dd2d · Jiaming Tang · 29f52eaa · bef5dd2d
Commit bef5dd2d authored Jun 25, 2023 by Jiaming Tang
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

awq/entry.py awq/entry.py +1 -1

No files found.
--- a/awq/entry.py
+++ b/awq/entry.py
@@ -80,7 +80,7 @@ def build_model_and_enc(model_path):
            model, args.load_quant, device_map="balanced",
            # TODO: can we remove this?
            no_split_module_classes=[
-                "OPTDecoderLayer", "LlamaDecoderLayer"]
+                "OPTDecoderLayer", "LlamaDecoderLayer", "BloomBlock", "MPTBlock", "DecoderLayer"]
        )
    else:  # fp16 to quantized
        kwargs = {"device_map": "balanced", "torch_dtype": torch.float16}