Add AWQ 4bit inference support
diff --git a/docs/awq.md b/docs/awq.md
new file mode 100644
index 0000000000..ac38ef3dd5
--- /dev/null
+++ b/docs/awq.md
@@ -0,0 +1,71 @@
+# AWQ 4bit Inference
+
+We integrated [AWQ](https://github.com/mit-han-lab/llm-awq) into FastChat to provide **efficient and accurate** 4bit LLM inference.
+
+## Install AWQ
+
+Setup environment (please refer to [this link](https://github.com/mit-han-lab/llm-awq#install) for more details):
+```bash
+conda create -n fastchat-awq python=3.10 -y
+conda activate fastchat-awq
+# cd /path/to/FastChat
+pip install --upgrade pip    # enable PEP 660 support
+pip install -e .             # install fastchat
+
+git clone https://github.com/mit-han-lab/llm-awq repositories/llm-awq
+cd repositories/llm-awq
+pip install -e .             # install awq package
+
+cd awq/kernels				
+python setup.py install	     # install awq CUDA kernels
+```
+
+## Chat with the CLI
+
+```bash
+# Download quantized model from huggingface
+# Make sure you have git-lfs installed (https://git-lfs.com)
+git lfs install
+git clone https://huggingface.co/mit-han-lab/vicuna-7b-v1.3-4bit-g128-awq
+
+# You can specify which quantized model to use by setting --awq-ckpt
+python3 -m fastchat.serve.cli \
+    --model-path models/vicuna-7b-v1.3-4bit-g128-awq \
+    --awq-wbits 4 \
+    --awq-groupsize 128 
+```
+
+## Benchmark
+
+* Through **4-bit weight quantization**, AWQ helps to run larger language models within the device memory restriction and prominently accelerates token generation. All benchmarks are done with group_size 128. 
+
+* Benchmark on NVIDIA RTX A6000:
+
+  | Model           | Bits | Max Memory (MiB) | Speed (ms/token) | AWQ Speedup |
+  | --------------- | ---- | ---------------- | ---------------- | ----------- |
+  | vicuna-7b       | 16   | 13543            | 26.06            | /           |
+  | vicuna-7b       | 4    | 5547             | 12.43            | 2.1x        |
+  | llama2-7b-chat  | 16   | 13543            | 27.14            | /           |
+  | llama2-7b-chat  | 4    | 5547             | 12.44            | 2.2x        |
+  | vicuna-13b      | 16   | 25647            | 44.91            | /           |
+  | vicuna-13b      | 4    | 9355             | 17.30            | 2.6x        |
+  | llama2-13b-chat | 16   | 25647            | 47.28            | /           |
+  | llama2-13b-chat | 4    | 9355             | 20.28            | 2.3x        |
+
+* NVIDIA RTX 4090:
+
+  | Model           | AWQ 4bit Speed (ms/token) | FP16 Speed (ms/token) | AWQ Speedup |
+  | --------------- | ------------------------- | --------------------- | ----------- |
+  | vicuna-7b       | 8.61                      | 19.09                 | 2.2x        |
+  | llama2-7b-chat  | 8.66                      | 19.97                 | 2.3x        |
+  | vicuna-13b      | 12.17                     | OOM                   | /           |
+  | llama2-13b-chat | 13.54                     | OOM                   | /           |
+
+* NVIDIA Jetson Orin:
+
+  | Model           | AWQ 4bit Speed (ms/token) | FP16 Speed (ms/token) | AWQ Speedup |
+  | --------------- | ------------------------- | --------------------- | ----------- |
+  | vicuna-7b       | 65.34                     | 93.12                 | 1.4x        |
+  | llama2-7b-chat  | 75.11                     | 104.71                | 1.4x        |
+  | vicuna-13b      | 115.40                    | OOM                   | /           |
+  | llama2-13b-chat | 136.81                    | OOM                   | /           |
diff --git a/fastchat/model/model_adapter.py b/fastchat/model/model_adapter.py
index 897a0b8aa2..b465ce68c5 100644
--- a/fastchat/model/model_adapter.py
+++ b/fastchat/model/model_adapter.py
@@ -26,6 +26,7 @@
 )
 
 from fastchat.modules.gptq import GptqConfig, load_gptq_quantized
+from fastchat.modules.awq import AWQConfig, load_awq_quantized
 from fastchat.conversation import Conversation, get_conv_template
 from fastchat.model.compression import load_compress_model
 from fastchat.model.model_chatglm import generate_stream_chatglm
@@ -150,11 +151,11 @@ def load_model(
     load_8bit: bool = False,
     cpu_offloading: bool = False,
     gptq_config: Optional[GptqConfig] = None,
+    awq_config: Optional[AWQConfig] = None,
     revision: str = "main",
     debug: bool = False,
 ):
     """Load a model from Hugging Face."""
-
     # get model adapter
     adapter = get_model_adapter(model_path)
 
@@ -219,6 +220,29 @@ def load_model(
                 torch_dtype=kwargs["torch_dtype"],
                 revision=revision,
             )
+    elif awq_config and awq_config.wbits < 16:
+        assert (
+            awq_config.wbits == 4
+        ), "Currently we only support 4-bit inference for AWQ."
+        model, tokenizer = load_awq_quantized(model_path, awq_config, device)
+        if num_gpus != 1:
+            device_map = accelerate.infer_auto_device_map(
+                model,
+                max_memory=kwargs["max_memory"],
+                no_split_module_classes=[
+                    "OPTDecoderLayer",
+                    "LlamaDecoderLayer",
+                    "BloomBlock",
+                    "MPTBlock",
+                    "DecoderLayer",
+                ],
+            )
+            model = accelerate.dispatch_model(
+                model, device_map=device_map, offload_buffers=True
+            )
+        else:
+            model.to(device)
+        return model, tokenizer
     elif gptq_config and gptq_config.wbits < 16:
         model, tokenizer = load_gptq_quantized(model_path, gptq_config)
         if num_gpus != 1:
@@ -370,6 +394,25 @@ def add_model_args(parser):
         action="store_true",
         help="Whether to apply the activation order GPTQ heuristic",
     )
+    parser.add_argument(
+        "--awq-ckpt",
+        type=str,
+        default=None,
+        help="Load quantized model. The path to the local AWQ checkpoint.",
+    )
+    parser.add_argument(
+        "--awq-wbits",
+        type=int,
+        default=16,
+        choices=[4, 16],
+        help="#bits to use for AWQ quantization",
+    )
+    parser.add_argument(
+        "--awq-groupsize",
+        type=int,
+        default=-1,
+        help="Groupsize to use for AWQ quantization; default uses full row.",
+    )
 
 
 def remove_parent_directory_name(model_path):
diff --git a/fastchat/modules/awq.py b/fastchat/modules/awq.py
new file mode 100644
index 0000000000..1f27be85c0
--- /dev/null
+++ b/fastchat/modules/awq.py
@@ -0,0 +1,85 @@
+from dataclasses import dataclass, field
+from pathlib import Path
+import sys
+
+import torch
+from transformers import AutoTokenizer, AutoConfig, AutoModelForCausalLM, modeling_utils
+
+
+@dataclass
+class AWQConfig:
+    ckpt: str = field(
+        default=None,
+        metadata={
+            "help": "Load quantized model. The path to the local AWQ checkpoint."
+        },
+    )
+    wbits: int = field(default=16, metadata={"help": "#bits to use for quantization"})
+    groupsize: int = field(
+        default=-1,
+        metadata={"help": "Groupsize to use for quantization; default uses full row."},
+    )
+
+
+def load_awq_quantized(model_name, awq_config: AWQConfig, device):
+    print("Loading AWQ quantized model...")
+
+    try:
+        from tinychat.utils import load_quant
+        from tinychat.modules import make_quant_norm, make_quant_attn, make_fused_mlp
+    except ImportError as e:
+        print(f"Error: Failed to import tinychat. {e}")
+        print("Please double check if you have successfully installed AWQ")
+        print("See https://github.com/lm-sys/FastChat/blob/main/docs/awq.md")
+        sys.exit(-1)
+
+    config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_name, use_fast=False, trust_remote_code=True
+    )
+
+    def skip(*args, **kwargs):
+        pass
+
+    torch.nn.init.kaiming_uniform_ = skip
+    torch.nn.init.kaiming_normal_ = skip
+    torch.nn.init.uniform_ = skip
+    torch.nn.init.normal_ = skip
+    modeling_utils._init_weights = False
+
+    torch.set_default_dtype(torch.half)
+    model = AutoModelForCausalLM.from_config(config, trust_remote_code=True)
+
+    if any(name in find_awq_ckpt(awq_config) for name in ["llama", "vicuna"]):
+        model = load_quant.load_awq_llama_fast(
+            model,
+            find_awq_ckpt(awq_config),
+            awq_config.wbits,
+            awq_config.groupsize,
+            device,
+        )
+        make_quant_attn(model, device)
+        make_quant_norm(model)
+        make_fused_mlp(model)
+    else:
+        model = load_quant.load_awq_model(
+            model,
+            find_awq_ckpt(awq_config),
+            awq_config.wbits,
+            awq_config.groupsize,
+            device,
+        )
+    return model, tokenizer
+
+
+def find_awq_ckpt(awq_config: AWQConfig):
+    if Path(awq_config.ckpt).is_file():
+        return awq_config.ckpt
+
+    for ext in ["*.pt", "*.safetensors"]:
+        matched_result = sorted(Path(awq_config.ckpt).glob(ext))
+        if len(matched_result) > 0:
+            return str(matched_result[-1])
+
+    print("Error: AWQ checkpoint not found")
+    sys.exit(1)
diff --git a/fastchat/serve/cli.py b/fastchat/serve/cli.py
index 67735b0d67..12b6146a40 100644
--- a/fastchat/serve/cli.py
+++ b/fastchat/serve/cli.py
@@ -25,6 +25,7 @@
 
 from fastchat.model.model_adapter import add_model_args
 from fastchat.modules.gptq import GptqConfig
