Added support for AWQ models

examples/batch_generation.py

@@ -1,6 +1,6 @@
 import sys
 sys.path.append("..")
-from models.llama import LLM
+from models.llama import LLM, LLMAwq
 import argparse
 import torch
 from transformers import AutoTokenizer
@@ -14,6 +14,7 @@
 parser.add_argument('--G', type=int, default=32, help='generation length')
 parser.add_argument('--K', type=int, default=10, help='K')
 parser.add_argument('--L', type=int, default=150, help='K')
+parser.add_argument('--awq', action='store_true', help='use LLMAwq')
 args = parser.parse_args()
 print(args)
 MAX_LEN = args.M
@@ -32,7 +33,13 @@
         data = item
         break
 
-llm = LLM(K=args.K, L=args.L, max_length=MAX_LEN, model_name=args.model, batch_size=BATCH_SIZE, device=DEVICE, dtype=DTYPE)
+if args.awq:
+    print("Using LLMAwq for AWQ optimization.")
+    llm = LLMAwq(K=args.K, L=args.L, max_length=MAX_LEN, model_name=args.model, batch_size=BATCH_SIZE, device=DEVICE, dtype=DTYPE)
+else:
+    print("Using standard LLM.")
+    llm = LLM(K=args.K, L=args.L, max_length=MAX_LEN, model_name=args.model, batch_size=BATCH_SIZE, device=DEVICE, dtype=DTYPE)
+
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 text = data["input"]
 input_ids = tokenizer.encode(text=text, return_tensors="pt").to(device=DEVICE)

examples/bench.py

@@ -1,6 +1,6 @@
 import sys
 sys.path.append("..")
-from models.llama import LLM
+from models.llama import LLM, LLMAwq
 import argparse
 import torch
 from transformers import AutoTokenizer
@@ -15,6 +15,7 @@
 parser.add_argument('--G', type=int, default=128, help='generation length')
 parser.add_argument('--K', type=int, default=10, help='K')
 parser.add_argument('--L', type=int, default=150, help='L')
+parser.add_argument('--awq', action='store_true', help='use LLMAwq')
 args = parser.parse_args()
 print(args)
 MAX_LEN = args.M
@@ -33,7 +34,13 @@
         data = item
         break
 
-llm = LLM(K=args.K, L=args.L, max_length=MAX_LEN, model_name=args.model, batch_size=B, device=DEVICE, dtype=DTYPE)
+if args.awq:
+    print("Using LLMAwq for AWQ optimization.")
+    llm = LLMAwq(K=args.K, L=args.L, max_length=MAX_LEN, model_name=args.model, batch_size=B, device=DEVICE, dtype=DTYPE)
+else:
+    print("Using standard LLM.")
+    llm = LLM(K=args.K, L=args.L, max_length=MAX_LEN, model_name=args.model, batch_size=B, device=DEVICE, dtype=DTYPE)
+
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 text = data["input"]
 input_ids = tokenizer.encode(text=text, return_tensors="pt").to(device=DEVICE)

examples/generation.py

@@ -1,6 +1,6 @@
 import sys
 sys.path.append("..")
-from models.llama import LLM
+from models.llama import LLM, LLMAwq
 import argparse
 import torch
 from transformers import AutoTokenizer
@@ -16,6 +16,7 @@
 parser.add_argument('--L', type=int, default=150, help='K')
 parser.add_argument('--data', type=str, default="../data/story.txt", help='source data file')
 parser.add_argument('--template', type=str, default="meta-llama3", help='chat template')
+parser.add_argument('--awq', action='store_true', help='use LLMAwq')
 args = parser.parse_args()
 print(args)
 MAX_LEN = args.M
@@ -25,7 +26,12 @@
 DTYPE = torch.bfloat16
 DEVICE = "cuda:0"
 chat_template = Templates[args.template]
-llm = LLM(K=args.K, L=args.L, max_length=MAX_LEN, model_name=args.model, batch_size=1, device=DEVICE, dtype=DTYPE, generation_buffer=args.G + 32)
+if args.awq:
+    print("Using LLMAwq for AWQ optimization.")
+    llm = LLMAwq(K=args.K, L=args.L, max_length=MAX_LEN, model_name=args.model, batch_size=1, device=DEVICE, dtype=DTYPE, generation_buffer=args.G + 32)
+else:
+    print("Using standard LLM.")
+    llm = LLM(K=args.K, L=args.L, max_length=MAX_LEN, model_name=args.model, batch_size=1, device=DEVICE, dtype=DTYPE, generation_buffer=args.G + 32)
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 with open(args.data, "r", encoding="utf-8") as file:
     content = file.read()