+from fastchat.modules.awq import AWQConfig
 from fastchat.serve.inference import ChatIO, chat_loop
 
 
@@ -206,6 +207,11 @@ def main(args):
                 groupsize=args.gptq_groupsize,
                 act_order=args.gptq_act_order,
             ),
+            AWQConfig(
+                ckpt=args.awq_ckpt or args.model_path,
+                wbits=args.awq_wbits,
+                groupsize=args.awq_groupsize,
+            ),
             args.revision,
             args.judge_sent_end,
             args.debug,
diff --git a/fastchat/serve/inference.py b/fastchat/serve/inference.py
index ee8387bd3e..26feff2235 100644
--- a/fastchat/serve/inference.py
+++ b/fastchat/serve/inference.py
@@ -34,6 +34,7 @@
     get_generate_stream_function,
 )
 from fastchat.modules.gptq import GptqConfig
+from fastchat.modules.awq import AWQConfig
 from fastchat.utils import is_partial_stop, is_sentence_complete, get_context_length
 
 
@@ -284,6 +285,7 @@ def chat_loop(
     max_new_tokens: int,
     chatio: ChatIO,
     gptq_config: GptqConfig,
+    awq_config: AWQConfig,
     revision: str,
     judge_sent_end: bool,
     debug: bool,
@@ -298,6 +300,7 @@ def chat_loop(
         load_8bit,
         cpu_offloading,
         gptq_config,
+        awq_config,
         revision,
         debug,
     )

<!-- Thank you for your contribution! -->

<!-- Please add a reviewer to the assignee section when you create a PR. If you don't have the access to it, we will shortly find a reviewer and assign them to your PR. -->

## Why are these changes needed?

<!-- Please give a short summary of the change and the problem this solves. -->

In this PR we introduce [AWQ](https://github.com/mit-han-lab/llm-awq) 4bit inference to FastChat. AWQ supports efficient and accurate low-bit weight quantization for LLMs, and helps to run larger language models within the device memory restriction and prominently accelerates token generation. For example, when running LLaMA-2-7b, AWQ delivers 2.3x and 1.4x speedups over the FP16 baseline on RTX 4090 and Jetson Orin, respectively.

## Checks

- [x] I've run `format.sh` to lint the changes in this PR.
- [x] I've included any doc changes needed.
- [x] I've made sure the relevant tests are passing (if applicable).
