Add files using upload-large-folder tool

llm-awq/scripts/DeepSeek_R1_Distill_example.sh

@@ -0,0 +1,25 @@
+MODEL=DeepSeek-R1-Distill-Llama-8B
+
+# run AWQ search (optional; we provided the pre-computed results)
+python -m awq.entry --model_path /dataset/models/$MODEL \
+    --w_bit 4 --q_group_size 128 \
+    --run_awq --dump_awq awq_cache/$MODEL-w4-g128.pt
+
+# evaluate the AWQ quantize model (simulated pseudo quantization)
+python -m awq.entry --model_path /dataset/models/$MODEL \
+    --tasks wikitext \
+    --w_bit 4 --q_group_size 128 \
+    --load_awq awq_cache/$MODEL-w4-g128.pt \
+    --q_backend fake
+
+# generate real quantized weights (w4)
+python -m awq.entry --model_path /dataset/models/$MODEL \
+    --w_bit 4 --q_group_size 128 \
+    --load_awq awq_cache/$MODEL-w4-g128.pt \
+    --q_backend real --dump_quant quant_cache/$MODEL-w4-g128-awq.pt
+
+# load and evaluate the real quantized model (smaller gpu memory usage)
+python -m awq.entry --model_path /dataset/models/$MODEL \
+    --tasks wikitext \
+    --w_bit 4 --q_group_size 128 \
+    --load_quant quant_cache/$MODEL-w4-g128-awq.pt

llm-awq/scripts/llava_example.sh

@@ -0,0 +1,12 @@
+MODEL=llava-13b-v0
+
+# run AWQ search (optional; we provided the pre-computed results)
+python -m awq.entry --model_path /dataset/llava-hf/$MODEL \
+    --w_bit 4 --q_group_size 128 \
+    --run_awq --dump_awq awq_cache/$MODEL-w4-g128.pt
+
+# generate real quantized weights (w4)
+python -m awq.entry --model_path /dataset/llava-hf/$MODEL \
+    --w_bit 4 --q_group_size 128 \
+    --load_awq awq_cache/$MODEL-w4-g128.pt \
+    --q_backend real --dump_quant quant_cache/$MODEL-w4-g128-awq.pt

llm-awq/scripts/nvila_example.sh

@@ -0,0 +1,12 @@
+# run AWQ search
+python -m awq.entry --model_path PATH/TO/NVILA \
+    --smooth_scale --media_path https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2-VL/space_woaudio.mp4  \
+    --act_scale_path awq_cache/NVILA-VT-smooth-scale.pt --vila-20 \
+    --w_bit 4 --q_group_size 128 \
+    --run_awq --dump_awq awq_cache/NVILA.pt
+
+# generate real quantized weights (w4)
+python -m awq.entry --model_path PATH/TO/NVILA/llm \
+    --w_bit 4 --q_group_size 128  \
+    --load_awq awq_cache/NVILA.pt \
+    --q_backend real --dump_quant quant_cache/NVILA-w4-g128-awq.pt --vila-20

llm-awq/scripts/vila15_example.sh

@@ -0,0 +1,14 @@
+MODEL=VILA1.5-7b
+
+# run AWQ search (optional; we provided the pre-computed results)
+# Note: vila checkpoints are now stored in 3 parts.
+# only llm folder will be quantized
+python -m awq.entry --model_path /dataset/vila-hf/$MODEL/llm \
+    --w_bit 4 --q_group_size 128 --vila-15 \
+    --run_awq --dump_awq awq_cache/$MODEL-w4-g128.pt
+
+# generate real quantized weights (w4)
+python -m awq.entry --model_path /dataset/vila-hf/$MODEL/llm \
+    --w_bit 4 --q_group_size 128 --vila-15 \
+    --load_awq awq_cache/$MODEL-w4-g128.pt \
+    --q_backend real --dump_quant /dataset/vila-hf/$MODEL-awq/llm/$MODEL-w4-g128-awq.pt

llm-awq/scripts/vila_example.sh

@@ -0,0 +1,12 @@
+MODEL=vila-7b
+
+# run AWQ search (optional; we provided the pre-computed results)
+python -m awq.entry --model_path /dataset/vila-hf/$MODEL \
+    --w_bit 4 --q_group_size 128 \
+    --run_awq --dump_awq awq_cache/$MODEL-w4-g128.pt
+
+# generate real quantized weights (w4)
+python -m awq.entry --model_path /dataset/vila-hf/$MODEL \
+    --w_bit 4 --q_group_size 128 \
+    --load_awq awq_cache/$MODEL-w4-g128.pt \
+    --q_backend real --dump_quant quant_cache/$MODEL-w4-g128-awq.pt

llm-awq/tinychat/README.md

@@ -0,0 +1,526 @@
+# TinyChat 2.0: Efficient and Lightweight Chatbot with AWQ
+
+We introduce TinyChat, a cutting-edge chatbot interface designed for lightweight resource consumption and fast inference speed on GPU platforms. It allows for seamless deployment on consumer-level GPUs such as 3090/4090 and low-power edge devices like the NVIDIA Jetson Orin, empowering users with a responsive conversational experience like never before.
+
+The current release supports:
+
+- DeepSeek-R1-Distill-Qwen-1.5B/7B
+
+- DeepSeek-R1-Distill-Llama-8B
+
+- Llama-3-8B/70B-instruct;
+
+- NVILA-3B/8B;
+
+- VILA-1.5-3B/8B/13B/40B;
+
+- VILA-7B/13B;
+
+- LLaVA-7B/13B;
+
+- Llama-2-7B/13B-chat;
+
+- Vicuna;
+
+## Contents
+
+- [Examples](#examples)
+
+- [Benchmarks](#benchmarks)
+
+- [Usage](#usage)
+
+- [Reference](#reference)
+
+## Examples
+
+**Thanks to AWQ, TinyChat can now deliver more prompt responses through 4-bit inference. The following examples showcase that TinyChat's W4A16 generation is up to 2.7x faster on RTX 4090 and 2.9x faster on Jetson Orin, compared to the FP16 baselines. (Tested with LLaMA-3-8b model.)**
+
+
+* TinyChat with LLaMA-3-8b on RTX 4090 (2.7x faster than FP16):
+
+![TinyChat with LLaMA-3-8b on RTX 4090: W4A16 is 2.7x faster than FP16](./figures/4090_example_new.gif)
+
+* TinyChat with LLaMA-3-8b on Jetson Orin (2.9x faster than FP16):
+
+![TinyChat with LLaMA-3-8b on Jetson Orin: W4A16 is 2.9x faster than FP16](./figures/orin_example_new.gif)
+
+**TinyChat also supports inference with visual language models (e.g., VILA, LLaVA, NVILA). In the following examples, W4A16 quantized models from VILA family are launched with TinyChat.**
+
+* TinyChat with NVILA-8B on RTX 4090 (single-image inputs):
+
+![TinyChat with NVILA on 4090 single image](./figures/4090_nvila_single.gif)
+
+* TinyChat with NVILA-8B on RTX 4090 (multi-image inputs):
+
+![TinyChat with NVILA on 4090 multiple images](./figures/4090_nvila_multi.gif)
+
+* TinyChat with video reasoning:
+
+https://github.com/user-attachments/assets/b68a7a0d-5175-4030-985b-5ae0ae94f874
+
+**Prompt:** What might be the next step according to the video?
+
+**Answer:** The next step in the video could be to place the shaped dough onto a baking sheet and let it rise before baking.
+
+**Online demo:** https://vila.hanlab.ai
+
+## Speed Benchmarks
+
+We benchmark TinyChat on NVIDIA RTX 4090 (desktop GPU), Orin (edge GPU), and A100 (server-class GPU).
+
+We use the default implementation from Huggingface for the FP16 baseline. The INT4 implementation applies AWQ and utilizes our fast W4A16 GPU kernel. We also apply additional optimization techniques in the latest release. For example, we fuse all the operations in MHA/GQA/MQA into a single kernel, and fuse positional embedding kernels into the attention kernel. We also pre-allocate key-value caches to avoid the online memory allocation overhead from Huggingface. For W4A16 GEMM, we introduce FP16 accumulation when applicable for higher throughputs.
+
+
+### Decoding Speed
+
+We benchmarked the per-token generation latency for the decoding stage in the following tables.
+
+
+#### RTX 4090 Results
+
+| Model       | FP16 latency (ms) | INT4 latency (ms) | Speedup |
+| ----------- |:-----------------:|:-----------------:|:-------:|
+| LLaMA-3-8B  | 17.07             | 6.39              | 2.69x   |
+| LLaMA-2-7B  | 15.50             | 5.28              | 2.94x   |
+| LLaMA-2-13B | OOM               | 9.19              | --      |
+| Vicuna-7B   | 15.81             | 5.33              | 2.97x   |
+| VILA-7B     | 17.09             | 5.95              | 2.87x   |
+| VILA-13B    | OOM               | 10.01             | --      |
+| NVILA-2B    | 5.26              | 4.27              | 1.23x   |
+| NVILA-8B    | 16.12             | 5.97              | 2.70x   |
+
+*: For the decoding speed of language models, we follow the benchmarking setting from exLLaMA (i.e. only 4 context tokens) for the sake of simplicity and fairness. For multi-modal LMs (VILA and [NVILA](https://arxiv.org/abs/2412.04468)), we benchmark the decoding speed with single image inputs. Specifically, for NVILA, we activate the lite mode during the benchmarking, where each image is correspond to 128 input tokens.
+
+<!-- | Model       | FP16 latency (ms) | INT4 latency (ms) | Speedup |
+| ----------- |:-----------------:|:-----------------:|:-------:|
+| LLaMA-3-8B  | 17.07             | 6.66              | 2.56x   |
+| LLaMA-2-7B  | 16.17             | 6.02*             | 2.68x   |
+| LLaMA-2-13B | OOM               | 10.35             | --      |
+| Vicuna-7B   | 15.81             | 5.33              | 2.97x   |
+| Vicuna-13B  | OOM               | 9.17              | --      |
+| MPT-7B      | 17.09             | 6.18              | 2.77x   |
+| MPT-30B     | OOM               | 20.60             | --      |
+| Falcon-7B   | 29.91             | 8.02              | 3.73x   |
+| VILA-7B     | 17.09             | 5.95              | 2.87x   |
+| VILA-13B    | OOM               | 10.01             | --      | -->
+
+<!-- *: The reason why LLaMA-2-7B is slower than Vicuna-7B is because we need a longer prompt (with > 500 tokens) to prevent the model from talking with itself. If we use the benchmarking strategy from exLLaMA (i.e. only 4 context tokens), our speed is around 195 tokens / second. -->
+
+<!-- ### A6000 Results
+| Model       | FP16 latency (ms) | INT4 latency (ms) | Speedup |
+| ----------- |:-----------------:|:-----------------:|:-------:|
+| LLaMA-3-8B  | 24.95             | 10.68             | 2.34x   |         
+| LLaMA-2-7B  | 22.75             | 8.71              | 2.61x   |
+| LLaMA-2-13B | 41.72             | 14.64             | 2.85x   |
+| Vicuna-7B   | 22.03             | 8.39              | 2.63x   | 
+| Vicuna-13B  | 38.97             | 13.46             | 2.90x   |
+| MPT-7B      | 22.79             | 7.99              | 2.85x   |
+| MPT-30B     | OOM               | 28.15             | --      |  
+| Falcon-7B   | 39.44             | 11.71             | 3.37x   |
+| VILA-7B     | 23.60             | 8.14              | 2.90x   |
+| VILA-13B    | 46.58             | 13.74             | 3.39x   | -->
+
+
+#### Jetson Orin Results
+
+| Model       | FP16 latency (ms) | INT4 latency (ms) | Speedup |
+| ----------- |:-----------------:|:-----------------:|:-------:|
+| LLaMA-3-8B  | 96.00             | 32.53             | 2.95x   |
+| LLaMA-2-7B  | 83.95             | 25.94             | 3.24x   | 
+| LLaMA-2-13B | 162.33            | 47.67             | 3.41x   |
+| Vicuna-7B   | 84.77             | 26.34             | 3.22x   |
+| VILA-7B     | 86.95             | 28.09             | 3.10x   |
+| VILA-13B    | OOM               | 57.14             | --      |
+| NVILA-2B    | 24.22             | 22.25             | 1.09x   |
+| NVILA-8B    | 86.24             | 30.48             | 2.83x   |
+
+<!-- | Model       | FP16 latency (ms) | INT4 latency (ms) | Speedup |
+| ----------- |:-----------------:|:-----------------:|:-------:|
+| LLaMA-3-8B  | 96.24             | 32.55             | 2.96x   |
+| LLaMA-2-7B  | 86.80             | 32.14*            | 2.70x   | 
+| LLaMA-2-13B | OOM               | 58.20             | --      |
+| Vicuna-7B   | 84.77             | 30.73             | 2.76x   |
+| Vicuna-13B  | OOM               | 54.98             | --      |
+| MPT-7B      | 89.85             | 31.22             | 2.88x   |
+| Falcon-7B   | 147.84            | 45.10             | 3.28x   |
+| VILA-7B     | 86.95             | 28.09             | 3.10x   |
+| VILA-13B    | OOM               | 57.14             | --      |
+
+*: We can similarly achieve 33 tokens / second on Orin if we use the benchmarking strategy from exLLaMA. -->
+
+#### A100 Results
+
+| Model       | FP16 latency (ms) | INT4 latency (ms) | Speedup |
+| ----------- |:-----------------:|:-----------------:|:-------:|
+| LLaMA-3-8B  | 12.37             | 6.29              | 1.96x   |
+| LLaMA-2-7B  | 10.77             | 5.71              | 1.89x   | 
+| LLaMA-2-13B | 19.08             | 7.90              | 2.41x   |
+| Vicuna-7B   | 10.54             | 5.87              | 1.80x   |
+| VILA-7B     | 13.35             | 5.92              | 2.26x   |
+| VILA-13B    | 19.64             | 8.63              | 2.28x   |
+| NVILA-2B    | 7.03              | 5.38              | 1.31x   |
+| NVILA-8B    | 11.90             | 5.50              | 2.16x   |
+
+
+### Prefilling Speed
+
+In TinyChat 2.0, we also introduce significant prefilling speed optimizations for Large Language Models (LLMs) and Visual Language Models (VLMs). Specifically, with the integration of latest flash attention and FP16 accumulation in GEMM kernels, TinyChat now achieves state-of-the-art prefilling speed on edge devices.
+
+#### RTX 4090 Results
+
+Time-To-First-Token (TTFT) of Llama-3-8B (Unit: Seconds):
+
+| Seq Len     | 256	    | 512     | 1024	| 2048    | 3072	| 4096    |
+| ----------- |:-------:|:-------:|:-------:|:-------:|:-------:|:-------:|
+| FP16        | 0.031   | 0.055   | 0.109   | 0.211   | 0.336   | 0.446   |
+| TinyChat    | 0.021   | 0.033   | 0.064   | 0.131   | 0.200   | 0.275   |
+| Speedup     | 1.52x   | 1.68x   | 1.69x   | 1.61x   | 1.68x   | 1.62x   |
+
+
+Time-To-First-Token (TTFT) of Llama-2-7B (Unit: Seconds):
+
+| Seq Len     | 256	    | 512     | 1024	| 2048    | 3072	| 4096    |
+| ----------- |:-------:|:-------:|:-------:|:-------:|:-------:|:-------:|
+| FP16        | 0.029	| 0.058   |	0.100   | 0.211   |	0.329   | 0.441   |
+| TinyChat    | 0.018   | 0.031   | 0.060   | 0.124   | 0.193   | 0.265   |
+| Speedup     | 1.57x   | 1.83x   | 1.66x   | 1.70x   | 1.70x   | 1.66x   |				
+
+
+#### Jetson Orin Results
+
+Time-To-First-Token (TTFT) of Llama-3-8B (Unit: Seconds):
+
+| Seq Len     | 256	    | 512     | 1024	| 2048    | 3072	| 4096    |
+| ----------- |:-------:|:-------:|:-------:|:-------:|:-------:|:-------:|
+| FP16        | 0.206   | 0.399   | 0.566   | 1.519   | 2.308   | 3.114   |
+| TinyChat    | 0.166   | 0.315   | 0.623   | 1.248   | 1.907   | 2.573   |
+| Speedup     | 1.24x   | 1.26x   | 0.91x   | 1.22x   | 1.21x   | 1.21x   |
+
+
+#### Comparison with Other Systems
+
+Time-To-First-Token (TTFT) of 4-bit weight-only quantized Llama3-8B on RTX 4090 across various systems (Unit: Seconds): 
+
+
+| Seq Len             | 256   | 512   | 1024  | 2048  | 4096   |
+|:-------------------:|:-----:|:-----:|:-----:|:-----:|:------:|
+| TensorRT-LLM        | 0.027 | 0.051 | 0.100 | 0.204 | 0.421  |
+| MLC                 | 0.028 | 0.042 | 0.081 | 0.166 | 0.350  |
+| llama.cpp           | 0.026 | 0.045 | 0.086 | 0.175 | 0.375  |
+| ExLlama v2          | 0.040 | 0.051 | 0.077 | 0.139 | 0.294  |
+| TinyChat (Legacy)   | 0.031 | 0.051 | 0.101 | 0.219 | 0.461  |
+| TinyChat 2.0        | 0.021 | 0.033 | 0.065 | 0.132 | 0.278  |
+
+Our approach outperforms all existing projects, achieving state-of-the-art speed.
+
+### Context Streaming: Efficient Multi-round Dialogues
+
+In TinyChat 2.0, we introduce chunk-prefilling optimization for multi-round dialogues. For multi-turn inputs, TinyChat will reuse the KV Cache from previous conversations without recomputing them. This optimization eliminates redundant computations and significantly reduce the Time To First Token (TTFT) for subsequent interaction rounds.
+
+#### RTX 4090 Results
+
+To evaluate Context Streaming, we measure the TTFT in multi-round conversations with a fixed question length of 32 tokens, and varying history lengths from 16 to 1024 tokens. Specifically, in TinyChat 2.0, all history tokens are already prefilled to the existing KV Cache when processing the current query, while baseline systems recompute the history tokens for each query.
+
+<!-- To demonstrate the effectiveness of Context Streaming, we measure the TTFT in multi-round conversation with a fixed question length of 32 and varying history lengths ranging from 16 to 1024 tokens. This setup means that a number of history tokens (based on the specified history length) are already input into the model. In this round, the question tokens (32 tokens) are also input, and the model takes TTFT to process these question tokens, prefill the KV cache, and generate the first token. All the tables below follows this setting. The speedup ratio in all the tables below refer to the acceleration achieved by the new method compared to FP16 inference. -->
+
+Time-To-First-Token (TTFT) of Llama-3-8B (Unit: ms):
+
+| History length            | 16    | 32    | 64    | 128   | 256   | 512   | 1024   |
+|---------------------------|:-----:|:-----:|:-----:|:-----:|:-----:|:-----:|:------:|
+| FP16                      | 21.49 | 21.38 | 23.51 | 40.82 | 47.15 | 75.41 | 162.27 | 
+| TinyChat (Legacy)         | 15.20 | 14.89 | 17.61 | 29.66 | 44.11 | 72.50 | 163.90 |
+| TinyChat 2.0              | 14.30 | 14.05 | 14.05 | 14.43 | 14.38 | 14.35 | 14.49  |
+| Speedup                   | 1.54x | 1.54x | 1.69x | 2.84x | 3.33x | 5.27x | 11.45x |
+
+<!-- Time-To-First-Token (TTFT) of VILA-1.5-8B (Unit: ms):
+
+| History length            | 16    | 32    | 64    | 128   | 256   | 512    | 1024   |
+|---------------------------|:-----:|:-----:|:-----:|:-----:|:-----:|:------:|:------:|
+| FP16 TTFT (ms)            | 22.20 | 22.00 | 24.17 | 41.85 | 62.97 | 101.84 | 217.57 | 
+| Legacy TinyChat TTFT (ms) | 16.14 | 15.98 | 18.28 | 30.72 | 59.67 | 98.52  | 219.19 |
+| New TinyChat TTFT (ms)    | 14.86 | 14.69 | 14.64 | 14.90 | 14.91 | 14.95  | 14.90  |
+| New TinyChat Speedup      | 1.49x | 1.50x | 1.65x | 2.81x | 4.22x | 6.81x  | 14.60x |
+
+NOTE: [TODO] @Yuming. The current setting is too complicated. Let's consider the case: each round, there is an image input. Baseline need to re-encode every image, while tinychat only need to encode 1 image.
+
+*: For Visual Language Models, the speedup of Context Streaming is more significant, since the model only decodes images during the first round. In the experiment, We assume that approximately 75% of the history tokens represent images, leading to the number of images in the table being 0, 0, 0, 0, 1, 2, 4. This assumption is reasonable to some extent, considering that a single image is decoded into 196 tokens. -->
+
+
+
+<!-- We have optimized the speed of the context stage and updated our code with several enhancements, including the adoption of FlashAttention and the elimination of redundant computations. The key optimizations include:
+1. Adopting the FlashAttention kernel. (Currently we only support single-batch operations to achieve better results)
+2. Computing only the last tokens in the final logits layer. (This method is used by default.)
+3. Utilizing history KV caches in the context stage to speed up. (chunk prefilling) -->
+
+<!-- These optimizations are orthogonal, enabling their combined application to achieve significant speedups. Under specific conditions, these enhancements can lead to up to an 14x speedup on 4090 GPUs and an 8x speedup on Orin GPUs in Time To First Token (TTFT) compared to the previous version of TinyChat and FP16. We conducted experiments using both Orin and 4090 GPUs, and detailed results are presented below.  -->
+
+
+<!-- ### Orin Results
+We follow the setup above and the results are as below.
+#### Llama-3-8B
+| History length            | 16     | 32     | 64     | 128    | 256    | 512    | 1024    |
+|---------------------------|:------:|:------:|:------:|:------:|:------:|:------:|:-------:|
+| FP16 TTFT (ms)            | 107.10 | 108.81 | 114.07 | 224.78 | 343.95 | 582.54 | 1048.11 |
+| Legacy TinyChat TTFT (ms) | 92.04  | 111.31 | 106.60 | 160.78 | 278.47 | 528.70 | 1145.35 |
+| New TinyChat TTFT (ms)    | 65.57  | 65.40  | 66.49  | 67.15  | 73.29  | 84.67  | 118.53  |
+| New TinyChat Speedup      | 1.52x  | 1.65x  | 1.70x  | 3.30x  | 4.51x  | 6.75x  | 8.65x   |  -->
+
+
+## Accuracy Evaluation
+
+
+AWQ also achieves decent performance on the Visual Language Models. We evaluate AWQ on VILA and the lastest NVILA models.
+
+| NVILA-8B   | AI2D       | ChartQA    | DocVQA     | MMMU_val   | SEED       | TextVQA    | VideoMME   | 
+| ---------- |:----------:|:----------:|:----------:|:----------:|:----------:|:----------:|:----------:|
+| FP16       | 91.0       | 84.8       | 91.7       | 50.7       | 76.3       | 78.1       | 63.9       |
+| AWQ-INT4   | 90.9       | 83.3       | 89.2       | 49.3       | 76.2       | 78.2       | 62.1       |
+
+<!--
+| NVILA-8B   | AI2D       | ChartQA    | DocVQA     | MMMU_val   | SEED       | TextVQA    | VideoMME-Short | VideoMME-Medium |  VideoMME-Long | VideoMME-Overall | 
+| ---------- |:----------:|:----------:|:----------:|:----------:|:----------:|:----------:|:----------:|:----------:|:----------:|:----------:|
+| FP16       | 91.0       | 84.8       | 91.7       | 50.7       | 76.3       | 78.1       | 74.9       | 62.1       | 54.7       | 63.9       |
+| AWQ-INT4   | 90.9       | 83.3       | 89.2       | 49.3       | 76.2       | 78.2       | 73.2       | 61.3       | 51.6       | 62.1       | -->
+
+
+
+| VILA-1.5-3B   | VQA-v2            | GQA               | VizWiz  | ScienceQA         | TextVQA           | POPE    | MME     | MMBench           | MMBench-CN    | SEED    |
+| ----------- |:-----------------:|:-----------------:|:-------:|:-----------------:|:-----------------:|:-------:|:-------:|:-----------------:|:-------------:|:-------:|
+| FP16        | 80.4  | 61.5 | 53.5   | 69.0  | 60.4  | 85.9 | 1442.4 | 63.4 | 52.7   | 60.9 |
+| AWQ-INT4    | 80.0  | 61.1 | 53.8   | 67.8  | 60.4  | 85.9 | 1437.3 | 63.3 | 51.4   | 59.8 | 
+
+| VILA-1.5-8B    | VQA-v2            | GQA               | VizWiz  | ScienceQA         | TextVQA           | POPE    | MME     | MMBench           | MMBench-CN    | SEED    |
+| ----------- |:-----------------:|:-----------------:|:-------:|:-----------------:|:-----------------:|:-------:|:-------:|:-----------------:|:-------------:|:-------:|
+| FP16        | 80.9  | 61.9 | 58.7   | 79.9  | 66.3  | 84.4 | 1577.01 | 72.3 | 66.2   | 64.2 |
+| AWQ-INT4    | 80.3  | 61.7 | 59.3   | 79.0  | 65.4  | 82.9 | 1593.65 | 71.0 | 64.9   | 64.0 |
+
+| VILA-1.5-13B    | VQA-v2            | GQA               | VizWiz  | ScienceQA         | TextVQA           | POPE    | MME     | MMBench           | MMBench-CN    | SEED    |
+| ----------- |:-----------------:|:-----------------:|:-------:|:-----------------:|:-----------------:|:-------:|:-------:|:-----------------:|:-------------:|:-------:|
+| FP16       | 82.8  | 64.3 | 62.6   | 80.1  | 65.0  | 86.3 | 1569.55 | 74.9 | 66.3   | 65.1 |
+| AWQ-INT4    | 82.7  | 64.5 | 63.3   | 79.7  | 64.7  | 86.7 | 1531.35 | 74.7 | 66.7   | 65.1 |
+
+
+| VILA-1.5-40B    | VQA-v2            | GQA               | VizWiz  | ScienceQA         | TextVQA           | POPE    | MME     | MMBench           | MMBench-CN    | SEED    |
+| ----------- |:-----------------:|:-----------------:|:-------:|:-----------------:|:-----------------:|:-------:|:-------:|:-----------------:|:-------------:|:-------:|
+| FP16      | 84.3  | 64.6 | 62.2   | 87.2  | 73.6  | 87.3 | 1726.82 | 82.4 | 80.2   | 69.1 |
+| AWQ-INT4   | 84.1  | 64.4 | 61.3   | 86.7  | 73.2  | 88.2 | 1714.79 | 83.2 | 79.6   | 68.9 | 
+
+AWQ has also demonstrated impressive performance on inference benchmarks, maintaining strong accuracy across a range of reasoning tasks.
+
+| DeepSeek-R1-Distill-Llama-8B | WikiText perplexity | AIME 2024 | Math-500 | 
+| ---------------------------- |:-------------------:|:-------------------:|:-------------------:|
+| FP16       | 13.13 | 43.33% | 83.00% |
+| AWQ-INT4   | 13.84 | 43.33% | 84.40% |
+
+| DeepSeek-R1-Distill-Qwen-7B | WikiText perplexity | AIME 2024 | Math-500 | 
+| --------------------------- |:-------------------:|:-------------------:|:-------------------:|
+| FP16       | 25.06 | 53.33% | 91.40% |
+| AWQ-INT4   | 27.45 | 53.33% | 89.60% |
+
+## Usage
+
+1. Please follow the [AWQ installation guidance](https://github.com/mit-han-lab/llm-awq#readme) to install AWQ and its dependencies. If you want to use FlashAttention, start by installing it with: ```pip install flash-attn --no-build-isolation```. However, for some GPUs such as Jetson Orin, there is no pre-built version available. You will need to build it from source. Follow these commands:
+```bash
+git clone https://github.com/Dao-AILab/flash-attention.git
+cd flash-attention
+sed -i '168 a\    cc_flag.append("-gencode")\n\    cc_flag.append("arch=compute_87,code=sm_87")' setup.py
+python setup.py install
+``` 
+This process may take some time as it involves compiling the code. Additionally, please note that these commands are just for Jetson Orin GPUs, whose CUDA compute capability is 87. For other GPUs, you may use ```nvidia-smi --query-gpu=compute_cap --format=csv``` to get the compute capability and merely change '87' to that. 
+
+2. Download the pretrained instruction-tuned LLMs:
+   
+   - For LLaMA-2-chat, please refer to [this link](https://huggingface.co/meta-llama/Llama-2-7b-chat-hf);
+   
+   - For Vicuna, please refer to [this link](https://huggingface.co/lmsys/);
+   
+   - For MPT-chat, please refer to [this link](https://huggingface.co/mosaicml/mpt-7b-chat);
+   
+   - For Falcon-instruct, please refer to [this link](https://huggingface.co/tiiuae/falcon-7b-instruct).
+
+3. Quantize instruction-tuned LLMs with AWQ:
+- We provide pre-computed AWQ search results for multiple model families, including LLaMA, OPT, Vicuna, VILA, and LLaVA. To get the pre-computed AWQ search results, run:
+
+```bash
+# git lfs install  # install git lfs if not already
+git clone https://huggingface.co/datasets/mit-han-lab/awq-model-zoo awq_cache
+```
+
+- You may run a one-line starter below:
+
+```bash
+./scripts/llama2_demo.sh
+```
+
+Alternatively, you may go through the process step by step. We will demonstrate the quantization process with LLaMA-2. For all other models except Falcon, one only needs to change the `model_path` and saving locations. For Falcon-7B, we also need to change `q_group_size` from 128 to 64.
+
+- Perform AWQ search and save search results (we already did it for you):
+
+```bash
+mkdir awq_cache
+python -m awq.entry --model_path /PATH/TO/LLAMA2/llama-2-7b-chat \
+    --w_bit 4 --q_group_size 128 \
+    --run_awq --dump_awq awq_cache/llama-2-7b-chat-w4-g128.pt
+```
+
+- Generate real quantized weights (INT4):
+
+```bash
+mkdir quant_cache
+python -m awq.entry --model_path /PATH/TO/LLAMA2/llama-2-7b-chat \
+    --w_bit 4 --q_group_size 128 \
+    --load_awq awq_cache/llama-2-7b-chat-w4-g128.pt \
+    --q_backend real --dump_quant quant_cache/llama-2-7b-chat-w4-g128-awq.pt
+```
+
+4. Run the TinyChat demo:
+
+```bash
+cd tinychat
+python demo.py --model_type llama \
+    --model_path /PATH/TO/LLAMA2/llama-2-7b-chat \
+    --q_group_size 128 --load_quant quant_cache/llama-2-7b-chat-w4-g128-awq.pt \ 
+    --precision W4A16
+```
+
+Note: if you use Falcon-7B-instruct, please remember to also change `q_group_size` to 64. You may also run the following command to execute the chatbot in FP16 to compare the speed and quality of language generation:
+
+```bash
+python demo.py --model_type llama \
+    --model_path /PATH/TO/LLAMA2/llama-2-7b-chat \
+    --precision W16A16
+```
+You can now try using FlashAttention along with chunk prefilling. Use the following two arguments when running demo: ```
+--flash --chunk_prefilling```. 
+
+The above command works well for most cloud and desktop GPUs, since their CPU and GPU memory space are separated. However, for edge GPUs with shared host and device memory, in order to run larger models (e.g. LLaMA-2-70B on 64GB Orin), it is necessary to break down the pretrained checkpoints into small pieces:
+
+```bash
+python split_ckpt.py --input_path quant_cache/llama-2-7b-chat-w4-g128-awq.pt \
+    --output_path quant_cache/llama-2-7b-chat-w4-g128-awq
+```
+
+Then, to run the demo, one can use the following command. The only changes compared with the demo command above are: 
+
+- We modify the `load_quant` argument;
+
+- We introduce another flag `mem_efficient_load`.
+
+```bash
+cd tinychat
+python demo.py --model_type llama \
+    --model_path /PATH/TO/LLAMA2/llama-2-7b-chat \
+    --q_group_size 128 --load_quant quant_cache/llama-2-7b-chat-w4-g128-awq \ 
+    --precision W4A16 --mem_efficient_load
+```
+
+5. (Optional) Run the benchmark script to get TTFT and decoding throughput:
+
+```bash
+cd tinychat
+python benchmark.py --flash \
+    --context_length 16 32 64 128 256 512 1024 2048 \
+    --model_path /PATH/TO/LLAMA2/llama-2-7b-chat --precision W4A16
+```
+To benchmark chunk prefilling, use:
+```bash
+python benchmark.py --chunk_prefilling \
+    --model_path /PATH/TO/LLAMA2/llama-2-7b-chat \
+    --question_length 32 --context_length 16 32 64 128 256 512 1024 --precision W4A16
+```
+Note: The kv caches in the current implementation are pre-allocated. So if you run out of memory, it might be the case that the kv cache is too large. To solve the problem, you may pass in `--max_seq_len [a smaller number]`.
+### Support Visual Language Models (VILA-1.5, VILA, LLaVA, NVILA)
+
+Our TinyChat also supports visual language models. Follow the instructions below to run VLMs on your own devices!
+
+Step 1-3 are same as the deployment for Language-only models.
+
+1. Follow the [AWQ installation guidance](https://github.com/mit-han-lab/llm-awq#readme) to install AWQ and its dependencies.
+
+2. Download the pretrained VLMs (VILA).
+
+3. Quantize the VLMs with AWQ and get the quantized checkpoint in `quant_cache`. We also provide a [sample script](../scripts/vila_example.sh) for this step.
+
+4. Run the TinyChat demo for VLMs (with vila15_demo.py for VILA-1.5, vila10_demo.py for VILA and LLaVA):
+
+```bash
+cd tinychat
+python vila15_demo.py \
+    --model-path /PATH/TO/VILA/VILA-1.5-13B \
+    --quant-path quant_cache/vila-1.5-13b-w4-g128-awq.pt \ 
+    --precision W4A16 \
+    --image-file /PATH/TO/INPUT/IMAGE \
+    --vis-image #Optional
+```
+
+Alternatively, one may also skip the quantization process and directy download the quantized VILA-1.5 checkpoints from [here](https://huggingface.co/Efficient-Large-Model). Take VILA-1.5-13B as an example, after running:
+
+```bash
+cd tinychat
+git clone https://huggingface.co/Efficient-Large-Model/VILA1.5-13b-AWQ
+```
+
+One may run:
+```bash
+python vila15_demo.py \
+    --model-path VILA1.5-13b-AWQ \
+    --quant-path VILA1.5-13b-AWQ/llm \ 
+    --precision W4A16 \
+    --image-file /PATH/TO/INPUT/IMAGE \
+    --vis-image #Optional
+```
+
+to run the terminal demo directly. You can also use``` --flash --chunk_prefilling``` to accelerate. We also support context stage benckmarking for VILA. 
+```bash
+python benchmark_context.py --flash --chunk_prefilling     \
+    --model_path PATH/TO/Llama-3-VILA1.5-8B     \
+    --question_length 32 --context_length 16 32 64 128 256 512 1024 \
+    --model_type vila --quant
+```
+Note: if you enable `--vis-image` mode, TinyChat will print input images directly in your terminal. You may need to install [termvisage](https://github.com/AnonymouX47/termvisage) to enable this mode. A [terminal emulator](https://github.com/AnonymouX47/termvisage?tab=readme-ov-file#requirements) is also required.
+
+Note: VILA model family supports multi-image inputs. You can input multiple images in `/PATH/TO/INPUT/IMAGE` above, each image should be seperated by `,`.
+
+5. TinyChat support NVILA now! We adopt W8A8 SmoothQuant for VisionTower and W4A16 quantization for LLM, achieving 1.3x -3.3x speedup for prefiling satge and nearly 1.5x higher throughput. You can use the commands below to prepare the your model and try four basic tasks of NVILA-video model.
+To prepared the needed act scale for awq and smoothquant, please run:
+```bash
+python -m awq.entry --model_path PATH/TO/NVILA \
+    --smooth_scale --media_path https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2-VL/space_woaudio.mp4  \
+    --act_scale_path awq_cache/NVILA-VT-smooth-scale.pt --vila-20 \
+    --w_bit 4 --q_group_size 128 \
+    --run_awq --dump_awq awq_cache/NVILA.pt
+```
+Then, please generate real quantized LLM with:
+```bash
+python -m awq.entry --model_path PATH/TO/NVILA/llm \
+    --w_bit 4 --q_group_size 128  \
+    --load_awq awq_cache/NVILA.pt \
+    --q_backend real --dump_quant quant_cache/NVILA-w4-g128-awq.pt --vila-20
+```
+Next, try chatting with it using the command below to experience shorter Time To First Token (TTFT) and higher decoding throughput.
+```bash
+python nvila_demo.py --model-path EPATH/TO/NVILA       \
+    --quant_path PATH/TO/NVILA-w4-g128-v2.pt      \
+    --media PATH/TO/MEDIA    \
+    --act_scale_path PATH/TO/NVILA-smooth-scale.pt \
+    --quant_llm --chunk --model_type nvila
+```
+
+
+## Team
+
+TinyChat is developed by the following wonderful team:
+
+- [Shang Yang](https://ys-2020.github.io/): Project Lead, TinyChat v1 and v2 Lead;
+- [Haotian Tang](http://kentang.net): Project Lead, TinyChat v1 Lead, v2 Mentor;
+- [Yuming Lou](<>): TinyChat v2 Lead;
+- [Junxian Guo](<>): TinyChat v2 Contributor;
+- [Song Han](https://hanlab.mit.edu/songhan): Project Advisor.
+
+Credits also go to AWQ algorithm leads: [Ji Lin](https://www.linji.me/) and [Jiaming Tang](https://jiamingtang.me/).
+
+## Reference
+
+TinyChat is inspired by the following open-source projects: [FasterTransformer](https://github.com/NVIDIA/FasterTransformer), [FlashAttention](https://github.com/Dao-AILab/flash-attention), [vLLM](https://github.com/vllm-project/vllm), [FastChat](https://github.com/lm-sys/FastChat), [llama_cu_awq](https://github.com/ankan-ban/llama_cu_awq), [LLaVA](https://github.com/haotian-liu/LLaVA), [termvisage](https://github.com/AnonymouX47/termvisage).
+

llm-awq/tinychat/internvl_demo.py

@@ -0,0 +1,270 @@
+import argparse
+
+from termcolor import colored
+
+import llava
+from llava.media import Image, Video
+import torch
+from awq.quantize import fake_quant
+from transformers import AutoConfig, AutoTokenizer
+from tinychat.utils.load_quant import load_awq_model
+from tinychat.utils.llava_image_processing import (
+    load_images,
+    vis_images,
+)
+
+
+def skip(*args, **kwargs):
+    pass
+
+
+from tinychat.utils.tune import (
+    device_warmup,
+    tune_all_wqlinears,
+    tune_llava_patch_embedding,
+)
+from tinychat.utils.prompt_templates import (
+    get_prompter,
+    get_stop_token_ids,
+    get_image_token,
+)
+from llava.utils.media import extract_media
+import tinychat.utils.constants
+from tinychat.stream_generators.internvl_stream_gen import InternVLStreamGenerator
+from tinychat.utils.conversation_utils import gen_params, stream_output, TimeStats
+
+import os
+
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+
+def tune_intern_patch_embedding(vision_model, device):
+    patch_embedding = vision_model.embeddings.patch_embedding
+    patch_embedding = patch_embedding.to(device)
+    
+    image = (
+        torch.randn((1, patch_embedding.in_channels, 336, 336))
+        .to(device)
+        .to(patch_embedding.weight.dtype)
+    )
+    for i in range(100):
+        patch_embedding(image)
+
+
+def main(args):
+    # Accelerate model initialization
+    setattr(torch.nn.Linear, "reset_parameters", lambda self: None)
+    setattr(torch.nn.LayerNorm, "reset_parameters", lambda self: None)
+    torch.nn.init.kaiming_uniform_ = skip
+    torch.nn.init.kaiming_normal_ = skip
+    torch.nn.init.uniform_ = skip
+    torch.nn.init.normal_ = skip
+    tinychat.utils.constants.max_seq_len = args.max_seq_len
+
+    # Prepare model
+    from tinychat.models import InternVL3
+    from tinychat.models.internvl.internvit import InternVisionModel
+    from transformers import AutoConfig
+    from tinychat.models.qwen2 import Qwen2ForCausalLM
+
+    config = AutoConfig.from_pretrained(args.model_path, trust_remote_code=True)
+    config.resume_path = args.model_path
+    if args.quant_llm or args.all:
+        model = InternVL3.from_pretrained(args.model_path, config=config).half()
+    else:
+        llm = Qwen2ForCausalLM.from_pretrained(args.model_path)
+        llm = llm.cpu()
+        tokenizer = AutoTokenizer.from_pretrained(
+            args.model_path, use_fast=False, trust_remote_code=True
+        )
+        llm.resize_token_embeddings(len(tokenizer))
+        model = InternVL3(config, language_model=llm).half()
+
+    if args.smooth_VT or args.all:
+        from awq.quantize import smooth_lm
+
+        act_scales = torch.load(args.act_scale_path)
+        smooth_lm(model.vision_tower, act_scales, 0.3)
+    if args.quant_llm or args.all:
+        from tinychat.modules import (
+            make_quant_norm,
+            make_quant_attn,
+            make_fused_mlp,
+            make_fused_vision_attn,
+        )
+
+        model = load_awq_model(model, args.quant_path, 4, 128, args.device)
+        make_quant_attn(model, args.device, True)
+        make_quant_norm(model)
+        model.cpu()
+        model.resize_token_embeddings(len(model.tokenizer))
+        pass
+
+    if args.quant_VT or args.all:
+        from tinychat.modules import QuantInternVisionEncoder
+        model.vision_model.encoder = QuantInternVisionEncoder(model.vision_model.encoder)
+        # model.vision_model.encoder = torch.compile(model.vision_model.encoder)
+    
+    model = model.cuda().eval()
+    device_warmup(args.device)
+    # tune_intern_patch_embedding(model.vision_model, device=args.device)
+
+    # Pre-prepare media
+    prompt = []
+    media_files = []
+    if args.media is not None:
+        for media in args.media or []:
+            if any(media.endswith(ext) for ext in [".jpg", ".jpeg", ".png"]):
+                media = Image(media)
+                media_files.append(media)
+                media_prompt = "<image>"
+            elif any(media.endswith(ext) for ext in [".mp4", ".mkv", ".webm"]):
+                media = Video(media)
+                media_files.append(media)
+                media_prompt = "<vila/video>"
+            else:
+                raise ValueError(f"Unsupported media type: {media}")
+            prompt.append(media)
+    media_num = len(media_files)
+    if args.vis_image:
+        print("=" * 50)
+        print("Input Image:")
+        vis_images(args.media)
+    
+    conversation = [{"from": "human", "value": prompt}]
+    media, media_cfg = model.prepare_media(conversation)
+    # Prepare streaming
+    stream_generator = InternVLStreamGenerator
+    # Prepare prompt
+    if args.max_seq_len <= 1024:
+        short_prompt = True
+    else:
+        short_prompt = False
+    model_prompter = get_prompter(
+        args.model_type, args.model_path, short_prompt, args.empty_prompt
+    )
+    stop_token_ids = get_stop_token_ids(args.model_type, args.model_path)
+    count = 0
+
+    if args.empty_prompt:
+        input_indicator = "Input: "
+        output_indicator = "Generated: "
+    else:
+        input_indicator = "USER: "
+        output_indicator = "ASSISTANT: "
+
+    count = 0
+    model.eval()
+    time_stats = TimeStats()
+    start_pos = 0
+    while True:
+        # Get input from the user
+        print("=" * 50)
+        input_prompt = input(input_indicator)
+        print("-" * 50)
+        if input_prompt == "":
+            print("EXIT...")
+            time_stats.show()
+            break
+        if count == 0:  # Insert media here
+            if args.media is not None:
+                if media_prompt in input_prompt:
+                    input_prompt = input_prompt
+                else:
+                    if media_prompt == "<image>":
+                        input_prompt = media_prompt * media_num + input_prompt
+                    elif media_prompt == "<vila/video>":
+                        video_prefix = ''.join([f'Frame{i+1}: <image>\n' for i in range(len(media_cfg))])
+                        input_prompt = video_prefix + input_prompt
+            
+            model_prompter.insert_prompt(input_prompt)
+        else:
+            model_prompter.insert_prompt(input_prompt)
+            if args.chunk_prefilling:
+                media = None
+                media_cfg = None
+        output_stream = stream_generator(
+            model,
+            gen_params,
+            model_prompter.model_input,
+            media,
+            media_cfg,
+            start_pos,
+            device=args.device,
+            stop_token_ids=stop_token_ids,
+            chunk_prefilling=args.chunk_prefilling,
+            quant_llm=args.quant_llm or args.all,
+        )
+        print(output_indicator, end="", flush=True)
+        if count == 0:
+            outputs, total_tokens = stream_output(output_stream, time_stats)
+        else:
+            outputs, total_tokens = stream_output(output_stream)
+        if args.chunk_prefilling:
+            start_pos += total_tokens
+        if (
+            args.single_round is not True and args.max_seq_len > 512
+        ):  # Only memorize previous conversations when kv_cache_size > 512
+            model_prompter.update_template(outputs, args.chunk_prefilling)
+        count += 1
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--model_type", type=str, default="LLaMa", help="type of the model"
+    )
+    parser.add_argument(
+        "--model-path", type=str, default="/data/llm/checkpoints/llava/llava-v1.5-7b"
+    )
+    parser.add_argument(
+        "--quant_path",
+        type=str,
+        default="/data/llm/checkpoints/llava/llava-v1.5-7b-w4-g128-awq.pt",
+    )
+    parser.add_argument(
+        "--act_scale_path",
+        type=str,
+        default="/PATH/TO/SCALE",
+    )
+    parser.add_argument(
+        "--media", type=str, nargs="+", help="Multi-modal input (Video or image path)"
+    )
+    parser.add_argument("--device", type=str, default="cuda")
+    parser.add_argument("--max_seq_len", type=int, default=4098)
+    parser.add_argument(
+        "--single_round",
+        action="store_true",
+        help="whether to memorize previous conversations",
+    )
+    parser.add_argument(
+        "--vis-image",
+        action="store_true",
+        help="whether to visualize the image while chatting",
+    )
+    parser.add_argument(
+        "--empty-prompt",
+        action="store_true",
+        help="whether to use empty prompt template",
+    )
+    parser.add_argument(
+        "--flash_attn",
+        action="store_true",
+        help="whether to use flash attention",
+    )
+    parser.add_argument(
+        "--chunk_prefilling",
+        action="store_true",
+        help="If used, in context stage, the history tokens will not be recalculated, greatly speeding up the calculation",
+    )
+    # smooth and quantization options
+    parser.add_argument("--quant_llm", action="store_true")
+    parser.add_argument("--quant_VT", action="store_true")
+    parser.add_argument("--smooth_VT", action="store_true")
+    parser.add_argument("--all", action="store_true")
+    parser.add_argument(
+        "--fakequant_VT",
+        action="store_true",
+        help="Use fake quant or real quant for VisionTower",
+    )
+    args = parser.parse_args()
+    main(args)

llm-awq/tinychat/models/__init__.py

@@ -0,0 +1,10 @@
+from .falcon import FalconForCausalLM
+from .llama import LlamaForCausalLM
+from .mpt import MPTForCausalLM
+from .llava_llama import LlavaLlamaForCausalLM
+from .qwen2 import Qwen2ForCausalLM
+try:
+    from .internvl3 import InternVL3
+except ImportError as e:
+    print("InternVL3 model import failure. To activate, please install VILA at https://github.com/NVlabs/VILA.")
+

llm-awq/tinychat/models/falcon.py

@@ -0,0 +1,304 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# This software may be used and distributed according to the terms of the GNU General Public License version 3.
+
+from typing import Optional, Tuple
+from dataclasses import dataclass
+import math
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+import awq_inference_engine
+
+import tinychat.utils.constants
+
+max_batch_size = tinychat.utils.constants.max_batch_size
+max_seq_len = tinychat.utils.constants.max_seq_len
+
+
+# rotary pos emb helpers (torch.jit.script does not seem to support staticmethod...)
+def rotate_half(x):
+    x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :]
+    return torch.cat(
+        (-x2, x1), dim=x1.ndim - 1
+    )  # dim=-1 triggers a bug in torch < 1.8.0
+
+
+class RotaryEmbedding(nn.Module):
+    """Implementation of RotaryEmbedding from GPT-NeoX.
+    This implementation is design to operate on queries and keys that are compatible with
+    [batch_size, n_heads_per_partition, seq_len, head_dim] (e.g. MinGPTAttention format).
+    """
+
+    def __init__(
+        self,
+        head_dim: int,
+        base=10000,
+    ):
+        super().__init__()
+        inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.head_dim = head_dim
+        self.seq_len_cached = None
+        self.batch_size_cached = None
+        self.cos_cached: torch.Tensor | None = None
+        self.sin_cached: torch.Tensor | None = None
+
+    def cos_sin(
+        self,
+        seq_len: int,
+        device="cuda",
+        dtype=torch.bfloat16,
+    ) -> torch.Tensor:
+        if seq_len != self.seq_len_cached:
+            self.seq_len_cached = seq_len
+            t = torch.arange(seq_len, device=device).type_as(self.inv_freq)
+            freqs = torch.einsum("i,j->ij", t, self.inv_freq)
+            emb = torch.cat((freqs, freqs), dim=-1).to(device)
+
+            if dtype in [torch.float16, torch.bfloat16]:
+                emb = emb.float()
+
+            self.cos_cached = emb.cos()[None, :, :]
+            self.sin_cached = emb.sin()[None, :, :]
+
+            self.cos_cached = self.cos_cached.type(dtype)
+            self.sin_cached = self.sin_cached.type(dtype)
+
+        return self.cos_cached, self.sin_cached
+
+    def forward(self, _q, _k):
+        batch, seq_len, num_heads, head_dim = _q.shape
+        q = _q.permute(0, 2, 1, 3).contiguous().reshape(-1, seq_len, head_dim)
+        k = _k.permute(0, 2, 1, 3).contiguous().reshape(-1, seq_len, head_dim)
+        cos, sin = self.cos_sin(seq_len, q.device, q.dtype)
+        return (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) * sin)
+
+
+class FalconAttentionFused(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+        self.args = args
+        self.n_local_heads = args.n_head
+        self.head_dim = args.hidden_size // args.n_head
+
+        self.query_key_value = nn.Linear(
+            args.hidden_size,
+            args.n_head * self.head_dim + 2 * self.head_dim,
+            bias=False,
+        )
+
+        self.dense = nn.Linear(
+            args.n_head * self.head_dim,
+            args.hidden_size,
+            bias=False,
+        )
+
+        # following fastertransformer definition
+
+        self.cache_v = (
+            torch.zeros(
+                (
+                    max_batch_size,
+                    1,
+                    max_seq_len,
+                    self.head_dim,
+                )
+            )
+            .cuda()
+            .half()
+        )  # added to half
+        # 8: pack 8 fp16 in FT, if fp32 then use 4
+        self.cache_k = (
+            torch.zeros(
+                (
+                    max_batch_size,
+                    1,
+                    self.head_dim // 8,
+                    max_seq_len,
+                    8,
+                )
+            )
+            .cuda()
+            .half()
+        )  # added to half
+
+        self.rotary_emb = RotaryEmbedding(self.head_dim)
+        self.rope_theta = args.rope_theta
+        self.rope_scaling = args.rope_scaling
+        if self.rope_scaling is None:
+            self.rope_scaling = 1.0
+        else:
+            self.rope_scaling = 1.0 / self.rope_scaling["factor"]
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        start_pos: int,
+        mask: Optional[torch.Tensor],
+    ):
+        bsz, seqlen, _ = x.shape
+
+        xqkv = self.query_key_value(x)
+        xqkv = xqkv.view(bsz, seqlen, self.n_local_heads + 2, self.head_dim)
+        xq = xqkv[:, :, :-2]
+        xk = xqkv[:, :, [-2]]
+        xv = xqkv[:, :, [-1]]
+
+        if seqlen > 1:
+            xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim)
+            xk = xk.view(bsz, seqlen, 1, self.head_dim)
+            xv = xv.view(bsz, seqlen, 1, self.head_dim)
+
+            xq, xk = self.rotary_emb(xq, xk)
+            xq = (
+                xq.reshape(bsz, self.n_local_heads, seqlen, self.head_dim)
+                .permute(0, 2, 1, 3)
+                .contiguous()
+            )
+            xk = (
+                xk.reshape(bsz, 1, seqlen, self.head_dim)
+                .permute(0, 2, 1, 3)
+                .contiguous()
+            )
+
+            self.cache_k = self.cache_k.to(xq)
+            self.cache_v = self.cache_v.to(xq)
+
+            values_store = xv.transpose(2, 1)
+            keys_store = (
+                xk.reshape(bsz, seqlen, 1, self.head_dim // 8, 8)
+                .permute(0, 2, 3, 1, 4)
+                .contiguous()
+            )
+
+            self.cache_v[:bsz, :, start_pos : start_pos + seqlen, :] = values_store
+            self.cache_k[:bsz, :, :, start_pos : start_pos + seqlen, :] = keys_store
+
+            keys = xk
+            values = xv
+
+            xq = xq.transpose(1, 2)
+            keys = keys.transpose(1, 2)
+            values = values.transpose(1, 2)
+            scores = torch.matmul(xq, keys.transpose(2, 3)) / math.sqrt(self.head_dim)
+            if mask is not None:
+                scores = scores + mask  # (bs, n_local_heads, slen, cache_len + slen)
+            scores = F.softmax(scores.float(), dim=-1).type_as(xq)
+            output = torch.matmul(scores, values)  # (bs, n_local_heads, slen, head_dim)
+            output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)
+        else:
+            # xq = xq[:, 0, :, :]
+            # xk = xk[:, 0, :, :]
+            # xv = xv[:, 0, :, :]
+            xq = xq.view(bsz, self.n_local_heads, self.head_dim)
+            xk = xk.view(bsz, 1, self.head_dim)
+            xv = xv.view(bsz, 1, self.head_dim)
+
+            output = awq_inference_engine.single_query_attention(
+                xq,
+                xk,
+                xv,
+                self.cache_k,
+                self.cache_v,
+                None,
+                # alibi position encodings
+                None,
+                start_pos,
+                self.head_dim,
+                self.rope_theta,
+                self.rope_scaling,
+                True,
+            )
+            output = output.reshape(bsz, 1, -1)
+
+        return self.dense(output)
+
+
+class FalconMLP(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+    ):
+        super().__init__()
+        self.dense_h_to_4h = nn.Linear(dim, 4 * dim, bias=False)
+        self.act = nn.GELU()
+        self.dense_4h_to_h = nn.Linear(4 * dim, dim, bias=False)
+
+    def forward(self, x):
+        x = self.act(self.dense_h_to_4h(x))
+        x = self.dense_4h_to_h(x)
+        return x
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, layer_id: int, args):
+        super().__init__()
+        self.n_heads = args.n_head
+        self.dim = args.hidden_size
+        self.head_dim = args.hidden_size // args.n_head
+        self.self_attention = FalconAttentionFused(args)
+        self.mlp = FalconMLP(dim=args.hidden_size)
+        self.layer_id = layer_id
+        self.input_layernorm = nn.LayerNorm(
+            args.hidden_size, eps=args.layer_norm_epsilon
+        )
+        # self.post_attention_layernorm = nn.LayerNorm(args.dim, eps=args.norm_eps)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        start_pos: int,
+        mask: Optional[torch.Tensor],
+    ):
+        layernorm_output = self.input_layernorm(x)
+        h_attn = x + self.self_attention.forward(layernorm_output, start_pos, mask)
+        h_mlp = self.mlp(layernorm_output)
+        out = h_attn + h_mlp
+        return out
+
+
+class Transformer(nn.Module):
+    def __init__(self, params):
+        super().__init__()
+        self.params = params
+        self.vocab_size = params.vocab_size
+        self.n_layers = params.n_layer
+
+        self.word_embeddings = nn.Embedding(params.vocab_size, params.hidden_size)
+
+        self.h = torch.nn.ModuleList()
+        for layer_id in range(params.n_layer):
+            self.h.append(TransformerBlock(layer_id, params))
+
+        self.ln_f = nn.LayerNorm(params.hidden_size, eps=params.layer_norm_epsilon)
+
+    @torch.inference_mode()
+    def forward(self, tokens: torch.Tensor, start_pos: int):
+        _bsz, seqlen = tokens.shape
+        h = self.word_embeddings(tokens)
+
+        mask = None
+        if seqlen > 1:
+            mask = torch.full(
+                (1, 1, seqlen, seqlen), float("-inf"), device=tokens.device
+            )
+            mask = torch.triu(mask, diagonal=start_pos + 1).type_as(h)
+        for layer in self.h:
+            h = layer(h, start_pos, mask)
+        h = self.ln_f(h)
+        return h
+
+
+class FalconForCausalLM(nn.Module):
+    def __init__(self, params):
+        super().__init__()
+        self.config = params
+        self.transformer = Transformer(params)
+        self.lm_head = nn.Linear(params.hidden_size, params.vocab_size, bias=False)
+
+    @torch.inference_mode()
+    def forward(self, tokens: torch.Tensor, start_pos: int):
+        h = self.transformer(tokens, start_pos)
+        output = self.lm_head(h)  # only compute last logits
+        return output.float()

llm-awq/tinychat/models/internvl/configuration_internvl.py

@@ -0,0 +1,204 @@
+# --------------------------------------------------------
+# InternVL
+# Copyright (c) 2024 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+import os
+from typing import Union
+import copy
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+from transformers import AutoConfig, LlamaConfig, Qwen2Config
+
+logger = logging.get_logger(__name__)
+
+
+class InternVisionConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`InternVisionModel`]. It is used to
+    instantiate a vision encoder according to the specified arguments, defining the model architecture.
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+    Args:
+        num_channels (`int`, *optional*, defaults to 3):
+            Number of color channels in the input images (e.g., 3 for RGB).
+        patch_size (`int`, *optional*, defaults to 14):
+            The size (resolution) of each patch.
+        image_size (`int`, *optional*, defaults to 224):
+            The size (resolution) of each image.
+        qkv_bias (`bool`, *optional*, defaults to `False`):
+            Whether to add a bias to the queries and values in the self-attention layers.
+        hidden_size (`int`, *optional*, defaults to 3200):
+            Dimensionality of the encoder layers and the pooler layer.
+        num_attention_heads (`int`, *optional*, defaults to 25):
+            Number of attention heads for each attention layer in the Transformer encoder.
+        intermediate_size (`int`, *optional*, defaults to 12800):
+            Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
+        qk_normalization (`bool`, *optional*, defaults to `True`):
+            Whether to normalize the queries and keys in the self-attention layers.
+        num_hidden_layers (`int`, *optional*, defaults to 48):
+            Number of hidden layers in the Transformer encoder.
+        use_flash_attn (`bool`, *optional*, defaults to `True`):
+            Whether to use flash attention mechanism.
+        hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`):
+            The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
+            `"relu"`, `"selu"` and `"gelu_new"` ``"gelu"` are supported.
+        layer_norm_eps (`float`, *optional*, defaults to 1e-6):
+            The epsilon used by the layer normalization layers.
+        dropout (`float`, *optional*, defaults to 0.0):
+            The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
+        drop_path_rate (`float`, *optional*, defaults to 0.0):
+            Dropout rate for stochastic depth.
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        initializer_factor (`float`, *optional*, defaults to 0.1):
+            A factor for layer scale.
+    """
+
+    model_type = 'intern_vit_6b'
+
+    def __init__(
+            self,
+            num_channels=3,
+            patch_size=14,
+            image_size=224,
+            qkv_bias=False,
+            hidden_size=3200,
+            num_attention_heads=25,
+            intermediate_size=12800,
+            qk_normalization=True,
+            num_hidden_layers=48,
+            use_flash_attn=True,
+            hidden_act='gelu',
+            norm_type='rms_norm',
+            layer_norm_eps=1e-6,
+            dropout=0.0,
+            drop_path_rate=0.0,
+            attention_dropout=0.0,
+            initializer_range=0.02,
+            initializer_factor=0.1,
+            **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.dropout = dropout
+        self.drop_path_rate = drop_path_rate
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_channels = num_channels
+        self.patch_size = patch_size
+        self.image_size = image_size
+        self.initializer_range = initializer_range
+        self.initializer_factor = initializer_factor
+        self.attention_dropout = attention_dropout
+        self.layer_norm_eps = layer_norm_eps
+        self.hidden_act = hidden_act
+        self.norm_type = norm_type
+        self.qkv_bias = qkv_bias
+        self.qk_normalization = qk_normalization
+        self.use_flash_attn = use_flash_attn
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs) -> 'PretrainedConfig':
+        config_dict, kwargs = cls.get_config_dict(pretrained_model_name_or_path, **kwargs)
+
+        if 'vision_config' in config_dict:
+            config_dict = config_dict['vision_config']
+
+        if 'model_type' in config_dict and hasattr(cls, 'model_type') and config_dict['model_type'] != cls.model_type:
+            logger.warning(
+                f"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
+                f'{cls.model_type}. This is not supported for all configurations of models and can yield errors.'
+            )
+
+        return cls.from_dict(config_dict, **kwargs)
+
+
+class InternVLChatConfig(PretrainedConfig):
+    model_type = 'internvl_chat'
+    is_composition = True
+
+    def __init__(
+            self,
+            vision_config=None,
+            llm_config=None,
+            use_backbone_lora=0,
+            use_llm_lora=0,
+            select_layer=-1,
+            force_image_size=None,
+            downsample_ratio=0.5,
+            template=None,
+            dynamic_image_size=False,
+            use_thumbnail=False,
+            ps_version='v1',
+            min_dynamic_patch=1,
+            max_dynamic_patch=6,
+            **kwargs):
+        super().__init__(**kwargs)
+
+        if vision_config is None:
+            vision_config = {'architectures': ['InternVisionModel']}
+            logger.info('vision_config is None. Initializing the InternVisionConfig with default values.')
+
+        if llm_config is None:
+            llm_config = {'architectures': ['Qwen2ForCausalLM']}
+            logger.info('llm_config is None. Initializing the LlamaConfig config with default values (`LlamaConfig`).')
+
+        self.vision_config = InternVisionConfig(**vision_config)
+        if llm_config.get('architectures')[0] == 'LlamaForCausalLM':
+            self.llm_config = LlamaConfig(**llm_config)
+        elif llm_config.get('architectures')[0] == 'Qwen2ForCausalLM':
+            self.llm_config = Qwen2Config(**llm_config)
+        else:
+            raise ValueError('Unsupported architecture: {}'.format(llm_config.get('architectures')[0]))
+        self.use_backbone_lora = use_backbone_lora
+        self.use_llm_lora = use_llm_lora
+        self.select_layer = select_layer
+        self.force_image_size = force_image_size
+        self.downsample_ratio = downsample_ratio
+        self.template = template
+        self.dynamic_image_size = dynamic_image_size
+        self.use_thumbnail = use_thumbnail
+        self.ps_version = ps_version  # pixel shuffle version
+        self.min_dynamic_patch = min_dynamic_patch
+        self.max_dynamic_patch = max_dynamic_patch
+        # By default, we use tie_word_embeddings=False for models of all sizes.
+        self.tie_word_embeddings = self.llm_config.tie_word_embeddings
+
+        logger.info(f'vision_select_layer: {self.select_layer}')
+        logger.info(f'ps_version: {self.ps_version}')
+        logger.info(f'min_dynamic_patch: {self.min_dynamic_patch}')
+        logger.info(f'max_dynamic_patch: {self.max_dynamic_patch}')
+
+    def to_dict(self):
+        """
+        Serializes this instance to a Python dictionary. Override the default [`~PretrainedConfig.to_dict`].
+
+        Returns:
+            `Dict[str, any]`: Dictionary of all the attributes that make up this configuration instance,
+        """
+        output = copy.deepcopy(self.__dict__)
+        output['vision_config'] = self.vision_config.to_dict()
+        output['llm_config'] = self.llm_config.to_dict()
+        output['model_type'] = self.__class__.model_type
+        output['use_backbone_lora'] = self.use_backbone_lora
+        output['use_llm_lora'] = self.use_llm_lora
+        output['select_layer'] = self.select_layer
+        output['force_image_size'] = self.force_image_size
+        output['downsample_ratio'] = self.downsample_ratio
+        output['template'] = self.template
+        output['dynamic_image_size'] = self.dynamic_image_size
+        output['use_thumbnail'] = self.use_thumbnail
+        output['ps_version'] = self.ps_version
+        output['min_dynamic_patch'] = self.min_dynamic_patch
+        output['max_dynamic_patch'] = self.max_dynamic_patch
+
+        return output

llm-awq/tinychat/models/internvl/conversation.py

@@ -0,0 +1,391 @@
+"""
+Conversation prompt templates.
+
+We kindly request that you import fastchat instead of copying this file if you wish to use it.
+If you have changes in mind, please contribute back so the community can benefit collectively and continue to maintain these valuable templates.
+
+Modified from https://github.com/lm-sys/FastChat/blob/main/fastchat/conversation.py
+"""
+
+import dataclasses
+from enum import IntEnum, auto
+from typing import Dict, List, Tuple, Union
+
+
+class SeparatorStyle(IntEnum):
+    """Separator styles."""
+
+    ADD_COLON_SINGLE = auto()
+    ADD_COLON_TWO = auto()
+    ADD_COLON_SPACE_SINGLE = auto()
+    NO_COLON_SINGLE = auto()
+    NO_COLON_TWO = auto()
+    ADD_NEW_LINE_SINGLE = auto()
+    LLAMA2 = auto()
+    CHATGLM = auto()
+    CHATML = auto()
+    CHATINTERN = auto()
+    DOLLY = auto()
+    RWKV = auto()
+    PHOENIX = auto()
+    ROBIN = auto()
+    FALCON_CHAT = auto()
+    CHATGLM3 = auto()
+    INTERNVL_ZH = auto()
+    MPT = auto()
+
+
+@dataclasses.dataclass
+class Conversation:
+    """A class that manages prompt templates and keeps all conversation history."""
+
+    # The name of this template
+    name: str
+    # The template of the system prompt
+    system_template: str = '{system_message}'
+    # The system message
+    system_message: str = ''
+    # The names of two roles
+    roles: Tuple[str] = ('USER', 'ASSISTANT')
+    # All messages. Each item is (role, message).
+    messages: List[List[str]] = ()
+    # The number of few shot examples
+    offset: int = 0
+    # The separator style and configurations
+    sep_style: SeparatorStyle = SeparatorStyle.ADD_COLON_SINGLE
+    sep: str = '\n'
+    sep2: str = None
+    # Stop criteria (the default one is EOS token)
+    stop_str: Union[str, List[str]] = None
+    # Stops generation if meeting any token in this list
+    stop_token_ids: List[int] = None
+
+    def get_prompt(self) -> str:
+        """Get the prompt for generation."""
+        system_prompt = self.system_template.format(system_message=self.system_message)
+        if self.sep_style == SeparatorStyle.ADD_COLON_SINGLE:
+            ret = system_prompt + self.sep
+            for role, message in self.messages:
+                if message:
+                    ret += role + ': ' + message + self.sep
+                else:
+                    ret += role + ':'
+            return ret
+        elif self.sep_style == SeparatorStyle.ADD_COLON_TWO:
+            seps = [self.sep, self.sep2]
+            ret = system_prompt + seps[0]
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += role + ': ' + message + seps[i % 2]
+                else:
+                    ret += role + ':'
+            return ret
+        elif self.sep_style == SeparatorStyle.ADD_COLON_SPACE_SINGLE:
+            ret = system_prompt + self.sep
+            for role, message in self.messages:
+                if message:
+                    ret += role + ': ' + message + self.sep
+                else:
+                    ret += role + ': '  # must be end with a space
+            return ret
+        elif self.sep_style == SeparatorStyle.ADD_NEW_LINE_SINGLE:
+            ret = '' if system_prompt == '' else system_prompt + self.sep
+            for role, message in self.messages:
+                if message:
+                    ret += role + '\n' + message + self.sep
+                else:
+                    ret += role + '\n'
+            return ret
+        elif self.sep_style == SeparatorStyle.NO_COLON_SINGLE:
+            ret = system_prompt
+            for role, message in self.messages:
+                if message:
+                    ret += role + message + self.sep
+                else:
+                    ret += role
+            return ret
+        elif self.sep_style == SeparatorStyle.NO_COLON_TWO:
+            seps = [self.sep, self.sep2]
+            ret = system_prompt
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += role + message + seps[i % 2]
+                else:
+                    ret += role
+            return ret
+        elif self.sep_style == SeparatorStyle.RWKV:
+            ret = system_prompt
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += (
+                        role
+                        + ': '
+                        + message.replace('\r\n', '\n').replace('\n\n', '\n')
+                    )
+                    ret += '\n\n'
+                else:
+                    ret += role + ':'
+            return ret
+        elif self.sep_style == SeparatorStyle.LLAMA2:
+            seps = [self.sep, self.sep2]
+            if self.system_message:
+                ret = system_prompt
+            else:
+                ret = '[INST] '
+            for i, (role, message) in enumerate(self.messages):
+                tag = self.roles[i % 2]
+                if message:
+                    if i == 0:
+                        ret += message + ' '
+                    else:
+                        ret += tag + ' ' + message + seps[i % 2]
+                else:
+                    ret += tag
+            return ret
+        elif self.sep_style == SeparatorStyle.CHATGLM:
+            # source: https://huggingface.co/THUDM/chatglm-6b/blob/1d240ba371910e9282298d4592532d7f0f3e9f3e/modeling_chatglm.py#L1302-L1308
+            # source2: https://huggingface.co/THUDM/chatglm2-6b/blob/e186c891cf64310ac66ef10a87e6635fa6c2a579/modeling_chatglm.py#L926
+            round_add_n = 1 if self.name == 'chatglm2' else 0
+            if system_prompt:
+                ret = system_prompt + self.sep
+            else:
+                ret = ''
+
+            for i, (role, message) in enumerate(self.messages):
+                if i % 2 == 0:
+                    ret += f'[Round {i//2 + round_add_n}]{self.sep}'
+
+                if message:
+                    ret += f'{role}：{message}{self.sep}'
+                else:
+                    ret += f'{role}：'
+            return ret
+        elif self.sep_style == SeparatorStyle.CHATML:
+            ret = '' if system_prompt == '' else system_prompt + self.sep + '\n'
+            for role, message in self.messages:
+                if message:
+                    ret += role + '\n' + message + self.sep + '\n'
+                else:
+                    ret += role + '\n'
+            return ret
+        elif self.sep_style == SeparatorStyle.CHATGLM3:
+            ret = ''
+            if self.system_message:
+                ret += system_prompt
+            for role, message in self.messages:
+                if message:
+                    ret += role + '\n' + ' ' + message
+                else:
+                    ret += role
+            return ret
+        elif self.sep_style == SeparatorStyle.CHATINTERN:
+            # source: https://huggingface.co/internlm/internlm-chat-7b-8k/blob/bd546fa984b4b0b86958f56bf37f94aa75ab8831/modeling_internlm.py#L771
+            seps = [self.sep, self.sep2]
+            ret = system_prompt
+            for i, (role, message) in enumerate(self.messages):
+                # if i % 2 == 0:
+                #     ret += "<s>"
+                if message:
+                    ret += role + ':' + message + seps[i % 2] + '\n'
+                else:
+                    ret += role + ':'
+            return ret
+        elif self.sep_style == SeparatorStyle.DOLLY:
+            seps = [self.sep, self.sep2]
+            ret = system_prompt
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += role + ':\n' + message + seps[i % 2]
+                    if i % 2 == 1:
+                        ret += '\n\n'
+                else:
+                    ret += role + ':\n'
+            return ret
+        elif self.sep_style == SeparatorStyle.PHOENIX:
+            ret = system_prompt
+            for role, message in self.messages:
+                if message:
+                    ret += role + ': ' + '<s>' + message + '</s>'
+                else:
+                    ret += role + ': ' + '<s>'
+            return ret
+        elif self.sep_style == SeparatorStyle.ROBIN:
+            ret = system_prompt + self.sep
+            for role, message in self.messages:
+                if message:
+                    ret += role + ':\n' + message + self.sep
+                else:
+                    ret += role + ':\n'
+            return ret
+        elif self.sep_style == SeparatorStyle.FALCON_CHAT:
+            ret = ''
+            if self.system_message:
+                ret += system_prompt + self.sep
+            for role, message in self.messages:
+                if message:
+                    ret += role + ': ' + message + self.sep
+                else:
+                    ret += role + ':'
+
+            return ret
+        elif self.sep_style == SeparatorStyle.INTERNVL_ZH:
+            seps = [self.sep, self.sep2]
+            ret = self.system_message + seps[0]
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += role + ': ' + message + seps[i % 2]
+                else:
+                    ret += role + ':'
+            return ret
+        elif self.sep_style == SeparatorStyle.MPT:
+            ret = system_prompt + self.sep
+            for role, message in self.messages:
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += role + message + self.sep
+                else:
+                    ret += role
+            return ret
+        else:
+            raise ValueError(f'Invalid style: {self.sep_style}')
+
+    def set_system_message(self, system_message: str):
+        """Set the system message."""
+        self.system_message = system_message
+
+    def append_message(self, role: str, message: str):
+        """Append a new message."""
+        self.messages.append([role, message])
+
+    def update_last_message(self, message: str):
+        """Update the last output.
+
+        The last message is typically set to be None when constructing the prompt,
+        so we need to update it in-place after getting the response from a model.
+        """
+        self.messages[-1][1] = message
+
+    def to_gradio_chatbot(self):
+        """Convert the conversation to gradio chatbot format."""
+        ret = []
+        for i, (role, msg) in enumerate(self.messages[self.offset :]):
+            if i % 2 == 0:
+                ret.append([msg, None])
+            else:
+                ret[-1][-1] = msg
+        return ret
+
+    def to_openai_api_messages(self):
+        """Convert the conversation to OpenAI chat completion format."""
+        ret = [{'role': 'system', 'content': self.system_message}]
+
+        for i, (_, msg) in enumerate(self.messages[self.offset :]):
+            if i % 2 == 0:
+                ret.append({'role': 'user', 'content': msg})
+            else:
+                if msg is not None:
+                    ret.append({'role': 'assistant', 'content': msg})
+        return ret
+
+    def copy(self):
+        return Conversation(
+            name=self.name,
+            system_template=self.system_template,
+            system_message=self.system_message,
+            roles=self.roles,
+            messages=[[x, y] for x, y in self.messages],
+            offset=self.offset,
+            sep_style=self.sep_style,
+            sep=self.sep,
+            sep2=self.sep2,
+            stop_str=self.stop_str,
+            stop_token_ids=self.stop_token_ids,
+        )
+
+    def dict(self):
+        return {
+            'template_name': self.name,
+            'system_message': self.system_message,
+            'roles': self.roles,
+            'messages': self.messages,
+            'offset': self.offset,
+        }
+
+
+# A global registry for all conversation templates
+conv_templates: Dict[str, Conversation] = {}
+
+
+def register_conv_template(template: Conversation, override: bool = False):
+    """Register a new conversation template."""
+    if not override:
+        assert (
+            template.name not in conv_templates
+        ), f'{template.name} has been registered.'
+
+    conv_templates[template.name] = template
+
+
+def get_conv_template(name: str) -> Conversation:
+    """Get a conversation template."""
+    return conv_templates[name].copy()
+
+
+# Both Hermes-2 and internlm2-chat are chatml-format conversation templates. The difference
+# is that during training, the preprocessing function for the Hermes-2 template doesn't add
+# <s> at the beginning of the tokenized sequence, while the internlm2-chat template does.
+# Therefore, they are completely equivalent during inference.
+register_conv_template(
+    Conversation(
+        name='Hermes-2',
+        system_template='<|im_start|>system\n{system_message}',
+        # note: The new system prompt was not used here to avoid changes in benchmark performance.
+        # system_message='我是书生·万象，英文名是InternVL，是由上海人工智能实验室、清华大学及多家合作单位联合开发的多模态大语言模型。',
+        system_message='你是由上海人工智能实验室联合商汤科技开发的书生多模态大模型，英文名叫InternVL, 是一个有用无害的人工智能助手。',
+        roles=('<|im_start|>user\n', '<|im_start|>assistant\n'),
+        sep_style=SeparatorStyle.MPT,
+        sep='<|im_end|>',
+        stop_str='<|endoftext|>',
+    )
+)
+
+
+register_conv_template(
+    Conversation(
+        name='internlm2-chat',
+        system_template='<|im_start|>system\n{system_message}',
+        # note: The new system prompt was not used here to avoid changes in benchmark performance.
+        # system_message='我是书生·万象，英文名是InternVL，是由上海人工智能实验室、清华大学及多家合作单位联合开发的多模态大语言模型。',
+        system_message='你是由上海人工智能实验室联合商汤科技开发的书生多模态大模型，英文名叫InternVL, 是一个有用无害的人工智能助手。',
+        roles=('<|im_start|>user\n', '<|im_start|>assistant\n'),
+        sep_style=SeparatorStyle.MPT,
+        sep='<|im_end|>',
+    )
+)
+
+
+register_conv_template(
+    Conversation(
+        name='phi3-chat',
+        system_template='<|system|>\n{system_message}',
+        # note: The new system prompt was not used here to avoid changes in benchmark performance.
+        # system_message='我是书生·万象，英文名是InternVL，是由上海人工智能实验室、清华大学及多家合作单位联合开发的多模态大语言模型。',
+        system_message='你是由上海人工智能实验室联合商汤科技开发的书生多模态大模型，英文名叫InternVL, 是一个有用无害的人工智能助手。',
+        roles=('<|user|>\n', '<|assistant|>\n'),
+        sep_style=SeparatorStyle.MPT,
+        sep='<|end|>',
+    )
+)
+
+
+register_conv_template(
+    Conversation(
+        name='internvl2_5',
+        system_template='<|im_start|>system\n{system_message}',
+        system_message='你是书生·万象，英文名是InternVL，是由上海人工智能实验室、清华大学及多家合作单位联合开发的多模态大语言模型。',
+        roles=('<|im_start|>user\n', '<|im_start|>assistant\n'),
+        sep_style=SeparatorStyle.MPT,
+        sep='<|im_end|>\n',
+    )
+)

llm-awq/tinychat/models/internvl/internvit.py

@@ -0,0 +1,425 @@
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from einops import rearrange
+from timm.layers import DropPath
+from torch import nn
+from transformers.activations import ACT2FN
+from transformers.modeling_outputs import (BaseModelOutput,
+                                           BaseModelOutputWithPooling)
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import logging
+
+from .configuration_internvl import InternVisionConfig
+
+try:
+    from flash_attn.bert_padding import pad_input, unpad_input
+    from flash_attn.flash_attn_interface import \
+        flash_attn_varlen_qkvpacked_func
+    has_flash_attn = True
+except:
+    print('FlashAttention2 is not installed.')
+    has_flash_attn = False
+
+logger = logging.get_logger(__name__)
+
+
+class FlashAttention(nn.Module):
+    """Implement the scaled dot product attention with softmax.
+    Arguments
+    ---------
+        softmax_scale: The temperature to use for the softmax attention.
+                      (default: 1/sqrt(d_keys) where d_keys is computed at
+                      runtime)
+        attention_dropout: The dropout rate to apply to the attention
+                           (default: 0.0)
+    """
+
+    def __init__(self, softmax_scale=None, attention_dropout=0.0, device=None, dtype=None):
+        super().__init__()
+        self.softmax_scale = softmax_scale
+        self.dropout_p = attention_dropout
+
+    def forward(self, qkv, key_padding_mask=None, causal=False, cu_seqlens=None,
+                max_s=None, need_weights=False):
+        """Implements the multihead softmax attention.
+        Arguments
+        ---------
+            qkv: The tensor containing the query, key, and value. (B, S, 3, H, D) if key_padding_mask is None
+                if unpadded: (nnz, 3, h, d)
+            key_padding_mask: a bool tensor of shape (B, S)
+        """
+        assert not need_weights
+        assert qkv.dtype in [torch.float16, torch.bfloat16]
+        assert qkv.is_cuda
+
+        if cu_seqlens is None:
+            batch_size = qkv.shape[0]
+            seqlen = qkv.shape[1]
+            if key_padding_mask is None:
+                qkv = rearrange(qkv, 'b s ... -> (b s) ...')
+                max_s = seqlen
+                cu_seqlens = torch.arange(0, (batch_size + 1) * seqlen, step=seqlen, dtype=torch.int32,
+                                          device=qkv.device)
+                output = flash_attn_varlen_qkvpacked_func(
+                    qkv, cu_seqlens, max_s, self.dropout_p if self.training else 0.0,
+                    softmax_scale=self.softmax_scale, causal=causal
+                )
+                output = rearrange(output, '(b s) ... -> b s ...', b=batch_size)
+            else:
+                nheads = qkv.shape[-2]
+                x = rearrange(qkv, 'b s three h d -> b s (three h d)')
+                x_unpad, indices, cu_seqlens, max_s = unpad_input(x, key_padding_mask)
+                x_unpad = rearrange(x_unpad, 'nnz (three h d) -> nnz three h d', three=3, h=nheads)
+                output_unpad = flash_attn_varlen_qkvpacked_func(
+                    x_unpad, cu_seqlens, max_s, self.dropout_p if self.training else 0.0,
+                    softmax_scale=self.softmax_scale, causal=causal
+                )
+                output = rearrange(pad_input(rearrange(output_unpad, 'nnz h d -> nnz (h d)'),
+                                             indices, batch_size, seqlen),
+                                   'b s (h d) -> b s h d', h=nheads)
+        else:
+            assert max_s is not None
+            output = flash_attn_varlen_qkvpacked_func(
+                qkv, cu_seqlens, max_s, self.dropout_p if self.training else 0.0,
+                softmax_scale=self.softmax_scale, causal=causal
+            )
+
+        return output, None
+
+
+class InternRMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+
+try:
+    from apex.normalization import FusedRMSNorm
+
+    InternRMSNorm = FusedRMSNorm  # noqa
+
+    logger.info('Discovered apex.normalization.FusedRMSNorm - will use it instead of InternRMSNorm')
+except ImportError:
+    # using the normal InternRMSNorm
+    pass
+except Exception:
+    logger.warning('discovered apex but it failed to load, falling back to InternRMSNorm')
+    pass
+
+
+NORM2FN = {
+    'rms_norm': InternRMSNorm,
+    'layer_norm': nn.LayerNorm,
+}
+
+
+class InternVisionEmbeddings(nn.Module):
+    def __init__(self, config: InternVisionConfig):
+        super().__init__()
+        self.config = config
+        self.embed_dim = config.hidden_size
+        self.image_size = config.image_size
+        self.patch_size = config.patch_size
+
+        self.class_embedding = nn.Parameter(
+            torch.randn(1, 1, self.embed_dim),
+        )
+
+        self.patch_embedding = nn.Conv2d(
+            in_channels=3, out_channels=self.embed_dim, kernel_size=self.patch_size, stride=self.patch_size
+        )
+
+        self.num_patches = (self.image_size // self.patch_size) ** 2
+        self.num_positions = self.num_patches + 1
+
+        self.position_embedding = nn.Parameter(torch.randn(1, self.num_positions, self.embed_dim))
+
+    def _get_pos_embed(self, pos_embed, H, W):
+        target_dtype = pos_embed.dtype
+        pos_embed = pos_embed.float().reshape(
+            1, self.image_size // self.patch_size, self.image_size // self.patch_size, -1).permute(0, 3, 1, 2)
+        pos_embed = F.interpolate(pos_embed, size=(H, W), mode='bicubic', align_corners=False). \
+            reshape(1, -1, H * W).permute(0, 2, 1).to(target_dtype)
+        return pos_embed
+
+    def forward(self, pixel_values: torch.FloatTensor) -> torch.Tensor:
+        target_dtype = self.patch_embedding.weight.dtype
+        patch_embeds = self.patch_embedding(pixel_values)  # shape = [*, channel, width, height]
+        batch_size, _, height, width = patch_embeds.shape
+        patch_embeds = patch_embeds.flatten(2).transpose(1, 2)
+        class_embeds = self.class_embedding.expand(batch_size, 1, -1).to(target_dtype)
+        embeddings = torch.cat([class_embeds, patch_embeds], dim=1)
+        position_embedding = torch.cat([
+            self.position_embedding[:, :1, :],
+            self._get_pos_embed(self.position_embedding[:, 1:, :], height, width)
+        ], dim=1)
+        embeddings = embeddings + position_embedding.to(target_dtype)
+        return embeddings
+
+
+class InternAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config: InternVisionConfig):
+        super().__init__()
+        self.config = config
+        self.embed_dim = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.use_flash_attn = config.use_flash_attn and has_flash_attn
+        if config.use_flash_attn and not has_flash_attn:
+            print('Warning: Flash Attention is not available, use_flash_attn is set to False.')
+        self.head_dim = self.embed_dim // self.num_heads
+        if self.head_dim * self.num_heads != self.embed_dim:
+            raise ValueError(
+                f'embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:'
+                f' {self.num_heads}).'
+            )
+
+        self.scale = self.head_dim ** -0.5
+        self.qkv = nn.Linear(self.embed_dim, 3 * self.embed_dim, bias=config.qkv_bias)
+        self.attn_drop = nn.Dropout(config.attention_dropout)
+        self.proj_drop = nn.Dropout(config.dropout)
+
+        self.qk_normalization = config.qk_normalization
+
+        if self.qk_normalization:
+            self.q_norm = InternRMSNorm(self.embed_dim, eps=config.layer_norm_eps)
+            self.k_norm = InternRMSNorm(self.embed_dim, eps=config.layer_norm_eps)
+
+        if self.use_flash_attn:
+            self.inner_attn = FlashAttention(attention_dropout=config.attention_dropout)
+        self.proj = nn.Linear(self.embed_dim, self.embed_dim)
+
+    def _naive_attn(self, x):
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv.unbind(0)  # make torchscript happy (cannot use tensor as tuple)
+
+        if self.qk_normalization:
+            B_, H_, N_, D_ = q.shape
+            q = self.q_norm(q.transpose(1, 2).flatten(-2, -1)).view(B_, N_, H_, D_).transpose(1, 2)
+            k = self.k_norm(k.transpose(1, 2).flatten(-2, -1)).view(B_, N_, H_, D_).transpose(1, 2)
+
+        attn = ((q * self.scale) @ k.transpose(-2, -1))
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+    def _flash_attn(self, x, key_padding_mask=None, need_weights=False):
+        qkv = self.qkv(x)
+        qkv = rearrange(qkv, 'b s (three h d) -> b s three h d', three=3, h=self.num_heads)
+
+        if self.qk_normalization:
+            q, k, v = qkv.unbind(2)
+            q = self.q_norm(q.flatten(-2, -1)).view(q.shape)
+            k = self.k_norm(k.flatten(-2, -1)).view(k.shape)
+            qkv = torch.stack([q, k, v], dim=2)
+
+        context, _ = self.inner_attn(
+            qkv, key_padding_mask=key_padding_mask, need_weights=need_weights, causal=False
+        )
+        outs = self.proj(rearrange(context, 'b s h d -> b s (h d)'))
+        outs = self.proj_drop(outs)
+        return outs
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        x = self._naive_attn(hidden_states) if not self.use_flash_attn else self._flash_attn(hidden_states)
+        return x
+
+
+class InternMLP(nn.Module):
+    def __init__(self, config: InternVisionConfig):
+        super().__init__()
+        self.config = config
+        self.act = ACT2FN[config.hidden_act]
+        self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)
+        self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.fc1(hidden_states)
+        hidden_states = self.act(hidden_states)
+        hidden_states = self.fc2(hidden_states)
+        return hidden_states
+
+
+class InternVisionEncoderLayer(nn.Module):
+    def __init__(self, config: InternVisionConfig, drop_path_rate: float):
+        super().__init__()
+        self.embed_dim = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.norm_type = config.norm_type
+
+        self.attn = InternAttention(config)
+        self.mlp = InternMLP(config)
+        self.norm1 = NORM2FN[self.norm_type](self.embed_dim, eps=config.layer_norm_eps)
+        self.norm2 = NORM2FN[self.norm_type](self.embed_dim, eps=config.layer_norm_eps)
+
+        self.ls1 = nn.Parameter(config.initializer_factor * torch.ones(self.embed_dim))
+        self.ls2 = nn.Parameter(config.initializer_factor * torch.ones(self.embed_dim))
+        self.drop_path1 = DropPath(drop_path_rate) if drop_path_rate > 0. else nn.Identity()
+        self.drop_path2 = DropPath(drop_path_rate) if drop_path_rate > 0. else nn.Identity()
+
+    def forward(
+            self,
+            hidden_states: torch.Tensor,
+    ) -> Tuple[torch.FloatTensor, Optional[torch.FloatTensor], Optional[Tuple[torch.FloatTensor]]]:
+        """
+        Args:
+            hidden_states (`Tuple[torch.FloatTensor, Optional[torch.FloatTensor]]`): input to the layer of shape `(batch, seq_len, embed_dim)`
+        """
+        hidden_states = hidden_states + self.drop_path1(self.attn(self.norm1(hidden_states).to(hidden_states.dtype)) * self.ls1)
+
+        hidden_states = hidden_states + self.drop_path2(self.mlp(self.norm2(hidden_states).to(hidden_states.dtype)) * self.ls2)
+
+        return hidden_states
+
+
+class InternVisionEncoder(nn.Module):
+    """
+    Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a
+    [`InternEncoderLayer`].
+
+    Args:
+        config (`InternConfig`):
+            The corresponding vision configuration for the `InternEncoder`.
+    """
+
+    def __init__(self, config: InternVisionConfig):
+        super().__init__()
+        self.config = config
+        # stochastic depth decay rule
+        dpr = [x.item() for x in torch.linspace(0, config.drop_path_rate, config.num_hidden_layers)]
+        self.layers = nn.ModuleList([
+            InternVisionEncoderLayer(config, dpr[idx]) for idx in range(config.num_hidden_layers)])
+        self.gradient_checkpointing = True
+
+    def forward(
+            self,
+            inputs_embeds,
+            output_hidden_states: Optional[bool] = None,
+            return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutput]:
+        r"""
+        Args:
+            inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+                Embedded representation of the inputs. Should be float, not int tokens.
+            output_hidden_states (`bool`, *optional*):
+                Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
+                for more detail.
+            return_dict (`bool`, *optional*):
+                Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+        """
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        encoder_states = () if output_hidden_states else None
+        hidden_states = inputs_embeds
+
+        for idx, encoder_layer in enumerate(self.layers):
+            if output_hidden_states:
+                encoder_states = encoder_states + (hidden_states,)
+            if self.gradient_checkpointing and self.training:
+                layer_outputs = torch.utils.checkpoint.checkpoint(
+                    encoder_layer,
+                    hidden_states)
+            else:
+                layer_outputs = encoder_layer(
+                    hidden_states,
+                )
+            hidden_states = layer_outputs
+
+        if output_hidden_states:
+            encoder_states = encoder_states + (hidden_states,)
+
+        if not return_dict:
+            return tuple(v for v in [hidden_states, encoder_states] if v is not None)
+        return BaseModelOutput(
+            last_hidden_state=hidden_states, hidden_states=encoder_states
+        )
+
+
+class InternVisionModel(PreTrainedModel):
+    main_input_name = 'pixel_values'
+    _supports_flash_attn_2 = True
+    supports_gradient_checkpointing = True
+    config_class = InternVisionConfig
+    _no_split_modules = ['InternVisionEncoderLayer']
+
+    def __init__(self, config: InternVisionConfig):
+        super().__init__(config)
+        self.config = config
+
+        self.embeddings = InternVisionEmbeddings(config)
+        self.encoder = InternVisionEncoder(config)
+
+    def resize_pos_embeddings(self, old_size, new_size, patch_size):
+        pos_emb = self.embeddings.position_embedding
+        _, num_positions, embed_dim = pos_emb.shape
+        cls_emb = pos_emb[:, :1, :]
+        pos_emb = pos_emb[:, 1:, :].reshape(1, old_size // patch_size, old_size // patch_size, -1).permute(0, 3, 1, 2)
+        pos_emb = F.interpolate(pos_emb.float(), size=new_size // patch_size, mode='bicubic', align_corners=False)
+        pos_emb = pos_emb.to(cls_emb.dtype).reshape(1, embed_dim, -1).permute(0, 2, 1)
+        pos_emb = torch.cat([cls_emb, pos_emb], dim=1)
+        self.embeddings.position_embedding = nn.Parameter(pos_emb)
+        self.embeddings.image_size = new_size
+        logger.info('Resized position embeddings from {} to {}'.format(old_size, new_size))
+
+    def get_input_embeddings(self):
+        return self.embeddings
+
+    def forward(
+            self,
+            pixel_values: Optional[torch.FloatTensor] = None,
+            output_hidden_states: Optional[bool] = None,
+            return_dict: Optional[bool] = None,
+            pixel_embeds: Optional[torch.FloatTensor] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPooling]:
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if pixel_values is None and pixel_embeds is None:
+            raise ValueError('You have to specify pixel_values or pixel_embeds')
+
+        if pixel_embeds is not None:
+            hidden_states = pixel_embeds
+        else:
+            if len(pixel_values.shape) == 4:
+                hidden_states = self.embeddings(pixel_values)
+            else:
+                raise ValueError(f'wrong pixel_values size: {pixel_values.shape}')
+        encoder_outputs = self.encoder(
+            inputs_embeds=hidden_states,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        last_hidden_state = encoder_outputs.last_hidden_state
+        pooled_output = last_hidden_state[:, 0, :]
+
+        if not return_dict:
+            return (last_hidden_state, pooled_output) + encoder_outputs[1:]
+
+        return BaseModelOutputWithPooling(
+            last_hidden_state=last_hidden_state,
+            pooler_output=pooled_output,
+            hidden_states=encoder_outputs.hidden_states,
+            attentions=encoder_outputs.attentions,
+        )

llm-awq/tinychat/models/internvl/media.py

@@ -0,0 +1,113 @@
+import math
+import numpy as np
+import torch
+import torchvision.transforms as T
+from decord import VideoReader, cpu
+from PIL import Image
+from torchvision.transforms.functional import InterpolationMode
+
+IMAGENET_MEAN = (0.485, 0.456, 0.406)
+IMAGENET_STD = (0.229, 0.224, 0.225)
+
+def build_transform(input_size):
+    MEAN, STD = IMAGENET_MEAN, IMAGENET_STD
+    transform = T.Compose([
+        T.Lambda(lambda img: img.convert('RGB') if img.mode != 'RGB' else img),
+        T.Resize((input_size, input_size), interpolation=InterpolationMode.BICUBIC),
+        T.ToTensor(),
+        T.Normalize(mean=MEAN, std=STD)
+    ])
+    return transform
+
+def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_size):
+    best_ratio_diff = float('inf')
+    best_ratio = (1, 1)
+    area = width * height
+    for ratio in target_ratios:
+        target_aspect_ratio = ratio[0] / ratio[1]
+        ratio_diff = abs(aspect_ratio - target_aspect_ratio)
+        if ratio_diff < best_ratio_diff:
+            best_ratio_diff = ratio_diff
+            best_ratio = ratio
+        elif ratio_diff == best_ratio_diff:
+            if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]:
+                best_ratio = ratio
+    return best_ratio
+
+def dynamic_preprocess(image, min_num=1, max_num=12, image_size=448, use_thumbnail=False):
+    orig_width, orig_height = image.size
+    aspect_ratio = orig_width / orig_height
+
+    # calculate the existing image aspect ratio
+    target_ratios = set(
+        (i, j) for n in range(min_num, max_num + 1) for i in range(1, n + 1) for j in range(1, n + 1) if
+        i * j <= max_num and i * j >= min_num)
+    target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])
+
+    # find the closest aspect ratio to the target
+    target_aspect_ratio = find_closest_aspect_ratio(
+        aspect_ratio, target_ratios, orig_width, orig_height, image_size)
+
+    # calculate the target width and height
+    target_width = image_size * target_aspect_ratio[0]
+    target_height = image_size * target_aspect_ratio[1]
+    blocks = target_aspect_ratio[0] * target_aspect_ratio[1]
+
+    # resize the image
+    resized_img = image.resize((target_width, target_height))
+    processed_images = []
+    for i in range(blocks):
+        box = (
+            (i % (target_width // image_size)) * image_size,
+            (i // (target_width // image_size)) * image_size,
+            ((i % (target_width // image_size)) + 1) * image_size,
+            ((i // (target_width // image_size)) + 1) * image_size
+        )
+        # split the image
+        split_img = resized_img.crop(box)
+        processed_images.append(split_img)
+    assert len(processed_images) == blocks
+    if use_thumbnail and len(processed_images) != 1:
+        thumbnail_img = image.resize((image_size, image_size))
+        processed_images.append(thumbnail_img)
+    return processed_images
+
+def load_image(image_file, input_size=448, max_num=12):
+    image = Image.open(image_file).convert('RGB')
+    transform = build_transform(input_size=input_size)
+    images = dynamic_preprocess(image, image_size=input_size, use_thumbnail=True, max_num=max_num)
+    pixel_values = [transform(image) for image in images]
+    pixel_values = torch.stack(pixel_values)
+    return pixel_values
+
+def get_index(bound, fps, max_frame, first_idx=0, num_segments=32):
+    if bound:
+        start, end = bound[0], bound[1]
+    else:
+        start, end = -100000, 100000
+    start_idx = max(first_idx, round(start * fps))
+    end_idx = min(round(end * fps), max_frame)
+    seg_size = float(end_idx - start_idx) / num_segments
+    frame_indices = np.array([
+        int(start_idx + (seg_size / 2) + np.round(seg_size * idx))
+        for idx in range(num_segments)
+    ])
+    return frame_indices
+
+def load_video(video_path, bound=None, input_size=448, max_num=1, num_segments=32):
+    vr = VideoReader(video_path, ctx=cpu(0), num_threads=1)
+    max_frame = len(vr) - 1
+    fps = float(vr.get_avg_fps())
+
+    pixel_values_list, num_patches_list = [], []
+    transform = build_transform(input_size=input_size)
+    frame_indices = get_index(bound, fps, max_frame, first_idx=0, num_segments=num_segments)
+    for frame_index in frame_indices:
+        img = Image.fromarray(vr[frame_index].asnumpy()).convert('RGB')
+        img = dynamic_preprocess(img, image_size=input_size, use_thumbnail=True, max_num=max_num)
+        pixel_values = [transform(tile) for tile in img]
+        pixel_values = torch.stack(pixel_values)
+        num_patches_list.append(pixel_values.shape[0])
+        pixel_values_list.append(pixel_values)
+    
+    return pixel_values_list, num_patches_list

llm-awq/tinychat/models/internvl3.py

@@ -0,0 +1,383 @@
+import os
+from collections import defaultdict, deque
+from typing import Dict, List, Optional, Tuple, Union, Any
+import warnings
+from time import time
+
+import torch
+import torch.utils.checkpoint
+from torch import nn
+from torch.nn import CrossEntropyLoss
+
+import transformers
+from transformers import (AutoConfig, 
+                          AutoModel, 
+                          AutoTokenizer,
+                          GenerationConfig,
+                          PretrainedConfig, 
+                          PreTrainedModel)
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import ModelOutput, logging
+from transformers import modeling_utils
+
+from .internvl.configuration_internvl import InternVisionConfig, InternVLChatConfig
+from .internvl.internvit import InternVisionModel
+from .internvl.conversation import get_conv_template
+from .internvl.media import load_image, load_video
+
+from llava.media import Image, Video
+
+from .qwen2 import Qwen2ForCausalLM
+from .llama import LlamaForCausalLM
+
+try:
+    import flash_attn
+    has_flash_attn = True
+except ImportError:
+    print('FlashAttention2 is not installed.')
+    has_flash_attn = False
+
+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
+
+
+logger = logging.get_logger(__name__)
+
+
+class InternVL3(PreTrainedModel):
+    config_class = InternVLChatConfig
+    main_input_name = 'pixel_values'
+    base_model_prefix = 'language_model'
+    _supports_flash_attn_2 = True
+    supports_gradient_checkpointing = True
+    _no_split_modules = ['InternVisionModel', 'LlamaDecoderLayer', 'Qwen2DecoderLayer']
+        
+    def __init__(self, config: InternVLChatConfig, vision_model=None, language_model=None, use_flash_attn=True):
+        super().__init__(config)
+        
+        self.tokenizer = AutoTokenizer.from_pretrained(config.name_or_path, trust_remote_code=True, use_fast=False)
+
+        image_size = config.force_image_size or config.vision_config.image_size
+        patch_size = config.vision_config.patch_size
+        self.patch_size = patch_size
+        self.select_layer = config.select_layer
+        self.template = config.template
+        self.num_image_token = int((image_size // patch_size) ** 2 * (config.downsample_ratio ** 2))
+        self.downsample_ratio = config.downsample_ratio
+        self.ps_version = config.ps_version
+        use_flash_attn = use_flash_attn if has_flash_attn else False
+        config.vision_config.use_flash_attn = True if use_flash_attn else False
+        config.llm_config._attn_implementation = 'flash_attention_2' if use_flash_attn else 'eager'
+
+        logger.info(f'num_image_token: {self.num_image_token}')
+        logger.info(f'ps_version: {self.ps_version}')
+        if vision_model is not None:
+            self.vision_model = vision_model
+        else:
+            self.vision_model = InternVisionModel(config.vision_config)
+        if language_model is not None:
+            self.language_model = language_model
+        else:
+            if config.llm_config.architectures[0] == 'LlamaForCausalLM':
+                self.language_model = LlamaForCausalLM(config.llm_config)
+            elif config.llm_config.architectures[0] == 'Qwen2ForCausalLM':
+                self.language_model = Qwen2ForCausalLM(config.llm_config)
+            else:
+                raise NotImplementedError(f'{config.llm_config.architectures[0]} is not implemented.')
+
+        vit_hidden_size = config.vision_config.hidden_size
+        llm_hidden_size = config.llm_config.hidden_size
+
+        self.mlp1 = nn.Sequential(
+            nn.LayerNorm(vit_hidden_size * int(1 / self.downsample_ratio) ** 2),
+            nn.Linear(vit_hidden_size * int(1 / self.downsample_ratio) ** 2, llm_hidden_size),
+            nn.GELU(),
+            nn.Linear(llm_hidden_size, llm_hidden_size)
+        )
+
+        self.img_context_token_id = None
+        self.conv_template = get_conv_template(self.template)
+        self.system_message = self.conv_template.system_message
+    
+    def freezed_module_patch(self):
+        self.vision_model.eval()
+        self.language_model.eval()
+        self.mlp1.eval()
+    
+    def pixel_shuffle(self, x, scale_factor=0.5):
+        n, w, h, c = x.size()
+        # N, W, H, C --> N, W, H * scale, C // scale
+        x = x.view(n, w, int(h * scale_factor), int(c / scale_factor))
+        # N, W, H * scale, C // scale --> N, H * scale, W, C // scale
+        x = x.permute(0, 2, 1, 3).contiguous()
+        # N, H * scale, W, C // scale --> N, H * scale, W * scale, C // (scale ** 2)
+        x = x.view(n, int(h * scale_factor), int(w * scale_factor),
+                   int(c / (scale_factor * scale_factor)))
+        if self.ps_version == 'v1':
+            warnings.warn("In ps_version 'v1', the height and width have not been swapped back, "
+                          'which results in a transposed image.')
+        else:
+            x = x.permute(0, 2, 1, 3).contiguous()
+        return x
+    
+    @torch.inference_mode()
+    def prepare_media(self, conversation):
+        prompt = conversation[0]["value"]
+        media = {"image": [], "video": []}
+        for item in prompt:
+            if isinstance(item, Image):
+                media["image"].append(load_image(item.path))
+            if isinstance(item, Video):
+                pixel_values, num_patches_list = load_video(item.path)
+                media["video"].extend(pixel_values)
+        
+        return media, num_patches_list if media["video"] else None
+        
+    @torch.inference_mode()
+    def extract_features(self, pixel_values):
+        if self.select_layer == -1:
+            vit_embeds = self.vision_model(
+                pixel_values=pixel_values,
+                output_hidden_states=False,
+                return_dict=True).last_hidden_state
+        else:
+            vit_embeds = self.vision_model(
+                pixel_values=pixel_values,
+                output_hidden_states=True,
+                return_dict=True).hidden_states[self.select_layer]
+        vit_embeds = vit_embeds[:, 1:, :]
+
+        h = w = int(vit_embeds.shape[1] ** 0.5)
+        vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], h, w, -1)
+        vit_embeds = self.pixel_shuffle(vit_embeds, scale_factor=self.downsample_ratio)
+        vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], -1, vit_embeds.shape[-1])
+        vit_embeds = self.mlp1(vit_embeds)
+        return vit_embeds
+        
+    @torch.inference_mode()
+    def _embed(
+        self,
+        input_ids: torch.Tensor,
+        media: Dict[str, List[torch.Tensor]],
+        media_config: Dict[str, Dict[str, Any]],
+        labels: Optional[torch.Tensor],
+        attention_mask: Optional[torch.Tensor],
+    ):
+        attention_mask = (
+            attention_mask
+            if attention_mask is not None
+            else torch.ones_like(input_ids, dtype=torch.bool)
+        )
+        
+        if media["image"]:
+            pixel_values = torch.cat(media["image"], dim=0).half().cuda()
+        elif media["video"]:
+            pixel_values = torch.cat(media["video"], dim=0).half().cuda()
+            
+        vit_embeds = self.extract_features(pixel_values)
+        
+        input_embeds = self.language_model.get_input_embeddings()(input_ids)
+        B, N, C = input_embeds.shape
+        input_embeds = input_embeds.reshape(B * N, C)
+        
+        input_ids = input_ids.reshape(B * N)
+        selected = (input_ids == self.img_context_token_id)
+        
+        input_embeds[selected] = input_embeds[selected] * 0.0 + vit_embeds.reshape(-1, C)
+        
+        input_embeds = input_embeds.reshape(B, N, C)
+        
+        return input_embeds, None, attention_mask
+    
+    @torch.inference_mode()
+    def benchmark(self, prompt: Union[str, List], quant_llm) -> None:
+        media = {"image": [], "video": []}
+        question = ""
+        for item in prompt:
+            if isinstance(item, str):
+                question += item
+            if isinstance(item, Image):
+                media["image"].append(load_image(item.path))
+            if isinstance(item, Video):
+                pixel_values, num_patches_list = load_video(item.path)
+                media["video"].extend(pixel_values)
+        
+        if media["image"]:
+            num_patches_list = [image.size(0) for image in media["image"]]
+        
+        if media["image"] and '<image>' not in question:
+            question = '<image>\n' + question
+        
+        if media["video"] and '<image>' not in question:
+            video_prefix = ''.join([f'Frame{i+1}: <image>\n' for i in range(len(num_patches_list))])
+            question = video_prefix + question
+        
+        template = get_conv_template(self.template)
+        template.system_message = self.system_message
+        eos_token_id = self.tokenizer.convert_tokens_to_ids(template.sep.strip())
+        
+        template.append_message(template.roles[0], question)
+        template.append_message(template.roles[1], None)
+        query = template.get_prompt()
+        
+        IMG_START_TOKEN = '<img>'
+        IMG_END_TOKEN = '</img>'
+        IMG_CONTEXT_TOKEN = '<IMG_CONTEXT>'
+        
+        img_context_token_id = self.tokenizer.convert_tokens_to_ids(IMG_CONTEXT_TOKEN)
+        self.img_context_token_id = img_context_token_id
+        
+        for num_patches in num_patches_list:
+            image_tokens = IMG_START_TOKEN + IMG_CONTEXT_TOKEN * self.num_image_token * num_patches + IMG_END_TOKEN
+            query = query.replace('<image>', image_tokens, 1)
+            
+        model_inputs = self.tokenizer(query, return_tensors='pt')
+        input_ids = model_inputs['input_ids'].to(self.device)
+        attention_mask = model_inputs['attention_mask'].to(self.device)
+        
+        for i in range(10):
+            torch.cuda.synchronize()
+            t_st = time()
+            inputs_embeds, _, attention_mask = self._embed(
+                input_ids=input_ids, 
+                media=media,
+                media_config=None,
+                labels=None,
+                attention_mask=attention_mask
+            )
+            torch.cuda.synchronize()
+            t_ed = time()
+            torch.cuda.empty_cache()
+            
+        if media["image"]:
+            print(
+                "Time of vision tower and others is {:.5f} s for {} images ({} x {} x {})".format(
+                    t_ed - t_st, sum(num_patches_list), media["image"][0].shape[1], media["image"][0].shape[2], media["image"][0].shape[3]
+                )
+            )
+        elif media["video"]:
+            print(
+                "Time of vision tower and others is {:.5f} s for {} video frames ({} x {} x {})".format(
+                    t_ed - t_st, sum(num_patches_list),  media["video"][0].shape[1], media["video"][0].shape[2], media["video"][0].shape[3]
+                )
+            )
+        output = self.language_model.benchmark(
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            quant_llm=quant_llm
+        )
+        response = self.tokenizer.decode(output[0], skip_special_tokens=True).strip()
+        
+        return response
+
+    @torch.inference_mode()
+    def stream_gen(
+        self,
+        input_ids,
+        media,
+        media_cfg,
+        start_pos,
+        chunk_prefilling,
+        quant_llm,
+        attention_mask=None,
+    ) -> str:
+        if media is None:
+            inputs_embeds = self.language_model.get_input_embeddings()(input_ids ).clone()
+        else:
+            inputs_embeds, _, _ = self._embed(input_ids, media, None, None, attention_mask)
+    
+        length = inputs_embeds.shape[1]
+        if quant_llm:
+            out = self.language_model(None, start_pos, inputs_embeds, chunk_prefilling)
+        else:
+            out = self.language_model.forwardfp16(None, start_pos, inputs_embeds, chunk_prefilling)
+        return out, length
+
+    @torch.inference_mode()
+    def forward(
+            self,
+            pixel_values: torch.FloatTensor,
+            input_ids: torch.LongTensor = None,
+            attention_mask: Optional[torch.Tensor] = None,
+            position_ids: Optional[torch.LongTensor] = None,
+            image_flags: Optional[torch.LongTensor] = None,
+            past_key_values: Optional[List[torch.FloatTensor]] = None,
+            labels: Optional[torch.LongTensor] = None,
+            use_cache: Optional[bool] = None,
+            output_attentions: Optional[bool] = None,
+            output_hidden_states: Optional[bool] = None,
+            return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        image_flags = image_flags.squeeze(-1)
+        input_embeds = self.language_model.get_input_embeddings()(input_ids).clone()
+
+        vit_embeds = self.extract_feature(pixel_values)
+        vit_embeds = vit_embeds[image_flags == 1]
+        vit_batch_size = pixel_values.shape[0]
+
+        B, N, C = input_embeds.shape
+        input_embeds = input_embeds.reshape(B * N, C)
+
+        if torch.distributed.is_initialized() and torch.distributed.get_rank() == 0:
+            print(f'dynamic ViT batch size: {vit_batch_size}, images per sample: {vit_batch_size / B}, dynamic token length: {N}')
+
+        input_ids = input_ids.reshape(B * N)
+        selected = (input_ids == self.img_context_token_id)
+        try:
+            input_embeds[selected] = input_embeds[selected] * 0.0 + vit_embeds.reshape(-1, C)
+        except Exception as e:
+            vit_embeds = vit_embeds.reshape(-1, C)
+            print(f'warning: {e}, input_embeds[selected].shape={input_embeds[selected].shape}, '
+                  f'vit_embeds.shape={vit_embeds.shape}')
+            n_token = min(selected.sum(), vit_embeds.size(0))
+            input_embeds[selected][:n_token] = input_embeds[selected][:n_token] * 0.0 + vit_embeds[:n_token]
+
+        input_embeds = input_embeds.reshape(B, N, C)
+
+        outputs = self.language_model(
+            inputs_embeds=input_embeds,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        logits = outputs.logits
+
+        loss = None
+        if labels is not None:
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            # Flatten the tokens
+            loss_fct = CrossEntropyLoss()
+            shift_logits = shift_logits.view(-1, self.language_model.config.vocab_size)
+            shift_labels = shift_labels.view(-1)
+            # Enable model parallelism
+            shift_labels = shift_labels.to(shift_logits.device)
+            loss = loss_fct(shift_logits, shift_labels)
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+    

llm-awq/tinychat/models/llama.py

@@ -0,0 +1,413 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# This software may be used and distributed according to the terms of the GNU General Public License version 3.
+
+from typing import Optional, Tuple
+from dataclasses import dataclass
+import math
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+import awq_inference_engine
+from transformers.models.llama.modeling_llama import LlamaRotaryEmbedding
+
+# from flash_attn.flash_attn_interface import flash_attn_unpadded_func
+
+import tinychat.utils.constants
+
+max_batch_size = tinychat.utils.constants.max_batch_size
+multiple_of = tinychat.utils.constants.llama_multiple_of
+max_seq_len = tinychat.utils.constants.max_seq_len
+from flash_attn import flash_attn_func
+
+
+class RMSNorm(torch.nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def _norm(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        output = torch.empty_like(x)
+        awq_inference_engine.layernorm_forward_cuda(x, self.weight, output, self.eps)
+        return output
+
+
+def precompute_freqs_cis(
+    dim: int, end: int, theta: float = 10000.0, scale: float = 1.0
+):
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
+    t = torch.arange(end, device=freqs.device)  # type: ignore
+    freqs = torch.outer(t * scale, freqs).float()  # type: ignore
+
+    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # complex64
+    return freqs_cis
+
+
+def precompute_freqs(
+    dim: int, end: int, theta: float = 10000.0, scale: float = 1.0, device=None
+):
+    inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2).float().to(device) / dim))
+    seq = torch.arange(end, dtype=inv_freq.dtype, device=device)
+    freqs = torch.einsum("i , j -> i j", seq, inv_freq)
+    freqs = freqs.reshape(freqs.shape[0], 1, 1, -1)
+    return torch.cat((freqs, freqs), dim=-1)
+
+
+def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
+    ndim = x.ndim
+    assert 0 <= 1 < ndim
+    assert freqs_cis.shape == (x.shape[1], x.shape[-1])
+    shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
+    return freqs_cis.view(*shape)
+
+
+def apply_rotary_emb(
+    xq: torch.Tensor,
+    xk: torch.Tensor,
+    freqs_cis: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    # xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
+    # k_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
+    xq_ = torch.view_as_complex(
+        xq.float().reshape(*xq.shape[:-1], 2, -1).transpose(-2, -1).contiguous()
+    )
+    xk_ = torch.view_as_complex(
+        xk.float().reshape(*xk.shape[:-1], 2, -1).transpose(-2, -1).contiguous()
+    )
+    freqs_cis = reshape_for_broadcast(freqs_cis, xq_)
+    xq_out = torch.view_as_real(xq_ * freqs_cis).transpose(-2, -1).flatten(3)
+    xk_out = torch.view_as_real(xk_ * freqs_cis).transpose(-2, -1).flatten(3)
+    return xq_out.type_as(xq), xk_out.type_as(xk)
+
+
+class LlamaAttentionFused(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+        self.args = args
+        self.n_local_heads = args.num_attention_heads
+        self.hidden_size = args.hidden_size
+        self.num_heads = args.num_attention_heads
+        self.head_dim = self.hidden_size // self.num_heads
+
+        self.num_key_value_heads = args.num_key_value_heads
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.max_position_embeddings = args.max_position_embeddings
+        self.rope_theta = args.rope_theta
+        self.rope_scaling = args.rope_scaling
+        if self.rope_scaling is None:
+            self.rope_scaling = 1.0
+        else:
+            self.rope_scaling = 1.0 / self.rope_scaling["factor"]
+        self.kv_max_seq_len = min(max_seq_len, self.max_position_embeddings)
+        self.q_proj = nn.Linear(
+            self.hidden_size,
+            self.num_heads * self.head_dim,
+            bias=False,
+        )
+        self.k_proj = nn.Linear(
+            self.hidden_size,
+            self.num_key_value_heads * self.head_dim,
+            bias=False,
+        )
+        self.v_proj = nn.Linear(
+            self.hidden_size,
+            self.num_key_value_heads * self.head_dim,
+            bias=False,
+        )
+        self.o_proj = nn.Linear(
+            self.num_heads * self.head_dim,
+            self.hidden_size,
+            bias=False,
+        )
+
+        # following fastertransformer definition
+        self.cache_v = (
+            torch.zeros(
+                (
+                    max_batch_size,
+                    self.num_key_value_heads,
+                    # args.max_position_embeddings,
+                    self.kv_max_seq_len,
+                    self.head_dim,
+                )
+            )
+            .cuda()
+            .half()
+        )  # added to half
+        # 8: pack 8 fp16 in FT, if fp32 then use 4
+        self.cache_k = (
+            torch.zeros(
+                (
+                    max_batch_size,
+                    self.num_key_value_heads,
+                    self.head_dim // 8,
+                    # args.max_position_embeddings,
+                    self.kv_max_seq_len,
+                    8,
+                )
+            )
+            .cuda()
+            .half()
+        )  # added to half
+        # dummy
+        self.rotary_emb = LlamaRotaryEmbedding(
+            self.head_dim, max_position_embeddings=2048, device="cuda:0"
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        start_pos: int,
+        freqs_cis: torch.Tensor,
+        mask: Optional[torch.Tensor],
+        chunk_prefilling: bool,
+    ):
+        bsz, seqlen, _ = x.shape
+        # xqkv = self.qkv_proj(x)
+        # xqkv = xqkv.view(bsz, seqlen, -1, self.n_local_heads, self.head_dim)
+        # xq = xqkv[:, :, 0]
+        # xk = xqkv[:, :, 1]
+        # xv = xqkv[:, :, 2]
+
+        xq, xk, xv = self.q_proj(x), self.k_proj(x), self.v_proj(x)
+
+        if seqlen > 1:
+            xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim)
+            xk = xk.view(bsz, seqlen, self.num_key_value_heads, self.head_dim)
+            xv = xv.view(bsz, seqlen, self.num_key_value_heads, self.head_dim)
+
+            xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)
+
+            self.cache_k = self.cache_k.to(xq)
+            self.cache_v = self.cache_v.to(xq)
+
+            values_store = xv.transpose(2, 1)
+            keys_store = (
+                xk.reshape(bsz, seqlen, self.num_key_value_heads, self.head_dim // 8, 8)
+                .permute(0, 2, 3, 1, 4)
+                .contiguous()
+            )
+
+            self.cache_v[:bsz, :, start_pos : start_pos + seqlen, :] = values_store
+            self.cache_k[:bsz, :, :, start_pos : start_pos + seqlen, :] = keys_store
+
+            if chunk_prefilling:
+                keys = self.cache_k[:, :, :, 0 : start_pos + seqlen, :]
+                keys = (
+                    keys.permute(0, 3, 1, 2, 4)
+                    .reshape(
+                        bsz, start_pos + seqlen, self.num_key_value_heads, self.head_dim
+                    )
+                    .contiguous()
+                )
+                values = self.cache_v[:, :, 0 : start_pos + seqlen, :]
+                values = (
+                    values.transpose(2, 1)
+                    .reshape(
+                        bsz, start_pos + seqlen, self.num_key_value_heads, self.head_dim
+                    )
+                    .contiguous()
+                )
+            else:
+                keys = xk
+                values = xv
+            output = flash_attn_func(
+                q=xq,
+                k=keys,
+                v=values,
+                causal=True,
+            )
+            output = output.contiguous().view(bsz, seqlen, -1)
+        else:
+            xq = xq.view(bsz, self.n_local_heads, self.head_dim)
+            xk = xk.view(bsz, self.num_key_value_heads, self.head_dim)
+            xv = xv.view(bsz, self.num_key_value_heads, self.head_dim)
+
+            output = awq_inference_engine.single_query_attention(
+                xq,
+                xk,
+                xv,
+                self.cache_k,
+                self.cache_v,
+                None,
+                # alibi position encodings
+                None,
+                start_pos,
+                self.head_dim,
+                self.rope_theta,
+                self.rope_scaling,
+                True,
+            )
+            output = output.reshape(bsz, 1, -1)
+
+        return self.o_proj(output)
+
+
+class LlamaMLP(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+        self.hidden_size = args.hidden_size
+        self.intermediate_size = args.intermediate_size
+
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+
+    def forward(self, x):
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, layer_id: int, args):
+        super().__init__()
+        self.n_heads = args.num_attention_heads
+        self.dim = args.hidden_size
+        self.head_dim = args.hidden_size // args.num_attention_heads
+        self.self_attn = LlamaAttentionFused(args)
+        self.mlp = LlamaMLP(args)
+        self.layer_id = layer_id
+        self.input_layernorm = RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        start_pos: int,
+        freqs_cis: torch.Tensor,
+        mask: Optional[torch.Tensor],
+        chunk_prefilling: bool,
+    ):
+        h = x + self.self_attn.forward(
+            self.input_layernorm(x), start_pos, freqs_cis, mask, chunk_prefilling
+        )
+        out = h + self.mlp.forward(self.post_attention_layernorm(h))
+        return out
+
+
+class Transformer(nn.Module):
+    def __init__(self, params):
+        super().__init__()
+        self.params = params
+        self.vocab_size = params.vocab_size
+        self.n_layers = params.num_hidden_layers
+
+        self.embed_tokens = nn.Embedding(params.vocab_size, params.hidden_size)
+
+        self.layers = torch.nn.ModuleList()
+        for layer_id in range(params.num_hidden_layers):
+            self.layers.append(TransformerBlock(layer_id, params))
+
+        self.norm = RMSNorm(params.hidden_size, eps=params.rms_norm_eps)
+
+        # Note (Haotian): rope_theta has to be defined here, otherwise context stage is wrong.
+        rope_scale = self.params.rope_scaling
+        if rope_scale is None:
+            rope_scale = 1.0
+        else:
+            rope_scale = 1.0 / rope_scale["factor"]
+        self.freqs = precompute_freqs(
+            self.params.hidden_size // self.params.num_attention_heads,
+            self.params.max_position_embeddings * 2,
+            self.params.rope_theta,
+            rope_scale,
+        )
+        self.freqs_cis = precompute_freqs_cis(
+            self.params.hidden_size // self.params.num_attention_heads,
+            self.params.max_position_embeddings * 2,
+            self.params.rope_theta,
+            rope_scale,
+        )
+
+    @torch.inference_mode()
+    def forward(
+        self,
+        tokens: torch.Tensor,
+        start_pos: int,
+        inputs_embeds: torch.Tensor = None,
+        chunk_prefilling: bool = False,
+    ):
+        if tokens is not None:
+            _bsz, seqlen = tokens.shape
+            h = self.embed_tokens(tokens)
+        else:
+            h = inputs_embeds
+            seqlen = inputs_embeds.shape[1]
+        self.freqs = self.freqs.to(h.device)
+        freqs = self.freqs[start_pos : start_pos + seqlen]
+
+        mask = None
+        if seqlen > 1:
+            mask = torch.full((1, 1, seqlen, seqlen), float("-inf"), device=h.device)
+            mask = torch.triu(mask, diagonal=1).type_as(h)
+            if chunk_prefilling:
+                mask_history = torch.zeros(
+                    (1, 1, seqlen, start_pos), dtype=torch.float16, device=h.device
+                ).type_as(h)
+                mask = torch.cat((mask_history, mask), dim=-1)
+        for layer in self.layers:
+            h = layer(h, start_pos, freqs, mask, chunk_prefilling)
+        h = h[:, -1:, :]  # Only the last token is useful
+        h = self.norm(h)
+        return h
+
+    @torch.inference_mode()
+    def forwardfp16(
+        self,
+        tokens: torch.Tensor,
+        start_pos: int,
+        inputs_embeds: torch.Tensor = None,
+        chunk_prefilling: bool = False,
+    ):
+        if tokens is not None:
+            _bsz, seqlen = tokens.shape
+            h = self.embed_tokens(tokens)
+        else:
+            h = inputs_embeds
+            seqlen = inputs_embeds.shape[1]
+        self.freqs_cis = self.freqs_cis.to(h.device)
+        freqs_cis = self.freqs_cis[start_pos : start_pos + seqlen]
+
+        mask = None
+        if seqlen > 1:
+            mask = torch.full((1, 1, seqlen, seqlen), float("-inf"), device=h.device)
+            mask = torch.triu(mask, diagonal=1).type_as(h)
+            if chunk_prefilling:
+                mask_history = torch.zeros(
+                    (1, 1, seqlen, start_pos), dtype=torch.float16, device=h.device
+                ).type_as(h)
+                mask = torch.cat((mask_history, mask), dim=-1)
+        for layer in self.layers:
+            h = layer(h, start_pos, freqs_cis, mask, chunk_prefilling)
+        h = h[:, -1:, :]  # Only the last token is useful
+        h = self.norm(h)
+        return h
+
+
+class LlamaForCausalLM(nn.Module):
+    def __init__(self, params):
+        super().__init__()
+        self.config = params
+        self.model = Transformer(params)
+        self.lm_head = nn.Linear(params.hidden_size, params.vocab_size, bias=False)
+
+    @torch.inference_mode()
+    def forward(
+        self,
+        tokens: torch.Tensor,
+        start_pos: int,
+        inputs_embeds: torch.Tensor = None,
+        chunk_prefilling=False,
+        quant=True,
+    ):
+        if quant:
+            h = self.model(tokens, start_pos, inputs_embeds, chunk_prefilling)
+        else:
+            h = self.model.forwardfp16(
+                tokens, start_pos, inputs_embeds, chunk_prefilling
+            )
+        output = self.lm_head(h)  # only compute last logits
+        return output.float()

llm-awq/tinychat/models/llava_base/llava_arch.py

@@ -0,0 +1,412 @@
+#    Modified from https://github.com/haotian-liu/LLaVA
+#    Copyright 2023 Haotian Liu
+#
+#    Licensed under the Apache License, Version 2.0 (the "License");
+#    you may not use this file except in compliance with the License.
+#    You may obtain a copy of the License at
+#
+#        http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS,
+#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#    See the License for the specific language governing permissions and
+#    limitations under the License.
+
+
+from abc import ABC, abstractmethod
+
+import torch
+import torch.nn as nn
+
+from .multimodal_encoder.builder import build_vision_tower
+from .multimodal_projector.builder import build_vision_projector
+
+from tinychat.utils.constants import (
+    LLAVA_IGNORE_INDEX,
+    LLAVA_DEFAULT_IMAGE_TOKEN_IDX,
+    LLAVA_DEFAULT_IMAGE_PATCH_TOKEN,
+    LLAVA_DEFAULT_IM_START_TOKEN,
+    LLAVA_DEFAULT_IM_END_TOKEN,
+)
+
+
+class LlavaMetaModel:
+    def __init__(self, config):
+        super(LlavaMetaModel, self).__init__(config)
+
+        if hasattr(config, "mm_vision_tower"):
+            self.vision_tower = build_vision_tower(config, delay_load=True)
+            self.mm_projector = build_vision_projector(config)
+
+    def get_vision_tower(self):
+        vision_tower = getattr(self, "vision_tower", None)
+        if type(vision_tower) is list:
+            vision_tower = vision_tower[0]
+        return vision_tower
+
+    def initialize_vision_modules(self, model_args, fsdp=None):
+        vision_tower = model_args.vision_tower
+        mm_vision_select_layer = model_args.mm_vision_select_layer
+        mm_vision_select_feature = model_args.mm_vision_select_feature
+        pretrain_mm_mlp_adapter = model_args.pretrain_mm_mlp_adapter
+
+        self.config.mm_vision_tower = vision_tower
+
+        if self.get_vision_tower() is None:
+            vision_tower = build_vision_tower(model_args)
+
+            if fsdp is not None and len(fsdp) > 0:
+                self.vision_tower = [vision_tower]
+            else:
+                self.vision_tower = vision_tower
+        else:
+            if fsdp is not None and len(fsdp) > 0:
+                vision_tower = self.vision_tower[0]
+            else:
+                vision_tower = self.vision_tower
+            vision_tower.load_model()
+
+        self.config.use_mm_proj = True
+        self.config.mm_projector_type = getattr(
+            model_args, "mm_projector_type", "linear"
+        )
+        self.config.mm_hidden_size = vision_tower.hidden_size
+        self.config.mm_vision_select_layer = mm_vision_select_layer
+        self.config.mm_vision_select_feature = mm_vision_select_feature
+
+        if getattr(self, "mm_projector", None) is None:
+            self.mm_projector = build_vision_projector(self.config)
+        else:
+            # In case it is frozen by LoRA
+            for p in self.mm_projector.parameters():
+                p.requires_grad = True
+
+        if pretrain_mm_mlp_adapter is not None:
+            mm_projector_weights = torch.load(
+                pretrain_mm_mlp_adapter, map_location="cpu"
+            )
+
+            def get_w(weights, keyword):
+                return {
+                    k.split(keyword + ".")[1]: v
+                    for k, v in weights.items()
+                    if keyword in k
+                }
+
+            self.mm_projector.load_state_dict(
+                get_w(mm_projector_weights, "mm_projector")
+            )
+
+
+class LlavaMetaForCausalLM(ABC):
+    @abstractmethod
+    def get_model(self):
+        pass
+
+    def get_vision_tower(self):
+        return self.get_model().get_vision_tower()
+
+    def encode_images(self, images):
+        vision_tower = self.get_model().get_vision_tower().half()
+        image_features = vision_tower(images)
+        image_features = self.get_model().mm_projector(image_features)
+        return image_features
+
+    def prepare_inputs_labels_for_multimodal(
+        self, input_ids, position_ids, attention_mask, past_key_values, labels, images
+    ):
+        vision_tower = self.get_vision_tower()
+        if vision_tower is None or images is None or input_ids.shape[1] == 1:
+            if (
+                past_key_values is not None
+                and vision_tower is not None
+                and images is not None
+                and input_ids.shape[1] == 1
+            ):
+                target_shape = past_key_values[-1][-1].shape[-2] + 1
+                attention_mask = torch.cat(
+                    (
+                        attention_mask,
+                        torch.ones(
+                            (
+                                attention_mask.shape[0],
+                                target_shape - attention_mask.shape[1],
+                            ),
+                            dtype=attention_mask.dtype,
+                            device=attention_mask.device,
+                        ),
+                    ),
+                    dim=1,
+                )
+                position_ids = torch.sum(attention_mask, dim=1).unsqueeze(-1) - 1
+            return (
+                input_ids,
+                position_ids,
+                attention_mask,
+                past_key_values,
+                None,
+                labels,
+            )
+
+        if type(images) is list or images.ndim == 5:
+            concat_images = torch.cat([image for image in images], dim=0)
+            image_features = self.encode_images(concat_images)
+            split_sizes = [image.shape[0] for image in images]
+            image_features = torch.split(image_features, split_sizes, dim=0)
+            image_features = [x.flatten(0, 1).to(self.device) for x in image_features]
+        else:
+            image_features = self.encode_images(images).to(self.device)
+
+        # TODO: image start / end is not implemented here to support pretraining.
+        if getattr(self.config, "tune_mm_mlp_adapter", False) and getattr(
+            self.config, "mm_use_im_start_end", False
+        ):
+            raise NotImplementedError
+
+        # Let's just add dummy tensors if they do not exist,
+        # it is a headache to deal with None all the time.
+        # But it is not ideal, and if you have a better idea,
+        # please open an issue / submit a PR, thanks.
+        _labels = labels
+        _position_ids = position_ids
+        _attention_mask = attention_mask
+        if attention_mask is None:
+            attention_mask = torch.ones_like(input_ids, dtype=torch.bool)
+        else:
+            attention_mask = attention_mask.bool()
+        if position_ids is None:
+            position_ids = torch.arange(
+                0, input_ids.shape[1], dtype=torch.long, device=input_ids.device
+            )
+        if labels is None:
+            labels = torch.full_like(input_ids, LLAVA_IGNORE_INDEX)
+
+        # remove the padding using attention_mask -- TODO: double check
+        input_ids = [
+            cur_input_ids[cur_attention_mask]
+            for cur_input_ids, cur_attention_mask in zip(input_ids, attention_mask)
+        ]
+        labels = [
+            cur_labels[cur_attention_mask]
+            for cur_labels, cur_attention_mask in zip(labels, attention_mask)
+        ]
+
+        new_input_embeds = []
+        new_labels = []
+        cur_image_idx = 0
+        for batch_idx, cur_input_ids in enumerate(input_ids):
+            num_images = (cur_input_ids == LLAVA_DEFAULT_IMAGE_TOKEN_IDX).sum()
+            if num_images == 0:
+                cur_image_features = image_features[cur_image_idx]
+                cur_input_embeds_1 = self.get_model().embed_tokens(cur_input_ids)
+                cur_input_embeds = torch.cat(
+                    [cur_input_embeds_1, cur_image_features[0:0]], dim=0
+                )
+                new_input_embeds.append(cur_input_embeds)
+                new_labels.append(labels[batch_idx])
+                cur_image_idx += 1
+                continue
+
+            image_token_indices = (
+                [-1]
+                + torch.where(cur_input_ids == LLAVA_DEFAULT_IMAGE_TOKEN_IDX)[
+                    0
+                ].tolist()
+                + [cur_input_ids.shape[0]]
+            )
+            cur_input_ids_noim = []
+            cur_labels = labels[batch_idx]
+            cur_labels_noim = []
+            for i in range(len(image_token_indices) - 1):
+                cur_input_ids_noim.append(
+                    cur_input_ids[
+                        image_token_indices[i] + 1 : image_token_indices[i + 1]
+                    ]
+                )
+                cur_labels_noim.append(
+                    cur_labels[image_token_indices[i] + 1 : image_token_indices[i + 1]]
+                )
+            split_sizes = [x.shape[0] for x in cur_labels_noim]
+            cur_input_embeds = self.get_model().embed_tokens(
+                torch.cat(cur_input_ids_noim)
+            )
+            cur_input_embeds_no_im = torch.split(cur_input_embeds, split_sizes, dim=0)
+            cur_new_input_embeds = []
+            cur_new_labels = []
+
+            for i in range(num_images + 1):
+                cur_new_input_embeds.append(cur_input_embeds_no_im[i])
+                cur_new_labels.append(cur_labels_noim[i])
+                if i < num_images:
+                    cur_image_features = image_features[cur_image_idx]
+                    cur_image_idx += 1
+                    cur_new_input_embeds.append(cur_image_features)
+                    cur_new_labels.append(
+                        torch.full(
+                            (cur_image_features.shape[0],),
+                            LLAVA_IGNORE_INDEX,
+                            device=cur_labels.device,
+                            dtype=cur_labels.dtype,
+                        )
+                    )
+
+            cur_new_input_embeds = torch.cat(cur_new_input_embeds)
+            cur_new_labels = torch.cat(cur_new_labels)
+
+            new_input_embeds.append(cur_new_input_embeds)
+            new_labels.append(cur_new_labels)
+
+        # Truncate sequences to max length as image embeddings can make the sequence longer
+        tokenizer_model_max_length = getattr(
+            self.config, "tokenizer_model_max_length", None
+        )
+        if tokenizer_model_max_length is not None:
+            new_input_embeds = [
+                x[:tokenizer_model_max_length] for x in new_input_embeds
+            ]
+            new_labels = [x[:tokenizer_model_max_length] for x in new_labels]
+
+        # Combine them
+        max_len = max(x.shape[0] for x in new_input_embeds)
+        batch_size = len(new_input_embeds)
+
+        new_input_embeds_padded = []
+        new_labels_padded = torch.full(
+            (batch_size, max_len),
+            LLAVA_IGNORE_INDEX,
+            dtype=new_labels[0].dtype,
+            device=new_labels[0].device,
+        )
+        attention_mask = torch.zeros(
+            (batch_size, max_len),
+            dtype=attention_mask.dtype,
+            device=attention_mask.device,
+        )
+        position_ids = torch.zeros(
+            (batch_size, max_len), dtype=position_ids.dtype, device=position_ids.device
+        )
+
+        for i, (cur_new_embed, cur_new_labels) in enumerate(
+            zip(new_input_embeds, new_labels)
+        ):
+            cur_len = cur_new_embed.shape[0]
+            if getattr(self.config, "tokenizer_padding_side", "right") == "left":
+                new_input_embeds_padded.append(
+                    torch.cat(
+                        (
+                            torch.zeros(
+                                (max_len - cur_len, cur_new_embed.shape[1]),
+                                dtype=cur_new_embed.dtype,
+                                device=cur_new_embed.device,
+                            ),
+                            cur_new_embed,
+                        ),
+                        dim=0,
+                    )
+                )
+                if cur_len > 0:
+                    new_labels_padded[i, -cur_len:] = cur_new_labels
+                    attention_mask[i, -cur_len:] = True
+                    position_ids[i, -cur_len:] = torch.arange(
+                        0, cur_len, dtype=position_ids.dtype, device=position_ids.device
+                    )
+            else:
+                new_input_embeds_padded.append(
+                    torch.cat(
+                        (
+                            cur_new_embed,
+                            torch.zeros(
+                                (max_len - cur_len, cur_new_embed.shape[1]),
+                                dtype=cur_new_embed.dtype,
+                                device=cur_new_embed.device,
+                            ),
+                        ),
+                        dim=0,
+                    )
+                )
+                if cur_len > 0:
+                    new_labels_padded[i, :cur_len] = cur_new_labels
+                    attention_mask[i, :cur_len] = True
+                    position_ids[i, :cur_len] = torch.arange(
+                        0, cur_len, dtype=position_ids.dtype, device=position_ids.device
+                    )
+
+        new_input_embeds = torch.stack(new_input_embeds_padded, dim=0)
+
+        if _labels is None:
+            new_labels = None
+        else:
+            new_labels = new_labels_padded
+
+        if _attention_mask is None:
+            attention_mask = None
+        else:
+            attention_mask = attention_mask.to(dtype=_attention_mask.dtype)
+
+        if _position_ids is None:
+            position_ids = None
+
+        return (
+            None,
+            position_ids,
+            attention_mask,
+            past_key_values,
+            new_input_embeds,
+            new_labels,
+        )
+
+    def initialize_vision_tokenizer(self, model_args, tokenizer):
+        if model_args.mm_use_im_patch_token:
+            tokenizer.add_tokens([LLAVA_DEFAULT_IMAGE_PATCH_TOKEN], special_tokens=True)
+            self.resize_token_embeddings(len(tokenizer))
+
+        if model_args.mm_use_im_start_end:
+            num_new_tokens = tokenizer.add_tokens(
+                [LLAVA_DEFAULT_IM_START_TOKEN, LLAVA_DEFAULT_IM_END_TOKEN],
+                special_tokens=True,
+            )
+            self.resize_token_embeddings(len(tokenizer))
+
+            if num_new_tokens > 0:
+                input_embeddings = self.get_input_embeddings().weight.data
+                output_embeddings = self.get_output_embeddings().weight.data
+
+                input_embeddings_avg = input_embeddings[:-num_new_tokens].mean(
+                    dim=0, keepdim=True
+                )
+                output_embeddings_avg = output_embeddings[:-num_new_tokens].mean(
+                    dim=0, keepdim=True
+                )
+
+                input_embeddings[-num_new_tokens:] = input_embeddings_avg
+                output_embeddings[-num_new_tokens:] = output_embeddings_avg
+
+            if model_args.tune_mm_mlp_adapter:
+                for p in self.get_input_embeddings().parameters():
+                    p.requires_grad = True
+                for p in self.get_output_embeddings().parameters():
+                    p.requires_grad = False
+
+            if model_args.pretrain_mm_mlp_adapter:
+                mm_projector_weights = torch.load(
+                    model_args.pretrain_mm_mlp_adapter, map_location="cpu"
+                )
+                embed_tokens_weight = mm_projector_weights["model.embed_tokens.weight"]
+                assert num_new_tokens == 2
+                if input_embeddings.shape == embed_tokens_weight.shape:
+                    input_embeddings[-num_new_tokens:] = embed_tokens_weight[
+                        -num_new_tokens:
+                    ]
+                elif embed_tokens_weight.shape[0] == num_new_tokens:
+                    input_embeddings[-num_new_tokens:] = embed_tokens_weight
+                else:
+                    raise ValueError(
+                        f"Unexpected embed_tokens_weight shape. Pretrained: {embed_tokens_weight.shape}. Current: {input_embeddings.shape}. Numer of new tokens: {num_new_tokens}."
+                    )
+        elif model_args.mm_use_im_patch_token:
+            if model_args.tune_mm_mlp_adapter:
+                for p in self.get_input_embeddings().parameters():
+                    p.requires_grad = False
+                for p in self.get_output_embeddings().parameters():
+                    p.requires_grad = False

llm-awq/tinychat/models/llava_base/multimodal_encoder/builder.py

@@ -0,0 +1,21 @@
+#    Modified from https://github.com/haotian-liu/LLaVA
+
+import os
+from .clip_encoder import CLIPVisionTower
+
+
+def build_vision_tower(vision_tower_cfg, **kwargs):
+    vision_tower = getattr(
+        vision_tower_cfg,
+        "mm_vision_tower",
+        getattr(vision_tower_cfg, "vision_tower", None),
+    )
+    is_absolute_path_exists = os.path.exists(vision_tower)
+    if (
+        is_absolute_path_exists
+        or vision_tower.startswith("openai")
+        or vision_tower.startswith("laion")
+    ):
+        return CLIPVisionTower(vision_tower, args=vision_tower_cfg, **kwargs)
+
+    raise ValueError(f"Unknown vision tower: {vision_tower}")

llm-awq/tinychat/models/llava_base/multimodal_encoder/clip_encoder.py

@@ -0,0 +1,97 @@
+#    Modified from https://github.com/haotian-liu/LLaVA
+
+import torch
+import torch.nn as nn
+
+from transformers import CLIPVisionModel, CLIPImageProcessor, CLIPVisionConfig
+
+
+class CLIPVisionTower(nn.Module):
+    def __init__(self, vision_tower, args, delay_load=False):
+        super().__init__()
+
+        self.is_loaded = False
+
+        self.vision_tower_name = vision_tower
+        self.select_layer = args.mm_vision_select_layer
+        self.select_feature = getattr(args, "mm_vision_select_feature", "patch")
+
+        if not delay_load:
+            self.load_model()
+        else:
+            self.cfg_only = CLIPVisionConfig.from_pretrained(self.vision_tower_name)
+
+    def load_model(self):
+        self.image_processor = CLIPImageProcessor.from_pretrained(
+            self.vision_tower_name
+        )
+        self.vision_tower = CLIPVisionModel.from_pretrained(self.vision_tower_name)
+        self.vision_tower.requires_grad_(False)
+
+        self.is_loaded = True
+
+    def feature_select(self, image_forward_outs):
+        image_features = image_forward_outs.hidden_states[self.select_layer]
+        if self.select_feature == "patch":
+            image_features = image_features[:, 1:]
+        elif self.select_feature == "cls_patch":
+            image_features = image_features
+        else:
+            raise ValueError(f"Unexpected select feature: {self.select_feature}")
+        return image_features
+
+    @torch.no_grad()
+    def forward(self, images):
+        if type(images) is list:
+            image_features = []
+            for image in images:
+                image_forward_out = self.vision_tower(
+                    image.to(device=self.device, dtype=self.dtype).unsqueeze(0),
+                    output_hidden_states=True,
+                )
+                image_feature = self.feature_select(image_forward_out).to(image.dtype)
+                image_features.append(image_feature)
+        else:
+            # import time
+            # torch.cuda.synchronize()
+            # image2 = images.to(self.dtype)
+            # print("model dtype:", self.dtype)
+            # print("image dtype:", images.dtype)
+            # st11 = time.time()
+            image_forward_outs = self.vision_tower(
+                images.to(device=self.device, dtype=self.dtype),
+                output_hidden_states=True,
+            )
+            # torch.cuda.synchronize()
+            # ed11 = time.time()
+            # print("hh2", (ed11 - st11)*1000)
+            image_features = self.feature_select(image_forward_outs).to(images.dtype)
+
+        return image_features
+
+    @property
+    def dummy_feature(self):
+        return torch.zeros(1, self.hidden_size, device=self.device, dtype=self.dtype)
+
+    @property
+    def dtype(self):
+        return self.vision_tower.dtype
+
+    @property
+    def device(self):
+        return self.vision_tower.device
+
+    @property
+    def config(self):
+        if self.is_loaded:
+            return self.vision_tower.config
+        else:
+            return self.cfg_only
+
+    @property
+    def hidden_size(self):
+        return self.config.hidden_size
+
+    @property
+    def num_patches(self):
+        return (self.config.image_size // self.config.patch_size) ** 2

llm-awq/tinychat/models/llava_base/multimodal_projector/builder.py

@@ -0,0 +1,72 @@
+#    Modified from https://github.com/haotian-liu/LLaVA
+
+import torch
+import torch.nn as nn
+import re
+
+
+class IdentityMap(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x, *args, **kwargs):
+        return x
+
+    @property
+    def config(self):
+        return {"mm_projector_type": "identity"}
+
+
+class SimpleResBlock(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.pre_norm = nn.LayerNorm(channels)
+
+        self.proj = nn.Sequential(
+            nn.Linear(channels, channels), nn.GELU(), nn.Linear(channels, channels)
+        )
+
+    def forward(self, x):
+        x = self.pre_norm(x)
+        return x + self.proj(x)
+
+
+def build_vision_projector(config, delay_load=False, **kwargs):
+    projector_type = getattr(config, "mm_projector_type", "linear")
+
+    if projector_type == "linear":
+        return nn.Linear(config.mm_hidden_size, config.hidden_size)
+
+    mlp_gelu_match = re.match(r"^mlp(\d+)x_gelu$", projector_type)
+    if mlp_gelu_match:
+        mlp_depth = int(mlp_gelu_match.group(1))
+        modules = [nn.Linear(config.mm_hidden_size, config.hidden_size)]
+        for _ in range(1, mlp_depth):
+            modules.append(nn.GELU())
+            modules.append(nn.Linear(config.hidden_size, config.hidden_size))
+        return nn.Sequential(*modules)
+
+    if projector_type == "identity":
+        return IdentityMap()
+
+    if projector_type == "linearclip":
+        # load min, max range
+        min_max_range = torch.load(config.min_max_range_path)
+        assert min_max_range is not None
+
+        class RangeClip(nn.Module):  # actually KNN projector
+            def __init__(self, min, max) -> None:
+                super().__init__()
+                self.register_buffer("min", min.detach().view(1, -1))
+                self.register_buffer("max", max.detach().view(1, -1))
+
+            def forward(self, x):
+                # dimension broadcast auto done
+                return torch.clamp(x, self.min.detach(), self.max.detach())
+
+        return nn.Sequential(
+            nn.Linear(config.mm_hidden_size, config.hidden_size),
+            RangeClip(*min_max_range),
+        )
+
+    raise ValueError(f"Unknown projector type: {projector_type}")

llm-awq/tinychat/models/llava_llama.py

@@ -0,0 +1,282 @@
+#    Modified from https://github.com/haotian-liu/LLaVA
+#    Copyright 2023 Haotian Liu
+#
+#    Licensed under the Apache License, Version 2.0 (the "License");
+#    you may not use this file except in compliance with the License.
+#    You may obtain a copy of the License at
+#
+#        http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS,
+#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#    See the License for the specific language governing permissions and
+#    limitations under the License.
+
+import os
+import warnings
+import shutil
+import torch
+import torch.nn as nn
+from typing import List, Optional, Tuple, Union
+
+from transformers import CLIPVisionModel
+
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from .llava_base.llava_arch import LlavaMetaModel, LlavaMetaForCausalLM
+from .llama import LlamaForCausalLM, Transformer
+
+
+class LlavaLlamaModel(LlavaMetaModel, Transformer):
+    def __init__(self, config):
+        super(LlavaLlamaModel, self).__init__(config)
+
+
+class LlavaLlamaForCausalLM(LlamaForCausalLM, LlavaMetaForCausalLM):
+    def __init__(self, config, dev="cuda"):
+        super(LlavaLlamaForCausalLM, self).__init__(config)
+        self.model = LlavaLlamaModel(config)
+        self.pretraining_tp = config.pretraining_tp
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.config = config
+        self.device = dev
+
+    def get_model(self):
+        return self.model
+
+    def default_inputs_embeds_for_multimodal(self, input_ids, inputs_embeds, images):
+        if inputs_embeds is None:
+            inputs_embeds = self.get_model().embed_tokens(input_ids)
+        vision_tower = self.get_vision_tower().vision_tower
+        from contextlib import nullcontext
+
+        if (
+            vision_tower is not None
+            and (input_ids.shape[1] != 1 or self.training)
+            and images is not None
+        ):
+            from tinychat.utils.constants import LLAVA_DEFAULT_IMAGE_PATCH_TOKEN_IDX
+
+            with (
+                nullcontext()
+                if getattr(self.config, "tune_vision_encoder", False)
+                else torch.no_grad()
+            ):
+                if type(images) is list:
+                    images = [
+                        image.unsqueeze(0) if len(image.shape) == 3 else image
+                        for image in images
+                    ]
+                    images = torch.cat(images, dim=0)
+                dtype = next(vision_tower.parameters()).dtype
+                if "visiontransformer" in vision_tower.__class__.__name__.lower():
+                    image_features = vision_tower(images.to(dtype))
+                else:
+                    image_forward_outs = vision_tower(
+                        images.to(dtype), output_hidden_states=True
+                    )
+                    select_hidden_state_layer = getattr(
+                        self.config, "mm_vision_select_layer", -1
+                    )
+                    if abs(select_hidden_state_layer) > 100:  # TOOD: find a better impl
+                        # -212 -> 12,
+                        idx1, idx2 = abs(select_hidden_state_layer) % 100, -(
+                            abs(select_hidden_state_layer) // 100
+                        )
+                        # print("selecting multiple indices", idx1, idx2)
+                        image_features = torch.cat(
+                            (
+                                image_forward_outs.hidden_states[idx1],
+                                image_forward_outs.hidden_states[idx2],
+                            ),
+                            dim=-1,
+                        )
+                    else:
+                        image_features = image_forward_outs.hidden_states[
+                            select_hidden_state_layer
+                        ]
+                if isinstance(vision_tower, CLIPVisionModel):  # clip case, not for sam
+                    image_features = image_features[:, 1:].to(images.dtype)  # (B, N, D)
+
+            image_features = self.model.mm_projector(image_features)
+
+            if hasattr(self.config, "neftune_alpha") and self.config.neftune_alpha > 0:
+                # print("using neftune tuning with alpha", self.config.neftune_alpha)
+                dims = torch.tensor(image_features.shape[-2] * image_features.shape[-1])
+                mag_norm = self.config.neftune_alpha / torch.sqrt(dims)
+                image_features = image_features + torch.zeros_like(
+                    image_features
+                ).uniform_(-mag_norm, mag_norm)
+
+            if self.config.mm_projector_type == "dsresampler":
+                dummy_feat_shape = (1, 1024, 1664)
+            elif self.config.mm_projector_type == "linear2":
+                dummy_feat_shape = (1, 256, self.config.mm_hidden_size * 2)
+            else:
+                dummy_feat_shape = (1, 256, self.config.mm_hidden_size)
+
+            dummy_image_features = torch.zeros(
+                *dummy_feat_shape,
+                device=inputs_embeds.device,
+                dtype=inputs_embeds.dtype,
+            )
+            dummy_image_features = self.model.mm_projector(dummy_image_features)[
+                0
+            ]  # (1, N, D)
+
+            new_input_embeds = []
+            cur_image_idx = 0
+
+            image_token_idx = []
+
+            num_patches = -1
+            for i_sample, (cur_input_ids, cur_input_embeds) in enumerate(
+                zip(input_ids, inputs_embeds)
+            ):
+                if (cur_input_ids == LLAVA_DEFAULT_IMAGE_PATCH_TOKEN_IDX).sum() == 0:
+                    # multimodal LLM, but the current sample is not multimodal
+                    cur_input_embeds = (
+                        cur_input_embeds + (0.0 * dummy_image_features).sum()
+                    )
+                    new_input_embeds.append(cur_input_embeds)
+                    # cur_image_idx += 1
+                    continue
+                # TODO: Need to fix if vision_tower.config.use_im_start_end == True
+                num_total_patches = (
+                    (cur_input_ids == LLAVA_DEFAULT_IMAGE_PATCH_TOKEN_IDX).sum().item()
+                )
+                masked_indices = torch.where(
+                    cur_input_ids == LLAVA_DEFAULT_IMAGE_PATCH_TOKEN_IDX
+                )[0]
+
+                while num_total_patches:
+                    if cur_image_idx >= image_features.shape[0]:  # SHOULD NOT HAPPEN!!!
+                        if self.training:
+                            print("%" * 20, "INDEXING ERROR!")
+                            break
+                        else:
+                            raise ValueError("INDEXING ERROR!")
+                    cur_image_features = image_features[cur_image_idx]
+                    num_patches = cur_image_features.shape[0]
+                    mask_index_start = masked_indices[0]
+                    masked_indices = masked_indices[num_patches:]
+
+                    image_token_idx.append(
+                        (
+                            i_sample,
+                            mask_index_start.item(),
+                            (mask_index_start + num_patches).item(),
+                        )
+                    )
+
+                    orig_embeds_params = None
+                    if orig_embeds_params is not None:
+                        cur_input_embeds = torch.cat(
+                            (
+                                cur_input_embeds[:mask_index_start].detach(),
+                                cur_image_features,
+                                cur_input_embeds[
+                                    mask_index_start + num_patches :
+                                ].detach(),
+                            ),
+                            dim=0,
+                        )
+                    else:
+                        cur_input_embeds = torch.cat(
+                            (
+                                cur_input_embeds[:mask_index_start],
+                                cur_image_features,
+                                cur_input_embeds[mask_index_start + num_patches :],
+                            ),
+                            dim=0,
+                        )
+                    num_total_patches -= num_patches
+                    assert num_total_patches >= 0, (num_total_patches, num_patches)
+                    cur_image_idx += 1
+
+                new_input_embeds.append(cur_input_embeds)
+                if self.training:
+                    if not masked_indices.numel() == 0:
+                        print("%" * 20, "ERROR! masked_indices not empty...")
+                else:
+                    assert masked_indices.numel() == 0
+
+            inputs_embeds = torch.stack(new_input_embeds, dim=0)
+
+        return inputs_embeds
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        start_pos: int = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        images: Optional[torch.FloatTensor] = None,
+        return_dict: Optional[bool] = None,
+        special_token: bool = False,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        if inputs_embeds is None:
+            if special_token:
+                (
+                    input_ids,
+                    position_ids,
+                    attention_mask,
+                    past_key_values,
+                    inputs_embeds,
+                    labels,
+                ) = self.prepare_inputs_labels_for_multimodal(
+                    input_ids,
+                    position_ids,
+                    attention_mask,
+                    past_key_values,
+                    labels,
+                    images,
+                )
+            else:
+                inputs_embeds = self.default_inputs_embeds_for_multimodal(
+                    input_ids, inputs_embeds, images
+                )
+                input_ids = None
+
+        if start_pos == None:
+            out = super().forward(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                inputs_embeds=inputs_embeds,
+                labels=labels,
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+            )
+        else:
+            out = super().forward(
+                tokens=input_ids,
+                start_pos=start_pos,
+                inputs_embeds=inputs_embeds,
+            )
+        return out
+
+    def prepare_inputs_for_generation(
+        self, input_ids, past_key_values=None, inputs_embeds=None, **kwargs
+    ):
+        images = kwargs.pop("images", None)
+        _inputs = super().prepare_inputs_for_generation(
+            input_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            **kwargs,
+        )
+        if images is not None:
+            _inputs["images"] = images
+        return _inputs

llm-awq/tinychat/models/mpt.py

@@ -0,0 +1,304 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# This software may be used and distributed according to the terms of the GNU General Public License version 3.
+
+from typing import Optional, Tuple
+from dataclasses import dataclass
+import math
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+import awq_inference_engine
+from transformers.models.llama.modeling_llama import LlamaRotaryEmbedding
+
+# from flash_attn.flash_attn_interface import flash_attn_unpadded_func
+
+import tinychat.utils.constants
+
+max_batch_size = tinychat.utils.constants.max_batch_size
+global_max_seq_len = tinychat.utils.constants.max_seq_len
+
+
+def gen_slopes(n_heads, alibi_bias_max=8):
+    _n_heads = 2 ** math.ceil(math.log2(n_heads))
+    m = torch.arange(1, _n_heads + 1, dtype=torch.float32)
+    m = m.mul(alibi_bias_max / _n_heads)
+    slopes = 1.0 / torch.pow(2, m)
+    if _n_heads != n_heads:
+        slopes = torch.concat([slopes[1::2], slopes[::2]])[:n_heads]
+    return slopes.view(1, n_heads, 1, 1)
+
+
+def build_alibi_bias(
+    n_heads, seq_len, full=False, alibi_bias_max=8, dtype=torch.float32
+):
+    alibi_bias = torch.arange(1 - seq_len, 1, dtype=torch.int32).view(1, 1, 1, seq_len)
+    if full:
+        alibi_bias = alibi_bias - torch.arange(1 - seq_len, 1, dtype=torch.int32).view(
+            1, 1, seq_len, 1
+        )
+        alibi_bias = alibi_bias.abs().mul(-1)
+    slopes = gen_slopes(n_heads, alibi_bias_max)
+    alibi_bias = alibi_bias * slopes
+    slopes = slopes.squeeze(0).squeeze(-1).squeeze(-1)
+    return slopes.to(dtype=dtype), alibi_bias.to(dtype=dtype)
+
+
+def _cast_if_autocast_enabled(tensor):
+    if torch.is_autocast_enabled():
+        if tensor.device.type == "cuda":
+            dtype = torch.get_autocast_gpu_dtype()
+        elif tensor.device.type == "cpu":
+            dtype = torch.get_autocast_cpu_dtype()
+        else:
+            raise NotImplementedError()
+        return tensor.to(dtype=dtype)
+    return tensor
+
+
+class LPLayerNorm(torch.nn.LayerNorm):
+    def __init__(self, normalized_shape, eps=1e-05, elementwise_affine=True):
+        super().__init__(
+            normalized_shape=normalized_shape,
+            eps=eps,
+            elementwise_affine=elementwise_affine,
+        )
+
+    def forward(self, x):
+        module_device = x.device
+        downcast_x = _cast_if_autocast_enabled(x)
+        downcast_weight = (
+            _cast_if_autocast_enabled(self.weight)
+            if self.weight is not None
+            else self.weight
+        )
+        downcast_bias = (
+            _cast_if_autocast_enabled(self.bias) if self.bias is not None else self.bias
+        )
+        with torch.autocast(enabled=False, device_type=module_device.type):
+            return torch.nn.functional.layer_norm(
+                downcast_x,
+                self.normalized_shape,
+                downcast_weight,
+                downcast_bias,
+                self.eps,
+            )
+
+
+class SharedEmbedding(nn.Embedding):
+    def forward(self, input: torch.Tensor, unembed: bool = False) -> torch.Tensor:
+        if unembed:
+            return F.linear(input, self.weight)
+        return super().forward(input)
+
+
+class MPTAttentionFused(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+        self.args = args
+        self.n_local_heads = args.n_heads
+        self.head_dim = args.d_model // args.n_heads
+        args.max_seq_len = min(args.max_seq_len, global_max_seq_len)
+
+        self.Wqkv = nn.Linear(
+            args.d_model,
+            args.n_heads * self.head_dim * 3,
+            bias=False,
+        )
+
+        self.out_proj = nn.Linear(
+            args.n_heads * self.head_dim,
+            args.d_model,
+            bias=False,
+        )
+
+        # following fastertransformer definition
+
+        self.cache_v = (
+            torch.zeros(
+                (
+                    max_batch_size,
+                    self.n_local_heads,
+                    args.max_seq_len,
+                    self.head_dim,
+                )
+            )
+            .cuda()
+            .half()
+        )  # added to half
+        # 8: pack 8 fp16 in FT, if fp32 then use 4
+        self.cache_k = (
+            torch.zeros(
+                (
+                    max_batch_size,
+                    self.n_local_heads,
+                    self.head_dim // 8,
+                    args.max_seq_len,
+                    8,
+                )
+            )
+            .cuda()
+            .half()
+        )  # added to half
+
+        alibi_slopes, alibi_bias = build_alibi_bias(
+            self.n_local_heads, args.max_seq_len
+        )
+        # TODO (Haotian): fix device
+        self.alibi_slopes = alibi_slopes.float().to("cuda:0")
+        self.alibi_bias = alibi_bias.to("cuda:0")
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        start_pos: int,
+        mask: Optional[torch.Tensor],
+    ):
+        bsz, seqlen, _ = x.shape
+        xqkv = self.Wqkv(x)
+        xqkv = xqkv.view(bsz, seqlen, -1, self.n_local_heads, self.head_dim)
+        xq = xqkv[:, :, 0]
+        xk = xqkv[:, :, 1]
+        xv = xqkv[:, :, 2]
+
+        if seqlen > 1:
+            xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim)
+            xk = xk.view(bsz, seqlen, self.n_local_heads, self.head_dim)
+            xv = xv.view(bsz, seqlen, self.n_local_heads, self.head_dim)
+
+            self.cache_k = self.cache_k.to(xq)
+            self.cache_v = self.cache_v.to(xq)
+
+            values_store = xv.transpose(2, 1)
+            keys_store = (
+                xk.reshape(bsz, seqlen, self.n_local_heads, self.head_dim // 8, 8)
+                .permute(0, 2, 3, 1, 4)
+                .contiguous()
+            )
+
+            self.cache_v[:bsz, :, start_pos : start_pos + seqlen, :] = values_store
+            self.cache_k[:bsz, :, :, start_pos : start_pos + seqlen, :] = keys_store
+
+            keys = xk
+            values = xv
+
+            xq = xq.transpose(1, 2)
+            keys = keys.transpose(1, 2)
+            values = values.transpose(1, 2)
+            scores = torch.matmul(xq, keys.transpose(2, 3)) / math.sqrt(self.head_dim)
+            scores += self.alibi_bias[..., :seqlen]
+            if mask is not None:
+                scores = scores + mask  # (bs, n_local_heads, slen, cache_len + slen)
+            scores = F.softmax(scores.float(), dim=-1).type_as(xq)
+            output = torch.matmul(scores, values)  # (bs, n_local_heads, slen, head_dim)
+            output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)
+        else:
+            # xq = xq[:, 0, :, :]
+            # xk = xk[:, 0, :, :]
+            # xv = xv[:, 0, :, :]
+            xq = xq.view(bsz, self.n_local_heads, self.head_dim)
+            xk = xk.view(bsz, self.n_local_heads, self.head_dim)
+            xv = xv.view(bsz, self.n_local_heads, self.head_dim)
+            output = awq_inference_engine.single_query_attention(
+                xq,
+                xk,
+                xv,
+                self.cache_k,
+                self.cache_v,
+                None,
+                # with alibi embedding
+                self.alibi_slopes.float(),
+                start_pos,
+                # rotary embed dim = 0 => no rotary embedding
+                0,
+                10000,
+                1.0,
+                True,
+            )
+            output = output.reshape(bsz, 1, -1)
+
+        return self.out_proj(output)
+
+
+class MPTMLP(nn.Module):
+    def __init__(self, d_model: int, expansion_ratio: int):
+        super().__init__()
+        self.up_proj = nn.Linear(d_model, expansion_ratio * d_model, bias=False)
+        self.act = nn.GELU(approximate="none")
+        self.down_proj = nn.Linear(expansion_ratio * d_model, d_model, bias=False)
+        self.down_proj._is_residual = True
+
+    def forward(self, x):
+        return self.down_proj(self.act(self.up_proj(x)))
+
+
+class MPTBlock(nn.Module):
+    def __init__(self, layer_id: int, args):
+        super().__init__()
+        self.n_heads = args.n_heads
+        self.dim = args.d_model
+        self.head_dim = args.d_model // args.n_heads
+        self.attn = MPTAttentionFused(args)
+        self.ffn = MPTMLP(d_model=args.d_model, expansion_ratio=4)
+        self.layer_id = layer_id
+        self.norm_1 = LPLayerNorm(args.d_model, eps=1e-6)
+        self.norm_2 = LPLayerNorm(args.d_model, eps=1e-6)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        start_pos: int,
+        mask: Optional[torch.Tensor],
+    ):
+        h = x + self.attn.forward(self.norm_1(x), start_pos, mask)
+        out = h + self.ffn.forward(self.norm_2(h))
+        return out
+
+
+class Transformer(nn.Module):
+    def __init__(self, params):
+        super().__init__()
+        self.params = params
+        self.vocab_size = params.vocab_size
+        self.n_layers = params.n_layers
+
+        self.wte = SharedEmbedding(params.vocab_size, params.d_model)
+
+        self.blocks = torch.nn.ModuleList()
+        for layer_id in range(params.n_layers):
+            self.blocks.append(MPTBlock(layer_id, params))
+
+        self.norm_f = LPLayerNorm(params.d_model, eps=1e-6)
+
+    @torch.inference_mode()
+    def forward(self, tokens: torch.Tensor, start_pos: int):
+        _bsz, seqlen = tokens.shape
+        h = self.wte(tokens)
+
+        mask = None
+        if seqlen > 1:
+            mask = torch.full(
+                (1, 1, seqlen, seqlen), float("-inf"), device=tokens.device
+            )
+            mask = torch.triu(mask, diagonal=start_pos + 1).type_as(h)
+        for layer in self.blocks:
+            h = layer(h, start_pos, mask)
+        h = self.norm_f(h)
+        return h
+
+
+class MPTForCausalLM(nn.Module):
+    def __init__(self, params):
+        super().__init__()
+        self.config = params
+        self.transformer = Transformer(params)
+        if params.no_bias:
+            for module in self.modules():
+                if hasattr(module, "bias") and isinstance(module.bias, nn.Parameter):
+                    module.register_parameter("bias", None)
+
+    @torch.inference_mode()
+    def forward(self, tokens: torch.Tensor, start_pos: int):
+        h = self.transformer(tokens, start_pos)
+        output = self.transformer.wte(h, unembed=True)  # only compute last logits
+        return output.float()

llm-awq/tinychat/models/nvila/builder.py

@@ -0,0 +1,291 @@
+# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import math
+import os
+import os.path as osp
+import warnings
+from dataclasses import asdict
+from typing import Tuple
+
+import torch
+from huggingface_hub import file_exists, repo_exists
+from huggingface_hub.utils import HFValidationError
+from transformers import (
+    AutoConfig,
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    PretrainedConfig,
+    PreTrainedModel,
+    PreTrainedTokenizer,
+)
+
+from llava.constants import MEDIA_TOKENS
+from llava.model.utils import packing
+from llava.utils.logging import logger
+from llava.utils.tokenizer import infer_stop_tokens
+
+
+def has_tokenizer(repo_id_or_path: str) -> bool:
+    # Check if the tokenizer is in a local directory
+    if osp.exists(osp.join(repo_id_or_path, "tokenizer_config.json")):
+        return True
+
+    # Check if the tokenizer is in a Hugging Face Hub repo
+    try:
+        return repo_exists(repo_id_or_path) and file_exists(
+            repo_id_or_path, "tokenizer_config.json"
+        )
+    except HFValidationError:
+        return False
+
+
+def context_length_extension(config):
+    orig_ctx_len = getattr(config, "max_position_embeddings", None)
+    model_max_length = getattr(config, "model_max_length", None)
+    if orig_ctx_len and model_max_length > orig_ctx_len:
+        print(f"Scaling RoPE from {orig_ctx_len} to {model_max_length}")
+        scaling_factor = float(math.ceil(model_max_length / orig_ctx_len))
+        config.rope_scaling = {"type": "linear", "factor": scaling_factor}
+    return config
+
+
+def build_llm_and_tokenizer(
+    model_name_or_path: str,
+    config: PretrainedConfig,
+    attn_implementation=None,
+    model_max_length=None,
+    *args,
+    **kwargs,
+) -> Tuple[PreTrainedModel, PreTrainedTokenizer]:
+    # print(model_name_or_path)
+    llm_cfg = AutoConfig.from_pretrained(model_name_or_path)
+    llm_cfg._attn_implementation = attn_implementation
+    llm_cfg.model_max_length = model_max_length
+    if model_max_length is not None:
+        context_length_extension(llm_cfg)
+
+    # Quantization related
+    quantization_restore_from_checkpoint = False
+    if kwargs.get("quantize_model_class") is not None:
+        assert kwargs.get("model_args") is not None
+        quantize_model_class = kwargs.pop("quantize_model_class", None)
+        model_args = kwargs.pop("model_args", None)
+
+        if (
+            quantize_model_class == "QLlamaForCausalLM"
+        ):  # TODO: Also change the name of this class
+            from .qllama import QLlamaConfig
+
+            llm_cfg.architectures = "QLlamaForCausalLM"
+            _attn_implementation = llm_cfg._attn_implementation
+            llm_cfg = QLlamaConfig(**llm_cfg.to_dict())
+            llm_cfg._attn_implementation = _attn_implementation
+        elif (
+            quantize_model_class == "QMemLlamaForCausalLM"
+        ):  # TODO: Also change the name of this class
+            from .qmemllama import QMemLlamaConfig
+
+            llm_cfg.architectures = "QMemLlamaForCausalLM"
+            llm_cfg = QMemLlamaConfig(**llm_cfg.to_dict())
+        elif quantize_model_class == "FP8LinearQwen2ForCausalLM":
+            from .configuration_quantize import QuantizationConfig
+            from .fp8linearqwen2 import FP8LinearQwen2Config
+
+            llm_cfg.architectures = "FP8LinearQwen2ForCausalLM"
+            coat_fp8_args = QuantizationConfig(**asdict(model_args))
+
+            # Remove the quantization args from llm_cfg and make it a independent config
+            model_args_dict = asdict(model_args)
+            for key in asdict(coat_fp8_args).keys():
+                model_args_dict.pop(key, None)
+
+            llm_cfg.coat_fp8_args = asdict(coat_fp8_args)
+            _attn_implementation = llm_cfg._attn_implementation
+
+            llm_cfg = FP8LinearQwen2Config(**llm_cfg.to_dict())
+            llm_cfg._attn_implementation = _attn_implementation
+
+        elif quantize_model_class == "FP8ActivationQwen2ForCausalLM":
+            from ..coat.activation.models._fp8_quantization_config import (
+                QuantizationConfig,
+            )
+            from .fp8activationqwen2 import FP8ActivationQwen2Config
+
+            quantization_restore_from_checkpoint = True
+
+            llm_cfg.architectures = "FP8ActivationQwen2ForCausalLM"
+            coat_fp8_args = QuantizationConfig(**asdict(model_args))
+
+            # Remove the quantization args from llm_cfg and make it a independent config
+            model_args_dict = asdict(model_args)
+            for key in asdict(coat_fp8_args).keys():
+                model_args_dict.pop(key, None)
+
+            llm_cfg.coat_fp8_args = asdict(coat_fp8_args)
+            _attn_implementation = llm_cfg._attn_implementation
+
+            llm_cfg = FP8ActivationQwen2Config(**llm_cfg.to_dict())
+            llm_cfg._attn_implementation = _attn_implementation
+
+        elif quantize_model_class == "FP8ActivationResidualQwen2ForCausalLM":
+            from ..coat.activation.models._fp8_quantization_config import (
+                QuantizationConfig,
+            )
+            from .fp8activationresidualqwen2 import FP8ActivationResidualQwen2Config
+
+            quantization_restore_from_checkpoint = True
+
+            llm_cfg.architectures = "FP8ActivationResidualQwen2ForCausalLM"
+            coat_fp8_args = QuantizationConfig(**asdict(model_args))
+
+            # Remove the quantization args from llm_cfg and make it a independent config
+            model_args_dict = asdict(model_args)
+            for key in asdict(coat_fp8_args).keys():
+                model_args_dict.pop(key, None)
+
+            llm_cfg.coat_fp8_args = asdict(coat_fp8_args)
+            _attn_implementation = llm_cfg._attn_implementation
+
+            llm_cfg = FP8ActivationResidualQwen2Config(**llm_cfg.to_dict())
+            llm_cfg._attn_implementation = _attn_implementation
+        else:
+            raise ValueError(
+                f"{quantize_model_class} is not supported quantize_model_class."
+            )
+
+        kwargs.pop("quantize_model_class", None)
+
+        if quantize_model_class in [
+            "FP8LinearQwen2ForCausalLM",
+            "FP8ActivationQwen2ForCausalLM",
+            "FP8ActivationResidualQwen2ForCausalLM",
+        ]:  # Remove the quantization args from llm_cfg and make it a independent config
+            llm_cfg.update(model_args_dict)
+        else:
+            llm_cfg.update(asdict(model_args))
+        # print(model_args)
+
+    if quantization_restore_from_checkpoint:
+        fp8_model_name_or_path = kwargs.pop("fp8_llm_cfg", None)
+
+        llm = AutoModelForCausalLM.from_pretrained(
+            fp8_model_name_or_path,
+            config=llm_cfg,
+            torch_dtype=eval(config.model_dtype),
+            *args,
+            **kwargs,
+        )
+
+    else:
+        llm = AutoModelForCausalLM.from_pretrained(
+            model_name_or_path,
+            config=llm_cfg,
+            torch_dtype=eval(config.model_dtype),
+            *args,
+            **kwargs,
+        )
+    packing.patch(llm)
+
+    # Locate the tokenizer.
+    llm_path = model_name_or_path
+    if not has_tokenizer(llm_path):
+        llm_path = osp.join(llm_path, "llm")
+    if not has_tokenizer(llm_path):
+        raise ValueError(f"Cannot find tokenizer in {llm_path}.")
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        llm_path, padding_side="right", use_fast=False, legacy=False
+    )
+    if model_max_length is not None:
+        tokenizer.model_max_length = model_max_length
+
+    # Load chat template if specified.
+    if getattr(config, "chat_template", None) is not None:
+        logger.info(f"Using chat template: {config.chat_template}")
+        fpath = os.path.join(
+            os.path.dirname(__file__), "chat_templates", f"{config.chat_template}.jinja"
+        )
+        with open(fpath) as fd:
+            chat_template = fd.read()
+        tokenizer.chat_template = chat_template.replace("    ", "").replace("\n", "")
+
+    # Set stop tokens for the tokenizer
+    tokenizer.stop_tokens = infer_stop_tokens(tokenizer)
+    tokenizer.stop_token_ids = tokenizer.convert_tokens_to_ids(tokenizer.stop_tokens)
+
+    # Add media tokens to the tokenizer
+    tokenizer.media_tokens = MEDIA_TOKENS
+    tokenizer.media_token_ids = {}
+    for name, token in MEDIA_TOKENS.items():
+        tokenizer.add_tokens([token], special_tokens=True)
+        tokenizer.media_token_ids[name] = tokenizer.convert_tokens_to_ids(token)
+
+    # TODO(ligeng): is this necessary for llava?
+    config.hidden_size = llm.config.hidden_size
+    return llm, tokenizer
+
+
+def build_tokenizer(
+    model_name_or_path: str,
+    config: PretrainedConfig,
+    attn_implementation=None,
+    model_max_length=None,
+    *args,
+    **kwargs,
+) -> Tuple[PreTrainedModel, PreTrainedTokenizer]:
+    # print(model_name_or_path)
+    llm_cfg = AutoConfig.from_pretrained(model_name_or_path)
+    llm_cfg._attn_implementation = attn_implementation
+    llm_cfg.model_max_length = model_max_length
+    if model_max_length is not None:
+        context_length_extension(llm_cfg)
+
+    # Locate the tokenizer.
+    llm_path = model_name_or_path
+    if not has_tokenizer(llm_path):
+        llm_path = osp.join(llm_path, "llm")
+    if not has_tokenizer(llm_path):
+        raise ValueError(f"Cannot find tokenizer in {llm_path}.")
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        llm_path, padding_side="right", use_fast=False, legacy=False
+    )
+    if model_max_length is not None:
+        tokenizer.model_max_length = model_max_length
+
+    # Load chat template if specified.
+    if getattr(config, "chat_template", None) is not None:
+        logger.info(f"Using chat template: {config.chat_template}")
+        fpath = os.path.join(
+            os.path.dirname(__file__), "chat_templates", f"{config.chat_template}.jinja"
+        )
+        with open(fpath) as fd:
+            chat_template = fd.read()
+        tokenizer.chat_template = chat_template.replace("    ", "").replace("\n", "")
+
+    # Set stop tokens for the tokenizer
+    tokenizer.stop_tokens = infer_stop_tokens(tokenizer)
+    tokenizer.stop_token_ids = tokenizer.convert_tokens_to_ids(tokenizer.stop_tokens)
+
+    # Add media tokens to the tokenizer
+    tokenizer.media_tokens = MEDIA_TOKENS
+    tokenizer.media_token_ids = {}
+    for name, token in MEDIA_TOKENS.items():
+        tokenizer.add_tokens([token], special_tokens=True)
+        tokenizer.media_token_ids[name] = tokenizer.convert_tokens_to_ids(token)
+
+    return tokenizer

llm-awq/tinychat/models/nvila/configuration_llava.py

@@ -0,0 +1,89 @@
+# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Optional
+
+from transformers import PretrainedConfig
+
+
+class LlavaConfig(PretrainedConfig):
+    model_type = "llava"
+
+    def __init__(
+        self,
+        llm_cfg=None,
+        vision_tower_cfg=None,
+        mm_projector_cfg=None,
+        architectures=None,
+        resume_path=None,
+        hidden_size=None,
+        mm_hidden_size=None,
+        image_aspect_ratio=None,
+        num_video_frames=None,
+        fps=None,
+        mm_vision_select_layer=None,
+        mm_vision_select_feature=None,
+        mm_use_im_start_end=False,
+        mm_use_im_patch_token=False,
+        mm_projector_lr=None,
+        vision_tower_lr=None,
+        vision_resolution=None,
+        interpolate_mode=None,
+        s2=None,
+        dynamic_s2=None,
+        s2_scales=None,
+        s2_max_split_size=None,
+        s2_resize_output_to_scale_idx=0,
+        min_tiles: Optional[int] = 1,
+        max_tiles: Optional[int] = 12,
+        num_time_tokens=None,
+        time_token_format=None,
+        image_encoder: str = '{"_target_": "llava.model.encoders.BasicImageEncoder"}',
+        video_encoder: str = '{"_target_": "llava.model.encoders.BasicVideoEncoder"}',
+        **kwargs,
+    ):
+        super().__init__()
+        self.architectures = architectures
+        self.llm_cfg = llm_cfg
+        self.vision_tower_cfg = vision_tower_cfg
+        self.mm_projector_cfg = mm_projector_cfg
+        self.resume_path = resume_path
+
+        self.hidden_size = hidden_size
+        self.mm_hidden_size = mm_hidden_size
+        self.image_aspect_ratio = image_aspect_ratio
+        self.num_video_frames = num_video_frames
+        self.fps = fps
+        self.mm_vision_select_layer = mm_vision_select_layer
+        self.mm_vision_select_feature = mm_vision_select_feature
+        self.mm_use_im_start_end = mm_use_im_start_end
+        self.mm_use_im_patch_token = mm_use_im_patch_token
+        self.mm_projector_lr = mm_projector_lr
+        self.vision_tower_lr = vision_tower_lr
+        self.vision_resolution = vision_resolution
+        self.interpolate_mode = interpolate_mode
+        self.s2 = s2
+        self.dynamic_s2 = dynamic_s2
+        self.s2_scales = s2_scales
+        self.s2_max_split_size = s2_max_split_size
+        self.s2_resize_output_to_scale_idx = s2_resize_output_to_scale_idx
+        self.min_tiles = min_tiles
+        self.max_tiles = max_tiles
+        self.num_time_tokens = num_time_tokens
+        self.time_token_format = time_token_format
+
+        self.image_encoder = image_encoder
+        self.video_encoder = video_encoder

llm-awq/tinychat/models/nvila/llava_arch.py

@@ -0,0 +1,909 @@
+#    Copyright 2023 Haotian Liu
+#
+#    Licensed under the Apache License, Version 2.0 (the "License");
+#    you may not use this file except in compliance with the License.
+#    You may obtain a copy of the License at
+#
+#        http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS,
+#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#    See the License for the specific language governing permissions and
+#    limitations under the License.
+
+import copy
+import json
+import logging
+import os
+import os.path as osp
+import warnings
+from abc import ABC
+from collections import OrderedDict, defaultdict, deque
+from itertools import chain
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+from einops import rearrange
+from hydra.utils import instantiate
+from transformers import AutoConfig, GenerationConfig, PreTrainedModel
+from transformers.modeling_utils import ContextManagers, no_init_weights
+from time import time
+from llava.constants import DEFAULT_IMAGE_TOKEN, IGNORE_INDEX
+from llava.mm_utils import process_image, process_images
+from llava.model.configuration_llava import LlavaConfig
+from llava.model.language_model.builder import build_llm_and_tokenizer
+from llava.model.multimodal_encoder.builder import build_vision_tower
+from llava.model.multimodal_projector.builder import build_mm_projector
+from llava.model.utils import get_model_config
+
+# from llava.train.sequence_parallel import get_pg_manager
+from llava.utils import distributed as dist
+from llava.utils.media import extract_media
+from llava.utils.tokenizer import tokenize_conversation
+from .builder import build_tokenizer
+
+
+class LlavaMetaModel(ABC):
+    def init_vlm(self, config, *args, **kwargs):
+        # TODO(ligeng): figure out how from_config and from_pretrained works in HF implementation.
+        if (
+            hasattr(self, "llm")
+            or hasattr(self, "vision_tower")
+            or hasattr(self, "mm_projector")
+        ):
+            # already initialized, skipped
+            return
+
+        model_dtype = getattr(config, "model_dtype", "torch.float16")
+        if not hasattr(config, "model_dtype"):
+            warnings.warn(
+                "model_dtype not found in config, defaulting to torch.float16."
+            )
+            config.model_dtype = model_dtype
+
+        cfgs = get_model_config(config)
+        if len(cfgs) == 3:
+            self.llm_cfg, vision_tower_cfg, mm_projector_cfg = cfgs
+        else:
+            raise ValueError(
+                "`llm_cfg` `mm_projector_cfg` `vision_tower_cfg` not found in the config."
+            )
+        self.tokenizer = build_tokenizer(self.llm_cfg, config, *args, **kwargs)
+        self.vision_tower = build_vision_tower(vision_tower_cfg, config)
+        self.mm_projector = build_mm_projector(mm_projector_cfg, config)
+
+        self.encoders = {}
+        for name in ["image", "video"]:
+            config = getattr(self.config, f"{name}_encoder")
+            if isinstance(config, str):
+                config = json.loads(config)
+            self.encoders[name] = instantiate(config, parent=self)
+
+        self.post_config()
+        self.is_loaded = True
+
+        assert (
+            self.vision_tower is not None or self.mm_projector is not None
+        ), "At least one of the components must be instantiated."
+
+    @classmethod
+    def load_from_config(cls, model_path_or_config, *args, **kwargs):
+        pass
+
+    ## FIXME we will use this function to load model in the future
+    @classmethod
+    def load_pretrained(cls, model_path_or_config, *args, **kwargs):
+        kwargs.pop("config", None)
+
+        if isinstance(model_path_or_config, str):
+            config = AutoConfig.from_pretrained(model_path_or_config)
+        elif isinstance(model_path_or_config, LlavaConfig):
+            config = model_path_or_config
+        else:
+            raise NotImplementedError(
+                f"wrong type, {type(model_path_or_config)} \
+                                      {isinstance(model_path_or_config, LlavaConfig)}"
+            )
+
+        model_dtype = getattr(config, "model_dtype", "torch.float16")
+        if not hasattr(config, "model_dtype"):
+            warnings.warn(
+                "model_dtype not found in config, defaulting to torch.float16."
+            )
+            config.model_dtype = model_dtype
+
+        cfgs = get_model_config(config)
+        if len(cfgs) == 3:
+            llm_cfg, vision_tower_cfg, mm_projector_cfg = cfgs
+        else:
+            raise ValueError(
+                "`llm_cfg` `mm_projector_cfg` `vision_tower_cfg` not found in the config."
+            )
+
+        # print(llm_cfg, vision_tower_cfg, mm_projector_cfg); input("DEBUG load_pretrained")
+        init_context = [
+            no_init_weights(_enable=True),
+        ]
+        # print("Before Init Context")
+        # if hasattr(config, "deepspeed") and "mics" in config.deepspeed:
+        #     print("Using MiCS_Init")
+        #     import deepspeed
+        #     init_context.append(deepspeed.zero.MiCS_Init(config_dict_or_path=config.deepspeed))
+        with ContextManagers(init_context):
+            vlm = cls(config, *args, **kwargs)
+        # print(llm_cfg, vision_tower_cfg, mm_projector_cfg); input("DEBUG load_pretrained finish")
+
+        if (
+            hasattr(vlm, "llm")
+            or hasattr(vlm, "vision_tower")
+            or hasattr(vlm, "mm_projector")
+        ):
+            if vlm.is_loaded:
+                return vlm
+
+        vlm.llm, vlm.tokenizer = build_llm_and_tokenizer(
+            llm_cfg, config, *args, **kwargs
+        )
+        vlm.vision_tower = build_vision_tower(vision_tower_cfg, config)
+        vlm.mm_projector = build_mm_projector(mm_projector_cfg, config)
+
+        self.post_config()
+        self.is_loaded = True
+
+        # FIXME(ligeng, yunhao): llm should never be none here.
+        assert (
+            vlm.llm is not None
+            or vlm.vision_tower is not None
+            or vlm.mm_projector is not None
+        ), "At least one of the components must be instantiated."
+        return vlm
+
+    ## FIXME we will use this function to save the model in the future
+    def save_pretrained(self, output_dir, state_dict=None):
+        if state_dict is None:
+            # other wise fetch from deepspeed
+            # state_dict = accelerator.get_state_dict(is_deepspeed_enabled)
+            state_dict = self.state_dict()
+
+        if getattr(self, "tokenizer", None):
+            self.tokenizer.save_pretrained(osp.join(output_dir, "llm"))
+
+        if self.get_llm():
+            print(f"saving llm to {osp.join(output_dir, 'llm')}")
+            self.llm.config._name_or_path = osp.join(output_dir, "llm")
+            llm_state_dict = OrderedDict(
+                {k.split("llm.")[-1]: v for k, v in state_dict.items() if "llm" in k}
+            )
+            self.llm.save_pretrained(
+                os.path.join(output_dir, "llm"), state_dict=llm_state_dict
+            )
+            self.config.llm_cfg = self.llm.config
+
+        if self.get_vision_tower():
+            print(f"saving vision_tower to {osp.join(output_dir, 'vision_tower')}")
+            self.vision_tower.config._name_or_path = osp.join(
+                output_dir, "vision_tower"
+            )
+            vision_tower_state_dict = OrderedDict(
+                {
+                    k.split("vision_tower.vision_tower.")[-1]: v
+                    for k, v in state_dict.items()
+                    if "vision_tower" in k
+                }
+            )
+            self.vision_tower.vision_tower.save_pretrained(
+                os.path.join(output_dir, "vision_tower"),
+                state_dict=vision_tower_state_dict,
+            )
+            self.vision_tower.image_processor.save_pretrained(
+                os.path.join(output_dir, "vision_tower")
+            )
+            self.config.vision_tower_cfg = self.vision_tower.config
+            if hasattr(self.config.vision_tower_cfg, "auto_map"):
+                if "radio" not in self.get_vision_tower().__class__.__name__.lower():
+                    delattr(self.config.vision_tower_cfg, "auto_map")
+
+        if self.get_mm_projector():
+            print(f"saving mm_projector to {osp.join(output_dir, 'mm_projector')}")
+            self.mm_projector.config._name_or_path = osp.join(
+                output_dir, "mm_projector"
+            )
+            mm_projector_state_dict = OrderedDict(
+                {
+                    k.split("mm_projector.")[-1]: v
+                    for k, v in state_dict.items()
+                    if "mm_projector" in k
+                }
+            )
+            self.mm_projector.save_pretrained(
+                os.path.join(output_dir, "mm_projector"),
+                state_dict=mm_projector_state_dict,
+            )
+            self.config.mm_projector_cfg = self.mm_projector.config
+        ## update and save top-level config
+        self.config._name_or_path = output_dir
+        self.config.architectures = [self.__class__.__name__]
+        self.config.save_pretrained(output_dir)
+
+    def get_llm(self):
+        llm = getattr(self, "llm", None)
+        if type(llm) is list:
+            llm = llm[0]
+        return llm
+
+    def get_lm_head(self):
+        lm_head = getattr(self.get_llm(), "lm_head", None)
+        return lm_head
+
+    def get_vision_tower(self):
+        vision_tower = getattr(self, "vision_tower", None)
+        if type(vision_tower) is list:
+            vision_tower = vision_tower[0]
+        return vision_tower
+
+    def get_mm_projector(self):
+        mm_projector = getattr(self, "mm_projector", None)
+        if type(mm_projector) is list:
+            mm_projector = mm_projector[0]
+        return mm_projector
+
+    def post_config(self):
+
+        if getattr(self.config, "vision_tower_cfg", None) is None:
+            self.config.vision_tower_cfg = self.vision_tower.config
+        if getattr(self.config, "mm_projector_cfg", None) is None:
+            self.config.mm_projector_cfg = self.mm_projector.config
+
+    @staticmethod
+    def merge_chessboard(x, num_split_h, num_split_w):
+        """
+        x: b * n * c or b * h * w * c
+        out: b * c * h * w
+        Assuming x contains num_split**2 sub-squares concatenated along batch dimension, merge the sub-squares back to the original whole square.
+        """
+        B = x.shape[0]
+        if x.dim() == 3:
+            N = x.shape[1]
+            x = rearrange(x, "b (h w) c -> b c h w", h=int(N**0.5), w=int(N**0.5))
+
+        assert B % (num_split_h * num_split_w) == 0
+        b = B // (num_split_h * num_split_w)
+
+        x_merge = torch.cat(
+            [
+                torch.cat(
+                    [
+                        x[(i * num_split_w + j) * b : (i * num_split_w + j + 1) * b]
+                        for j in range(num_split_w)
+                    ],
+                    dim=-1,
+                )
+                for i in range(num_split_h)
+            ],
+            dim=-2,
+        )
+
+        return x_merge
+
+    @staticmethod
+    def split_chessboard(x, num_split_h, num_split_w):
+        """
+        x: b * c * h * w
+        out: b * c * h * w
+        Deividing x into num_split**2 sub-squares, and concatenate all the sub-squares on the batch dimension
+        """
+        B, C, H, W = x.shape
+        assert H % num_split_h == 0 and W % num_split_w == 0
+        h, w = H // num_split_h, W // num_split_w
+        x_split = torch.cat(
+            [
+                x[:, :, i * h : (i + 1) * h, j * w : (j + 1) * w]
+                for i in range(num_split_h)
+                for j in range(num_split_w)
+            ],
+            dim=0,
+        )
+        return x_split
+
+    def merge_features_for_dynamic_s2(self, image_features, block_sizes):
+        scales = self.get_vision_tower().scales
+        resize_output_to_scale_idx = self.get_vision_tower().resize_output_to_scale_idx
+
+        image_features_each_image = []
+        new_block_sizes = []
+        block_cnt = 0
+        for block_size_each_image in block_sizes:
+            if block_size_each_image is None:
+                cur_features = image_features[block_cnt : block_cnt + 1]
+                cur_features = rearrange(
+                    cur_features,
+                    "1 (h w) c -> 1 c h w",
+                    h=int(cur_features.shape[1] ** 0.5),
+                )
+                cur_features = cur_features.repeat(1, len(scales), 1, 1)
+                image_features_each_image.append(cur_features)
+                new_block_sizes.append((1, 1))
+                block_cnt += 1
+            else:
+                cur_features_each_scale = []
+                for scale in scales[:-1]:
+                    num_blocks_this_scale = (scale // scales[0]) ** 2
+                    cur_features_each_scale.append(
+                        self.merge_chessboard(
+                            image_features[
+                                block_cnt : block_cnt + num_blocks_this_scale
+                            ],
+                            num_split_h=scale // scales[0],
+                            num_split_w=scale // scales[0],
+                        )
+                    )  # 1 * C * H * W
+                    block_cnt += num_blocks_this_scale
+                num_blocks_last_scale = (
+                    block_size_each_image[0] * block_size_each_image[1]
+                )
+                cur_features_each_scale.append(
+                    self.merge_chessboard(
+                        image_features[block_cnt : block_cnt + num_blocks_last_scale],
+                        num_split_h=block_size_each_image[0],
+                        num_split_w=block_size_each_image[1],
+                    )
+                )  # 1 * C * H * W
+                block_cnt += num_blocks_last_scale
+
+                # resize and concat features from different scales
+                output_size = cur_features_each_scale[resize_output_to_scale_idx].shape[
+                    -2:
+                ]
+                cur_features = torch.cat(
+                    [
+                        F.interpolate(
+                            cur_features_each_scale[i].to(torch.float32),
+                            size=output_size,
+                            mode="area",
+                        ).to(cur_features_each_scale[i].dtype)
+                        for i in range(len(cur_features_each_scale))
+                    ],
+                    dim=1,
+                )
+                # cur_features = rearrange(cur_features, "1 c h w -> (h w) c")
+
+                image_features_each_image.append(cur_features)
+
+                if (
+                    resize_output_to_scale_idx == len(scales) - 1
+                    or resize_output_to_scale_idx == -1
+                ):
+                    new_block_sizes.append(block_size_each_image)
+                else:
+                    new_block_sizes.append(
+                        (
+                            scales[resize_output_to_scale_idx] // scales[0],
+                            scales[resize_output_to_scale_idx] // scales[0],
+                        )
+                    )
+
+        assert block_cnt == len(image_features)
+
+        return image_features_each_image, new_block_sizes
+
+    def encode_images(
+        self, images, block_sizes: Optional[Optional[Tuple[int, ...]]] = None
+    ):
+        if block_sizes is None:
+            block_sizes = [None] * len(images)
+        if getattr(self.config, "dynamic_s2", False):
+            image_features = self.get_vision_tower()(images)
+            image_features, new_block_sizes = self.merge_features_for_dynamic_s2(
+                image_features, block_sizes
+            )
+
+            image_features = [
+                self.split_chessboard(x, block_size[0], block_size[1])
+                for x, block_size in zip(image_features, new_block_sizes)
+            ]  # list of B * C * H * W tensors
+            image_features = torch.cat(
+                [rearrange(x, "b c h w -> b (h w) c") for x in image_features], dim=0
+            )  # B * N * C
+            image_features = self.get_mm_projector()(image_features)
+            image_features = list(
+                image_features.split(
+                    [block_size[0] * block_size[1] for block_size in new_block_sizes],
+                    dim=0,
+                )
+            )
+            image_features = [
+                self.merge_chessboard(x, block_size[0], block_size[1])
+                for x, block_size in zip(image_features, new_block_sizes)
+            ]  # list of 1 * C * H * W tensors
+            image_features = [
+                rearrange(x, "1 c h w -> (h w) c") for x in image_features
+            ]  # list of N * C tensors
+            image_features = torch.stack(image_features, dim=0)
+        else:
+            image_features = self.get_vision_tower()(images)
+            image_features = self.get_mm_projector()(image_features)
+        return image_features
+
+    ## @yunhao: is there a better way to handle function call and attributes for llm?
+    ## support beam search
+    def _temporary_reorder_cache(self, past_key_values, sorted_idx):
+        return self.get_llm()._temporary_reorder_cache(past_key_values, sorted_idx)
+
+    def get_input_embeddings(self):
+        return self.get_llm().get_input_embeddings()
+
+    def get_output_embeddings(self):
+        return self.get_llm().get_output_embeddings()
+
+    def resize_token_embeddings(self, embed_size):
+        self.get_llm().resize_token_embeddings(embed_size)
+
+
+class LlavaMetaForCausalLM(ABC):
+    def _embed(
+        self,
+        input_ids: torch.Tensor,
+        media: Dict[str, List[torch.Tensor]],
+        media_config: Dict[str, Dict[str, Any]],
+        labels: Optional[torch.Tensor],
+        attention_mask: Optional[torch.Tensor],
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        labels = (
+            labels if labels is not None else torch.full_like(input_ids, IGNORE_INDEX)
+        )
+        attention_mask = (
+            attention_mask
+            if attention_mask is not None
+            else torch.ones_like(input_ids, dtype=torch.bool)
+        )
+
+        # Extract text and media embeddings
+        text_embeds = self.llm.model.embed_tokens(input_ids)
+        media_embeds = self.__embed_media_tokens(media, media_config)
+
+        # This is a workaround to make sure the dummy embeddings are consumed
+        while media_embeds.get("dummy"):
+            dummy_embed = media_embeds["dummy"].popleft()
+            text_embeds += torch.sum(dummy_embed) * 0
+        # Remove padding
+        batch_size = labels.shape[0]
+        text_embeds = [text_embeds[k][attention_mask[k]] for k in range(batch_size)]
+        labels = [labels[k][attention_mask[k]] for k in range(batch_size)]
+
+        # Build inverse mapping from token ID to media name
+        media_tokens = {}
+        for name, token_id in self.tokenizer.media_token_ids.items():
+            media_tokens[token_id] = name
+
+        # Fuse text and media embeddings
+        inputs_m, labels_m = [], []
+        for k in range(batch_size):
+            inputs_mk, labels_mk = [], []
+            pos = 0
+            while pos < len(labels[k]):
+                if input_ids[k][pos].item() in media_tokens:
+                    end = pos + 1
+                    name = media_tokens[input_ids[k][pos].item()]
+                    input = media_embeds[name].popleft()
+                    label = torch.full(
+                        [input.shape[0]],
+                        IGNORE_INDEX,
+                        device=labels[k].device,
+                        dtype=labels[k].dtype,
+                    )
+                else:
+                    end = pos
+                    while (
+                        end < len(labels[k])
+                        and input_ids[k][end].item() not in media_tokens
+                    ):
+                        end += 1
+                    input = text_embeds[k][pos:end]
+                    label = labels[k][pos:end]
+                inputs_mk.append(input)
+                labels_mk.append(label)
+                pos = end
+            inputs_m.append(torch.cat(inputs_mk, dim=0))
+            labels_m.append(torch.cat(labels_mk, dim=0))
+        inputs, labels = inputs_m, labels_m
+
+        # Check if all media embeddings are consumed
+        for name in media_embeds:
+            if media_embeds[name]:
+                raise ValueError(f"Not all {name} embeddings are consumed!")
+
+        # Truncate sequences to `model_max_length` as media embeddings are inserted
+        inputs, labels = self.__truncate_sequence(inputs, labels)
+
+        # Pad sequences to the longest one in the batch
+        return self.__batchify_sequence(inputs, labels)
+
+    def __embed_media_tokens(
+        self,
+        media: Dict[str, List[torch.Tensor]],
+        media_config: Dict[str, Dict[str, Any]],
+    ) -> Dict[str, List[torch.Tensor]]:
+        embeds = defaultdict(deque)
+        for name in media:
+            embeds[name] = deque(self.encoders[name](media[name], media_config[name]))
+        return embeds
+
+    def __truncate_sequence(
+        self, inputs: List[torch.Tensor], labels: List[torch.Tensor]
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        if any(len(input) > self.tokenizer.model_max_length for input in inputs):
+            warnings.warn(
+                f"Truncating sequences to `model_max_length` ({self.tokenizer.model_max_length})."
+            )
+            inputs = [input[: self.tokenizer.model_max_length] for input in inputs]
+            labels = [label[: self.tokenizer.model_max_length] for label in labels]
+        return inputs, labels
+
+    def __batchify_sequence(
+        self, inputs: List[torch.Tensor], labels: List[torch.Tensor]
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        batch_size = len(inputs)
+        device = inputs[0].device
+        hidden_size = inputs[0].shape[1]
+        max_length = max(inputs[k].shape[0] for k in range(batch_size))
+        attention_mask = torch.ones(
+            (batch_size, max_length), dtype=torch.bool, device=device
+        )
+
+        inputs_p, labels_p = [], []
+        for k in range(batch_size):
+            size_pk = max_length - inputs[k].shape[0]
+            inputs_pk = torch.zeros(
+                (size_pk, hidden_size), dtype=inputs[k].dtype, device=device
+            )
+            labels_pk = torch.full(
+                (size_pk,), IGNORE_INDEX, dtype=labels[k].dtype, device=device
+            )
+            if self.tokenizer.padding_side == "right":
+                attention_mask[k, inputs[k].shape[0] :] = False
+                inputs_pk = torch.cat([inputs[k], inputs_pk], dim=0)
+                labels_pk = torch.cat([labels[k], labels_pk], dim=0)
+            else:
+                attention_mask[k, : -inputs[k].shape[0]] = False
+                inputs_pk = torch.cat([inputs_pk, inputs[k]], dim=0)
+                labels_pk = torch.cat([labels_pk, labels[k]], dim=0)
+            inputs_p.append(inputs_pk)
+            labels_p.append(labels_pk)
+
+        inputs = torch.stack(inputs_p, dim=0)
+        labels = torch.stack(labels_p, dim=0)
+        return inputs, labels, attention_mask
+
+    @torch.inference_mode()
+    def generate(
+        self,
+        input_ids: Optional[torch.FloatTensor] = None,
+        media: Optional[Dict[str, List[torch.Tensor]]] = None,
+        media_config: Dict[str, Dict[str, Any]] = None,
+        attention_mask: Optional[torch.LongTensor] = None,
+        quant_llm: Optional[bool] = True,
+        **generation_kwargs,
+    ):
+        inputs_embeds, _, attention_mask = self._embed(
+            input_ids, media, media_config, None, attention_mask
+        )
+        return self.llm.generate(
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            quant_llm=quant_llm,
+            **generation_kwargs,
+        )
+
+    @torch.inference_mode()
+    def generate_content(
+        self,
+        prompt: Union[str, List],
+        generation_config: Optional[GenerationConfig] = None,
+        quant_llm: Optional[bool] = True,
+    ) -> str:
+        # TODO(zhijianl): Support directly taking conversation as input
+        conversation = [{"from": "human", "value": prompt}]
+
+        # Extract media from the conversation
+
+        # TODO (extract and preprocess should be done together, as the preprocess of image and video can be different, i.e. when dynamic res is used)
+        media = extract_media(conversation, self.config)
+
+        # Process media
+        media_config = defaultdict(dict)
+        for name in media:
+            if name == "image":
+                if len(media["image"]) == 1 and self.config.image_aspect_ratio in [
+                    "dynamic",
+                    "dynamic_s2",
+                ]:
+                    self.config.image_processor = self.vision_tower.image_processor
+                    if self.config.image_aspect_ratio == "dynamic":
+                        images = process_image(
+                            media["image"][0],
+                            self.config,
+                            None,
+                            enable_dynamic_res=True,
+                        ).half()
+                        conversation[0]["value"] = conversation[0]["value"].replace(
+                            DEFAULT_IMAGE_TOKEN,
+                            f"{DEFAULT_IMAGE_TOKEN}\n" * images.shape[0],
+                        )
+                    else:
+                        if type(self.config.s2_scales) is str:
+                            self.config.s2_scales = list(
+                                map(int, self.config.s2_scales.split(","))
+                            )
+                        images, block_sizes = process_image(
+                            media["image"][0], self.config, None, enable_dynamic_s2=True
+                        )
+                        images = images.half()
+                        media_config[name]["block_sizes"] = [block_sizes]
+                else:
+                    images = process_images(
+                        media["image"], self.vision_tower.image_processor, self.config
+                    ).half()
+                media[name] = [image for image in images]
+            elif name == "video":
+                media[name] = [
+                    process_images(
+                        images, self.vision_tower.image_processor, self.config
+                    ).half()
+                    for images in media[name]
+                ]
+            else:
+                raise ValueError(f"Unsupported media type: {name}")
+
+        # Tokenize the conversation
+        input_ids = (
+            tokenize_conversation(
+                conversation, self.tokenizer, add_generation_prompt=True
+            )
+            .cuda()
+            .unsqueeze(0)
+        )
+
+        # Set up the generation config
+        generation_config = generation_config or self.default_generation_config
+        # Generate the response
+        try:
+            output_ids = self.generate(
+                input_ids=input_ids,
+                media=media,
+                media_config=media_config,
+                generation_config=generation_config,
+                quant_llm=quant_llm,
+            )
+        except ValueError:
+            if not generation_config.do_sample:
+                raise
+            # FIXME(zhijianl): This is a temporary workaround for the sampling issue
+            logging.warning(
+                "Generation failed with sampling, retrying with greedy decoding."
+            )
+            generation_config.do_sample = False
+            output_ids = self.generate(
+                input_ids=input_ids,
+                media=media,
+                media_config=media_config,
+                generation_config=generation_config,
+            )
+
+        # Decode the response
+        response = self.tokenizer.decode(
+            output_ids[0], skip_special_tokens=True
+        ).strip()
+        return response
+
+    @torch.inference_mode()
+    def benchmark(self, prompt: Union[str, List], quant_llm) -> None:
+        # TODO(zhijianl): Support directly taking conversation as input
+        conversation = [{"from": "human", "value": prompt}]
+
+        # Extract media from the conversation
+
+        # TODO (extract and preprocess should be done together, as the preprocess of image and video can be different, i.e. when dynamic res is used)
+        media = extract_media(conversation, self.config)
+
+        # Process media
+        media_config = defaultdict(dict)
+        image_num = 0
+        for name in media:
+            if name == "image":
+                if len(media["image"]) == 1 and self.config.image_aspect_ratio in [
+                    "dynamic",
+                    "dynamic_s2",
+                ]:
+                    self.config.image_processor = self.vision_tower.image_processor
+                    if self.config.image_aspect_ratio == "dynamic":
+                        images = process_image(
+                            media["image"][0],
+                            self.config,
+                            None,
+                            enable_dynamic_res=True,
+                        ).half()
+                        if len(images.shape) == 3:
+                            images = images.reshape(1, *images.shape)
+                        image_num += images.shape[0]
+                        size = images.shape[1:]
+                        conversation[0]["value"] = conversation[0]["value"].replace(
+                            DEFAULT_IMAGE_TOKEN,
+                            f"{DEFAULT_IMAGE_TOKEN}\n" * images.shape[0],
+                        )
+                    else:
+                        if type(self.config.s2_scales) is str:
+                            self.config.s2_scales = list(
+                                map(int, self.config.s2_scales.split(","))
+                            )
+                        images, block_sizes = process_image(
+                            media["image"][0], self.config, None, enable_dynamic_s2=True
+                        )
+                        images = images.half()
+                        if len(images.shape) == 3:
+                            images = images.reshape(1, *images.shape)
+                        image_num += images.shape[0]
+                        size = images.shape[1:]
+                        media_config[name]["block_sizes"] = [block_sizes]
+                else:
+                    images = process_images(
+                        media["image"], self.vision_tower.image_processor, self.config
+                    ).half()
+                    image_num += images.shape[0]
+                    size = images.shape[1:]
+                media[name] = [image for image in images]
+            elif name == "video":
+                media[name] = [
+                    process_images(
+                        images, self.vision_tower.image_processor, self.config
+                    ).half()
+                    for images in media[name]
+                ]
+                for images in media[name]:
+                    image_num += images.shape[0]
+                    size = images.shape[1:]
+            else:
+                raise ValueError(f"Unsupported media type: {name}")
+
+        # Tokenize the conversation
+        input_ids = (
+            tokenize_conversation(
+                conversation, self.tokenizer, add_generation_prompt=True
+            )
+            .cuda()
+            .unsqueeze(0)
+        )
+
+        # Set up the generation config
+        for i in range(10):
+            torch.cuda.synchronize()
+            t_st = time()
+            inputs_embeds, _, attention_mask = self._embed(
+                input_ids, media, media_config, None, None
+            )
+            torch.cuda.synchronize()
+            t_ed = time()
+            torch.cuda.empty_cache()
+        print(
+            "Time of vision tower and others is {:.5f} s for {} images ({} x {} x {})".format(
+                t_ed - t_st, image_num, size[0], size[1], size[2]
+            )
+        )
+        output = self.llm.benchmark(
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            quant_llm=quant_llm,
+        )
+        # response = self.tokenizer.decode(output, skip_special_tokens=True).strip()
+        return
+
+    @property
+    def default_generation_config(self) -> GenerationConfig:
+        generation_config = copy.deepcopy(self.generation_config or GenerationConfig())
+        if self.tokenizer.eos_token_id is None:
+            raise ValueError("Tokenizer must have an EOS token")
+        if generation_config.max_length == GenerationConfig().max_length:
+            generation_config.max_length = self.tokenizer.model_max_length
+        if generation_config.pad_token_id is None:
+            generation_config.pad_token_id = (
+                self.tokenizer.pad_token_id or self.tokenizer.eos_token_id
+            )
+        if generation_config.bos_token_id is None:
+            generation_config.bos_token_id = (
+                self.tokenizer.bos_token_id or self.tokenizer.eos_token_id
+            )
+        if generation_config.eos_token_id is None:
+            generation_config.eos_token_id = self.tokenizer.stop_token_ids
+        return generation_config
+
+    # Prepare media
+
+    # Process media
+    @torch.inference_mode()
+    def prepare_media(self, conversation):
+        media = extract_media(conversation, self.config)
+
+        # Process media
+        media_config = defaultdict(dict)
+        for name in media:
+            if name == "image":
+                if len(media["image"]) == 1 and self.config.image_aspect_ratio in [
+                    "dynamic",
+                    "dynamic_s2",
+                ]:
+                    self.config.image_processor = self.vision_tower.image_processor
+                    if self.config.image_aspect_ratio == "dynamic":
+                        images = process_image(
+                            media["image"][0],
+                            self.config,
+                            None,
+                            enable_dynamic_res=True,
+                        ).half()
+                        conversation[0]["value"] = conversation[0]["value"].replace(
+                            DEFAULT_IMAGE_TOKEN,
+                            f"{DEFAULT_IMAGE_TOKEN}\n" * images.shape[0],
+                        )
+                    else:
+                        if type(self.config.s2_scales) is str:
+                            self.config.s2_scales = list(
+                                map(int, self.config.s2_scales.split(","))
+                            )
+                        images, block_sizes = process_image(
+                            media["image"][0], self.config, None, enable_dynamic_s2=True
+                        )
+                        images = images.half()
+                        media_config[name]["block_sizes"] = [block_sizes]
+                else:
+                    images = process_images(
+                        media["image"], self.vision_tower.image_processor, self.config
+                    ).half()
+                media[name] = [image for image in images]
+            elif name == "video":
+                media[name] = [
+                    process_images(
+                        images, self.vision_tower.image_processor, self.config
+                    ).half()
+                    for images in media[name]
+                ]
+            else:
+                raise ValueError(f"Unsupported media type: {name}")
+        return media, media_config
+
+    @torch.inference_mode()
+    def stream_gen(
+        self,
+        input_ids,
+        media,
+        media_cfg,
+        start_pos,
+        chunk_prefilling,
+        quant_llm,
+        attention_mask=None,
+    ) -> str:
+        if media is None:
+            inputs_embeds = self.llm.model.embed_tokens(input_ids)
+        else:
+            image_num = torch.sum(input_ids == 151649)
+            if image_num == 1 and self.config.image_aspect_ratio == "dynamic":
+                patch_num = len(media["image"])
+                new_input_ids = []
+                for i, id in enumerate(input_ids[0]):
+                    if id == 151649:
+                        new_input_ids.extend(input_ids[0, 0:i])
+                        new_input_ids.extend([198, 151649, 198] * patch_num)
+                        new_input_ids.extend(input_ids[0, i + 1 :])
+                        break
+                input_ids = torch.tensor(
+                    [new_input_ids], dtype=torch.int, device="cuda"
+                )
+            inputs_embeds, _, _ = self._embed(
+                input_ids, media, media_cfg, None, attention_mask=None
+            )
+        length = inputs_embeds.shape[1]
+        if quant_llm:
+            out = self.llm(None, start_pos, inputs_embeds, chunk_prefilling)
+        else:
+            out = self.llm.forwardfp16(None, start_pos, inputs_embeds, chunk_prefilling)
+        return out, length

llm-awq/tinychat/models/nvila_qwen2.py

@@ -0,0 +1,157 @@
+#    Copyright 2023 Haotian Liu
+#
+#    Licensed under the Apache License, Version 2.0 (the "License");
+#    you may not use this file except in compliance with the License.
+#    You may obtain a copy of the License at
+#
+#        http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS,
+#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#    See the License for the specific language governing permissions and
+#    limitations under the License.
+
+# This file is modified from https://github.com/haotian-liu/LLaVA/
+
+
+import os
+from collections import defaultdict
+from typing import Dict, List, Optional, Tuple, Union
+
+import torch
+from transformers import AutoConfig, AutoModel, PretrainedConfig, PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+
+from .nvila.configuration_llava import LlavaConfig
+from .nvila.llava_arch import LlavaMetaForCausalLM, LlavaMetaModel
+from .qwen2 import Qwen2ForCausalLM
+
+
+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
+from transformers import modeling_utils
+
+modeling_utils._init_weights = False
+
+
+class LlavaLlamaConfig(LlavaConfig):
+    model_type = "llava_llama"
+
+
+class NVILAQwen2(LlavaMetaModel, LlavaMetaForCausalLM, PreTrainedModel):
+    config_class = LlavaLlamaConfig
+    main_input_name = "input_embeds"
+    supports_gradient_checkpointing = True
+    _supports_flash_attn_2 = True
+
+    def __init__(
+        self, config: LlavaLlamaConfig = None, llm=True, *args, **kwargs
+    ) -> None:
+        super().__init__(config)
+        self.init_vlm(config=config, *args, **kwargs)
+        # TODO: Skip the weight loading to save time
+        self.llm_cfg = AutoConfig.from_pretrained(self.llm_cfg, init_weights=False)
+        if llm:
+            self.llm = Qwen2ForCausalLM.from_pretrained(self.llm_cfg._name_or_path)
+            self.llm = self.llm.cpu()
+            self.llm.resize_token_embeddings(len(self.tokenizer))
+        else:
+            self.llm = None
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_name_or_path: Optional[Union[str, os.PathLike]],
+        *model_args,
+        config: Optional[Union[PretrainedConfig, str, os.PathLike]] = None,
+        cache_dir: Optional[Union[str, os.PathLike]] = None,
+        ignore_mismatched_sizes: bool = False,
+        force_download: bool = False,
+        local_files_only: bool = False,
+        token: Optional[Union[str, bool]] = None,
+        revision: str = "main",
+        use_safetensors: bool = None,
+        **kwargs,
+    ):
+        if hasattr(cls, "load_pretrained"):
+            return cls.load_pretrained(
+                pretrained_model_name_or_path,
+                *model_args,
+                config=config,
+                cache_dir=cache_dir,
+                ignore_mismatched_sizes=ignore_mismatched_sizes,
+                force_download=force_download,
+                local_files_only=local_files_only,
+                token=token,
+                revision=revision,
+                use_safetensors=use_safetensors,
+                **kwargs,
+            )
+        return super(NVILAQwen2).from_pretrained(
+            pretrained_model_name_or_path,
+            *model_args,
+            config=config,
+            cache_dir=cache_dir,
+            ignore_mismatched_sizes=ignore_mismatched_sizes,
+            force_download=force_download,
+            local_files_only=local_files_only,
+            token=token,
+            revision=revision,
+            use_safetensors=use_safetensors,
+            **kwargs,
+        )
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        media: Optional[Dict[str, List[torch.Tensor]]] = None,
+        images: Optional[torch.FloatTensor] = None,
+        media_config: Optional[List] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        packing: bool = True,
+        seqlens_in_batch: Optional[torch.LongTensor] = None,
+        dpo_forward: bool = False,
+        **kwargs,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        self.freezed_module_patch()
+
+        if images is not None:
+            if media is not None:
+                raise ValueError(
+                    "Both 'media' and 'images' are provided. Please provide only one."
+                )
+            media = {"image": images}
+
+        if media_config is None:
+            media_config = defaultdict(dict)
+
+        if inputs_embeds is None:
+            inputs_embeds, labels, attention_mask = self._embed(
+                input_ids, media, media_config, labels, attention_mask
+            )
+
+        outputs = self.llm(
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            labels=labels,
+            **kwargs,
+        )
+
+        if dpo_forward:
+            return outputs.logits, labels
+
+        return outputs

llm-awq/tinychat/models/qwen2.py

@@ -0,0 +1,511 @@
+# Modified from https://github.com/huggingface/transformers/blob/main/src/transformers/models/qwen2/modeling_qwen2.py
+"""PyTorch Qwen2 model."""
+
+import math
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.utils.checkpoint
+from torch import nn
+import awq_inference_engine
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+from transformers.models.qwen2.configuration_qwen2 import Qwen2Config
+from transformers.activations import ACT2FN
+import tinychat
+import torch.nn.functional as F
+import time
+from tqdm import tqdm
+from transformers import GenerationMixin
+from transformers.models.qwen2 import Qwen2ForCausalLM
+from flash_attn import flash_attn_func
+
+max_batch_size = tinychat.utils.constants.max_batch_size
+max_seq_len = tinychat.utils.constants.max_seq_len
+
+
+class Qwen2RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def _norm(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        output = torch.empty_like(x)
+        awq_inference_engine.layernorm_forward_cuda(x, self.weight, output, self.eps)
+        return output
+
+
+def precompute_freqs_cis(
+    dim: int, end: int, theta: float = 10000.0, scale: float = 1.0
+):
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
+    t = torch.arange(end, device=freqs.device)  # type: ignore
+    freqs = torch.outer(t * scale, freqs).float()  # type: ignore
+
+    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # complex64
+    return freqs_cis
+
+
+def precompute_freqs(
+    dim: int, end: int, theta: float = 10000.0, scale: float = 1.0, device=None
+):
+    inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2).float().to(device) / dim))
+    seq = torch.arange(end, dtype=inv_freq.dtype, device=device)
+    freqs = torch.einsum("i , j -> i j", seq, inv_freq)
+    freqs = freqs.reshape(freqs.shape[0], 1, 1, -1)
+    return torch.cat((freqs, freqs), dim=-1)
+
+
+def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
+    ndim = x.ndim
+    assert 0 <= 1 < ndim
+    assert freqs_cis.shape == (x.shape[1], x.shape[-1])
+    shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
+    return freqs_cis.view(*shape)
+
+
+def apply_rotary_emb(
+    xq: torch.Tensor,
+    xk: torch.Tensor,
+    freqs_cis: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    # xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
+    # k_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
+    xq_ = torch.view_as_complex(
+        xq.float().reshape(*xq.shape[:-1], 2, -1).transpose(-2, -1).contiguous()
+    )
+    xk_ = torch.view_as_complex(
+        xk.float().reshape(*xk.shape[:-1], 2, -1).transpose(-2, -1).contiguous()
+    )
+    freqs_cis = reshape_for_broadcast(freqs_cis, xq_)
+    xq_out = torch.view_as_real(xq_ * freqs_cis).transpose(-2, -1).flatten(3)
+    xk_out = torch.view_as_real(xk_ * freqs_cis).transpose(-2, -1).flatten(3)
+    return xq_out.type_as(xq), xk_out.type_as(xk)
+
+
+class Qwen2MLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, hidden_state):
+        return self.down_proj(
+            self.act_fn(self.gate_proj(hidden_state)) * self.up_proj(hidden_state)
+        )
+
+
+# Copied from transformers.models.llama.modeling_llama.repeat_kv
+def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    """
+    batch, num_key_value_heads, slen, head_dim = x.shape
+    if n_rep == 1:
+        return x
+    x = x[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+    return x.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+class Qwen2AttentionFused(nn.Module):
+    """
+    Multi-headed attention from 'Attention Is All You Need' paper. Modified to use sliding window attention: Longformer
+    and "Generating Long Sequences with Sparse Transformers".
+    """
+
+    def __init__(self, config: Qwen2Config, layer_idx: Optional[int] = None):
+        super().__init__()
+        self.args = config
+        self.layer_idx = layer_idx
+        if layer_idx is None:
+            print(
+                f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will "
+                "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "
+                "when creating this class."
+            )
+
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = self.hidden_size // self.num_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.max_position_embeddings = config.max_position_embeddings
+        self.rope_theta = config.rope_theta
+        self.is_causal = True
+        self.attention_dropout = config.attention_dropout
+        self.rope_scaling = config.rope_scaling
+        if self.rope_scaling is None:
+            self.rope_scaling = 1.0
+        elif isinstance(self.rope_scaling, dict):
+            self.rope_scaling = self.rope_scaling.get("factor", 1.0)
+
+        if (self.head_dim * self.num_heads) != self.hidden_size:
+            raise ValueError(
+                f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
+                f" and `num_heads`: {self.num_heads})."
+            )
+        self.q_proj = nn.Linear(
+            self.hidden_size, self.num_heads * self.head_dim, bias=True
+        )
+        self.k_proj = nn.Linear(
+            self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True
+        )
+        self.v_proj = nn.Linear(
+            self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True
+        )
+        self.o_proj = nn.Linear(
+            self.num_heads * self.head_dim, self.hidden_size, bias=False
+        )
+        self.kv_max_seq_len = min(max_seq_len, self.max_position_embeddings)
+        # following fastertransformer definition
+        self.cache_v = (
+            torch.zeros(
+                (
+                    max_batch_size,
+                    self.num_key_value_heads,
+                    # args.max_position_embeddings,
+                    self.kv_max_seq_len,
+                    self.head_dim,
+                )
+            )
+            .cuda()
+            .half()
+        )  # added to half
+        # 8: pack 8 fp16 in FT, if fp32 then use 4
+        self.cache_k = (
+            torch.zeros(
+                (
+                    max_batch_size,
+                    self.num_key_value_heads,
+                    self.head_dim // 8,
+                    # args.max_position_embeddings,
+                    self.kv_max_seq_len,
+                    8,
+                )
+            )
+            .cuda()
+            .half()
+        )  # added to half
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        start_pos: int,
+        freqs: torch.Tensor,
+        mask: Optional[torch.Tensor],
+        chunk_prefilling: bool = False,
+    ):
+        bsz, seqlen, _ = x.shape
+
+        query_states = self.q_proj(x)
+        key_states = self.k_proj(x)
+        value_states = self.v_proj(x)
+
+        if seqlen > 1:
+            xq = query_states.view(bsz, seqlen, self.num_heads, self.head_dim)
+            xk = key_states.view(bsz, seqlen, self.num_key_value_heads, self.head_dim)
+            xv = value_states.view(bsz, seqlen, self.num_key_value_heads, self.head_dim)
+
+            xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs)
+
+            self.cache_k = self.cache_k.to(xq)
+            self.cache_v = self.cache_v.to(xq)
+
+            values_store = xv.transpose(2, 1)
+
+            keys_store = (
+                xk.reshape(bsz, seqlen, self.num_key_value_heads, self.head_dim // 8, 8)
+                .permute(0, 2, 3, 1, 4)
+                .contiguous()
+            )
+
+            self.cache_v[:bsz, :, start_pos : start_pos + seqlen, :] = values_store
+            self.cache_k[:bsz, :, :, start_pos : start_pos + seqlen, :] = keys_store
+            if chunk_prefilling:
+                keys = self.cache_k[:, :, :, 0 : start_pos + seqlen, :]
+                keys = (
+                    keys.permute(0, 3, 1, 2, 4)
+                    .reshape(
+                        bsz, start_pos + seqlen, self.num_key_value_heads, self.head_dim
+                    )
+                    .contiguous()
+                )
+                values = self.cache_v[:, :, 0 : start_pos + seqlen, :]
+                values = (
+                    values.transpose(2, 1)
+                    .reshape(
+                        bsz, start_pos + seqlen, self.num_key_value_heads, self.head_dim
+                    )
+                    .contiguous()
+                )
+            else:
+                keys = xk
+                values = xv
+            output = flash_attn_func(
+                q=xq,
+                k=keys,
+                v=values,
+                causal=True,
+            )
+            output = output.contiguous().view(bsz, seqlen, -1)
+        else:
+            xq = query_states.view(bsz, self.num_heads, self.head_dim)
+            xk = key_states.view(bsz, self.num_key_value_heads, self.head_dim)
+            xv = value_states.view(bsz, self.num_key_value_heads, self.head_dim)
+            
+            output = awq_inference_engine.single_query_attention(
+                xq,
+                xk,
+                xv,
+                self.cache_k,
+                self.cache_v,
+                None,
+                # alibi position encodings
+                None,
+                start_pos,
+                self.head_dim,
+                self.rope_theta,
+                self.rope_scaling,
+                True,
+            )
+            output = output.reshape(bsz, 1, -1)
+
+        return self.o_proj(output)
+
+
+class Qwen2DecoderLayer(nn.Module):
+    def __init__(self, config: Qwen2Config, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+
+        self.self_attn = Qwen2AttentionFused(config, layer_idx)
+
+        self.mlp = Qwen2MLP(config)
+        self.input_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = Qwen2RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        start_pos: int,
+        freqs: torch.Tensor,
+        mask: Optional[torch.Tensor],
+        chunk_prefilling: bool = False,
+    ):
+        residual = x
+        x = self.input_layernorm(x)
+
+        # Self Attention
+        x = self.self_attn(
+            x=x,
+            start_pos=start_pos,
+            freqs=freqs,
+            mask=mask,
+            chunk_prefilling=chunk_prefilling,
+        )
+        x = residual + x
+
+        # Fully Connected
+        residual = x
+        x = self.post_attention_layernorm(x)
+        x = self.mlp(x)
+        x = residual + x
+        return x
+
+
+class Qwen2Model(nn.Module):
+    def __init__(self, config: Qwen2Config):
+        super().__init__()
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = nn.Embedding(
+            config.vocab_size, config.hidden_size, self.padding_idx
+        )
+        self.layers = nn.ModuleList(
+            [
+                Qwen2DecoderLayer(config, layer_idx)
+                for layer_idx in range(config.num_hidden_layers)
+            ]
+        )
+        self.norm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        # Note (Haotian): rope_theta has to be defined here, otherwise context stage is wrong.
+        rope_scale = config.rope_scaling
+        if rope_scale is None:
+            rope_scale = 1.0
+        else:
+            rope_scale = 1.0 / rope_scale["factor"]
+        self.freqs = precompute_freqs(
+            config.hidden_size // config.num_attention_heads,
+            config.max_position_embeddings * 2,
+            config.rope_theta,
+            rope_scale,
+        )
+        self.freqs_cis = precompute_freqs_cis(
+            config.hidden_size // config.num_attention_heads,
+            config.max_position_embeddings * 2,
+            config.rope_theta,
+            rope_scale,
+        )
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        start_pos: Optional[int] = 0,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        chunk_prefilling: bool = False,
+    ):
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+        seqlen = inputs_embeds.shape[1]
+
+        self.freqs = self.freqs.to(inputs_embeds.device)
+        freqs = self.freqs[start_pos : start_pos + seqlen]
+
+        mask = None
+        if seqlen > 1:
+            mask = torch.full(
+                (1, 1, seqlen, seqlen), float("-inf"), device=inputs_embeds.device
+            )
+            mask = torch.triu(mask, diagonal=1).type_as(inputs_embeds)
+            if chunk_prefilling:
+                mask_history = torch.zeros(
+                    (1, 1, seqlen, start_pos),
+                    dtype=torch.float16,
+                    device=inputs_embeds.device,
+                ).type_as(inputs_embeds)
+                mask = torch.cat((mask_history, mask), dim=-1)
+        x = inputs_embeds
+
+        for decoder_layer in self.layers:
+            x = decoder_layer(x, start_pos, freqs, mask, chunk_prefilling)
+        x = x[:, -1:, :]
+        x = self.norm(x)
+
+        return x
+
+    def forwardfp16(
+        self,
+        input_ids: torch.LongTensor = None,
+        start_pos: Optional[int] = 0,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        chunk_prefilling: bool = False,
+    ):
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+        seqlen = inputs_embeds.shape[1]
+
+        self.freqs_cis = self.freqs_cis.to(inputs_embeds.device)
+        freqs_cis = self.freqs_cis[start_pos : start_pos + seqlen]
+
+        mask = None
+        if seqlen > 1:
+            mask = torch.full(
+                (1, 1, seqlen, seqlen), float("-inf"), device=inputs_embeds.device
+            )
+            mask = torch.triu(mask, diagonal=1).type_as(inputs_embeds)
+            if chunk_prefilling:
+                mask_history = torch.zeros(
+                    (1, 1, seqlen, start_pos),
+                    dtype=torch.float16,
+                    device=inputs_embeds.device,
+                ).type_as(inputs_embeds)
+                mask = torch.cat((mask_history, mask), dim=-1)
+        x = inputs_embeds
+
+        for decoder_layer in self.layers:
+            x = decoder_layer(x, start_pos, freqs_cis, mask, chunk_prefilling)
+        x = x[:, -1:, :]
+        x = self.norm(x)
+
+        return x
+
+
+class Qwen2ForCausalLM(Qwen2ForCausalLM):
+    def __init__(self, config):
+
+        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
+        from transformers import modeling_utils
+
+        modeling_utils._init_weights = False
+
+        super().__init__(config)
+        self.model = Qwen2Model(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.config = config
+
+    @torch.inference_mode()
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        start_pos: int = 0,
+        inputs_embeds: torch.Tensor = None,
+        chunk_prefilling: bool = False,
+        quant=True,
+    ):
+        if quant:
+            outputs = self.model(
+                input_ids=input_ids,
+                inputs_embeds=inputs_embeds,
+                start_pos=start_pos,
+                chunk_prefilling=chunk_prefilling,
+            )
+        else:
+            outputs = self.model.forwardfp16(
+                input_ids=input_ids,
+                inputs_embeds=inputs_embeds,
+                start_pos=start_pos,
+                chunk_prefilling=chunk_prefilling,
+            )
+        logits = self.lm_head(outputs)
+        return logits
+
+    def benchmark(self, inputs_embeds, attention_mask, max_output=128, quant_llm=True):
+        output_list = []
+        start_pos = 0
+        for i in range(10):
+            torch.cuda.synchronize()
+            tst = time.time()
+            token = self.forward(None, start_pos, inputs_embeds, quant=quant_llm)
+            torch.cuda.synchronize()
+            ted = time.time()
+        print(
+            "LLM TTFT: {:.6f} s for {} tokens".format(
+                (ted - tst), inputs_embeds.shape[1]
+            )
+        )
+        start_pos = inputs_embeds.shape[1]
+        token = torch.argmax(token, keepdim=True)[0]
+        output_list.append(token)
+
+        torch.cuda.synchronize()
+        tst = time.time()
+        for _ in range(max_output):
+            token = self.forward(token, start_pos)
+            token = torch.argmax(token, keepdim=True)[
+                0
+            ]  # Only fixed-length eager decoding is supported now
+            output_list.append(token)
+            start_pos += 1
+        torch.cuda.synchronize()
+        ted = time.time()
+        print("Decoding througput: {:.6f} tokens/s".format(max_output / (ted - tst)))
+
+        return torch.cat(output_list, dim=1)

llm-awq/tinychat/models/vila_llama.py

@@ -0,0 +1,109 @@
+import os
+import warnings
+import shutil
+import torch
+import torch.nn as nn
+from typing import List, Optional, Tuple, Union
+import time
+
+from transformers import AutoConfig, PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from llava.model.utils import get_model_config
+from llava.model.language_model.builder import build_llm_and_tokenizer
+from llava.model.multimodal_encoder.builder import build_vision_tower
+from llava.model.multimodal_projector.builder import build_mm_projector
+from llava.model.llava_arch import LlavaMetaModel, LlavaMetaForCausalLM
+from .llama import LlamaForCausalLM, Transformer
+
+
+class VilaLlamaForCausalLM(LlavaMetaModel, LlavaMetaForCausalLM, PreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.init_vlm(config)
+
+    def init_vlm(self, config=None, *args, **kwargs):
+        if (
+            hasattr(self, "llm")
+            or hasattr(self, "vision_tower")
+            or hasattr(self, "mm_projector")
+        ):
+            # already initialized, skipped
+            return
+
+        model_dtype = getattr(config, "model_dtype", "torch.float16")
+        if not hasattr(config, "model_dtype"):
+            warnings.warn(
+                "model_dtype not found in config, defaulting to torch.float16."
+            )
+            config.model_dtype = model_dtype
+
+        # print("init_vlm(): config", config); input("DEBUG init_vlm")
+        cfgs = get_model_config(config)
+        if len(cfgs) == 3:
+            llm_cfg, vision_tower_cfg, mm_projector_cfg = cfgs
+        else:
+            raise ValueError(
+                "`llm_cfg` `mm_projector_cfg` `vision_tower_cfg` not found in the config."
+            )
+        # print("init_vlm():", cfgs); input("DEBUG init_vlm")
+        llm_cfg = AutoConfig.from_pretrained(llm_cfg)
+
+        # self.llm, self.tokenizer = build_llm_and_tokenizer(llm_cfg, config, *args, **kwargs)
+        self.llm = LlamaForCausalLM(llm_cfg)
+        self.vision_tower = build_vision_tower(vision_tower_cfg, config)
+        self.mm_projector = build_mm_projector(mm_projector_cfg, config)
+
+        self.post_config()
+        self.is_loaded = True
+
+        assert (
+            self.llm is not None
+            or self.vision_tower is not None
+            or self.mm_projector is not None
+        ), "At least one of the components must be instantiated."
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        start_pos: int = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        images: Optional[torch.FloatTensor] = None,
+        return_dict: Optional[bool] = None,
+        special_token: bool = False,
+        chunk_prefilling: bool = False,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        self.freezed_module_patch()
+        if inputs_embeds is None:
+            (
+                _,
+                _,
+                _,
+                _,
+                inputs_embeds,
+                _,
+            ) = self.prepare_inputs_labels_for_multimodal(
+                input_ids, position_ids, attention_mask, past_key_values, labels, images
+            )
+        if inputs_embeds is not None:
+            outputs = self.llm.forward(
+                tokens=None,
+                start_pos=start_pos,
+                inputs_embeds=inputs_embeds,
+                chunk_prefilling=chunk_prefilling,
+            )
+        else:  # tokens
+            outputs = self.llm.forward(
+                tokens=input_ids,
+                start_pos=start_pos,
+                inputs_embeds=None,
+                chunk_prefilling=chunk_prefilling,
+            )
+        return outputs

llm-awq/tinychat/modules/__init__.py

@@ -0,0 +1,9 @@
+from .fused_norm import *
+from .fused_attn import *
+from .fused_mlp import *
+from .fused_vision_attn import *
+try:
+    from .fused_siglipdecoder import *
+    from .fused_internencoder import *
+except ImportError as e:
+    print("InternVL3 model import failure. To activate, please install VILA at https://github.com/NVlabs/VILA.")

llm-awq/tinychat/modules/fused_attn.py

@@ -0,0 +1,634 @@
+import math
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from transformers.models.llama.modeling_llama import (
+    LlamaAttention,
+    LlamaRotaryEmbedding,
+    apply_rotary_pos_emb,
+)
+from typing import Optional
+from awq.quantize.qmodule import WQLinear
+import awq_inference_engine
+from tinychat.models.llama import apply_rotary_emb
+import gc
+
+import tinychat.utils.constants
+from flash_attn import flash_attn_func
+from tinychat.models.llama import LlamaAttentionFused
+from tinychat.models.qwen2 import Qwen2AttentionFused
+
+max_batch_size = tinychat.utils.constants.max_batch_size
+max_seq_len = tinychat.utils.constants.max_seq_len
+
+
+class QuantLlamaRotaryEmbedding(nn.Module):
+    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
+        super().__init__()
+
+        self.dim = dim
+        self.max_position_embeddings = max_position_embeddings
+        self.base = base
+        inv_freq = 1.0 / (
+            self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim)
+        )
+        self.register_buffer("inv_freq", inv_freq)
+        # Build here to make `torch.jit.trace` work.
+        self._set_cos_sin_cache(
+            seq_len=max_position_embeddings,
+            device=self.inv_freq.device,
+            dtype=torch.get_default_dtype(),
+        )
+
+    def _set_cos_sin_cache(self, seq_len, device, dtype):
+        self.max_seq_len_cached = seq_len
+        t = torch.arange(
+            self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype
+        )
+
+        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
+        # Different from paper, but it uses a different permutation in order to obtain the same calculation
+        # emb = torch.cat((freqs, freqs), dim=-1)
+
+        cos = freqs.cos()
+        sin = freqs.sin()
+        cache = torch.cat((cos, sin), dim=-1)
+
+        # self.register_buffer("cos_cached", emb.cos()[None, None, :, :].to(dtype), persistent=False)
+        # self.register_buffer("sin_cached", emb.sin()[None, None, :, :].to(dtype), persistent=False)
+        self.register_buffer("cos_sin_cache", cache.half(), persistent=False)
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        positions: torch.Tensor,
+    ):
+        # Apply rotary embedding to the query and key before passing them
+        # to the attention op.
+        # print(positions.shape, query.shape, key.shape, self.cos_sin_cache.shape)
+        query = query.contiguous()
+        key = key.contiguous()
+        awq_inference_engine.rotary_embedding_neox(
+            positions,
+            query,
+            key,
+            self.dim,
+            self.cos_sin_cache,
+        )
+        return query, key
+
+
+class QuantLlamaAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, hidden_size, num_heads, qkv_proj, o_proj, dev):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        self.head_dim = hidden_size // num_heads
+
+        if (self.head_dim * num_heads) != self.hidden_size:
+            raise ValueError(
+                f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
+                f" and `num_heads`: {num_heads})."
+            )
+        self.qkv_proj = qkv_proj
+        self.o_proj = o_proj
+        self.rotary_emb = QuantLlamaRotaryEmbedding(
+            self.head_dim, max_position_embeddings=2048, device=dev
+        )
+
+    def forward(
+        self,
+        hidden_states,
+        past_key_value=None,
+        attention_mask=None,
+        position_ids=None,
+        output_attentions=False,
+        use_cache=False,
+    ):
+        """Input shape: Batch x Time x Channel"""
+
+        bsz, q_len, _ = hidden_states.size()
+
+        qkv_states = self.qkv_proj(hidden_states)
+        qkv_states = qkv_states.view(bsz, q_len, 3, self.num_heads, self.head_dim)
+
+        # This updates the query and key states in-place, saving VRAM.
+        query_states, key_states, value_states = torch.split(qkv_states, 1, dim=2)
+        query_states, key_states = self.rotary_emb(
+            query_states, key_states, position_ids
+        )
+
+        del qkv_states
+        query_states = query_states.view(
+            bsz, q_len, self.num_heads, self.head_dim
+        ).transpose(1, 2)
+        key_states = key_states.view(
+            bsz, q_len, self.num_heads, self.head_dim
+        ).transpose(1, 2)
+        value_states = value_states.view(
+            bsz, q_len, self.num_heads, self.head_dim
+        ).transpose(1, 2)
+
+        is_causal = past_key_value is None
+
+        kv_seq_len = q_len
+        if past_key_value is not None:
+            kv_seq_len += past_key_value[0].shape[-2]
+
+        value_states = value_states.to("cuda:0")
+
+        if past_key_value is not None:
+            # reuse k, v, self_attention
+            key_states = torch.cat([past_key_value[0], key_states], dim=2)
+            value_states = torch.cat([past_key_value[1], value_states], dim=2)
+
+        if use_cache:
+            # Since qkv_proj is fused, query_states etc will hold a reference to the original qkv_states tensor
+            # which can cause excessive memory usage by the cache. `contiguous` is a convenient way to workaround this.
+            key_states = key_states.contiguous()
+            value_states = value_states.contiguous()
+            query_states = query_states.contiguous()
+
+        past_key_value = (key_states, value_states) if use_cache else None
+
+        # with torch.backends.cuda.sdp_kernel(enable_math=False):
+        attn_output = F.scaled_dot_product_attention(
+            query_states, key_states, value_states, is_causal=is_causal
+        )
+        del query_states, key_states, value_states
+
+        attn_output = attn_output.transpose(1, 2).reshape(bsz, q_len, self.hidden_size)
+        attn_output = self.o_proj(attn_output)
+
+        return attn_output, None, past_key_value
+
+
+class QuantLlamaAttentionFused(nn.Module):
+    def __init__(
+        self, hidden_size, num_heads, kv_max_seq_len, qkv_layer, o_proj, dev, args
+    ):
+        super().__init__()
+
+        self.args = args
+        self.n_local_heads = args.num_attention_heads
+        self.hidden_size = args.hidden_size
+        self.num_heads = args.num_attention_heads
+        self.head_dim = self.hidden_size // self.num_heads
+
+        self.num_key_value_heads = args.num_key_value_heads
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.max_position_embeddings = args.max_position_embeddings
+        self.rope_theta = args.rope_theta
+        self.rope_scaling = args.rope_scaling
+        if self.rope_scaling is None:
+            self.rope_scaling = 1.0
+        if isinstance(self.rope_scaling, dict):
+            self.rope_scaling = self.rope_scaling.get("factor", 1.0)
+
+        self.qkv_proj = qkv_layer
+        self.o_proj = o_proj
+
+        self.kv_max_seq_len = kv_max_seq_len
+
+        # following fastertransformer definition
+        self.cache_v = (
+            torch.zeros(
+                (
+                    max_batch_size,
+                    self.num_key_value_heads,
+                    # args.max_position_embeddings,
+                    self.kv_max_seq_len,
+                    self.head_dim,
+                )
+            )
+            .to(dev)
+            .half()
+        )  # added to half
+        # 8: pack 8 fp16 in FT, if fp32 then use 4
+        self.cache_k = (
+            torch.zeros(
+                (
+                    max_batch_size,
+                    self.num_key_value_heads,
+                    self.head_dim // 8,
+                    # args.max_position_embeddings,
+                    self.kv_max_seq_len,
+                    8,
+                )
+            )
+            .to(dev)
+            .half()
+        )  # added to half
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        start_pos: int,
+        freqs: torch.Tensor,
+        mask: Optional[torch.Tensor],
+        chunk_prefilling: bool = False,
+    ):
+        bsz, seqlen, _ = x.shape
+        xqkv = self.qkv_proj(x)
+        xqkv = xqkv.view(
+            bsz,
+            seqlen,
+            self.n_local_heads + self.num_key_value_heads * 2,
+            self.head_dim,
+        )
+        xq = xqkv[:, :, 0 : self.n_local_heads]
+        xk = xqkv[
+            :, :, self.n_local_heads : (self.n_local_heads + self.num_key_value_heads)
+        ]
+        xv = xqkv[:, :, -self.num_key_value_heads :]
+
+        if seqlen > 1:
+            xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim)
+            xk = xk.view(bsz, seqlen, self.num_key_value_heads, self.head_dim)
+            xv = xv.view(bsz, seqlen, self.num_key_value_heads, self.head_dim)
+
+            xq = awq_inference_engine.fused_rope_with_pos_forward_func(xq, freqs, True)
+            xk = awq_inference_engine.fused_rope_with_pos_forward_func(xk, freqs, True)
+
+            self.cache_k = self.cache_k.to(xq)
+            self.cache_v = self.cache_v.to(xq)
+
+            values_store = xv.transpose(2, 1)
+            keys_store = (
+                xk.reshape(bsz, seqlen, self.num_key_value_heads, self.head_dim // 8, 8)
+                .permute(0, 2, 3, 1, 4)
+                .contiguous()
+            )
+
+            self.cache_v[:bsz, :, start_pos : start_pos + seqlen, :] = values_store
+            self.cache_k[:bsz, :, :, start_pos : start_pos + seqlen, :] = keys_store
+            if chunk_prefilling:
+                keys = self.cache_k[:, :, :, 0:start_pos, :]
+                keys = (
+                    keys.permute(0, 3, 1, 2, 4)
+                    .reshape(bsz, start_pos, self.num_key_value_heads, self.head_dim)
+                    .contiguous()
+                )
+                keys = torch.cat((keys, xk), dim=1)
+                values = self.cache_v[:, :, 0:start_pos, :]
+                values = (
+                    values.transpose(2, 1)
+                    .reshape(bsz, start_pos, self.num_key_value_heads, self.head_dim)
+                    .contiguous()
+                )
+                values = torch.cat((values, xv), dim=1)
+            else:
+                keys = xk
+                values = xv
+
+            keys = torch.repeat_interleave(
+                keys, dim=2, repeats=self.num_key_value_groups
+            )
+            values = torch.repeat_interleave(
+                values, dim=2, repeats=self.num_key_value_groups
+            )
+
+            xq = xq.transpose(1, 2)
+            keys = keys.transpose(1, 2)
+            values = values.transpose(1, 2)
+            scores = torch.matmul(xq, keys.transpose(2, 3)) / math.sqrt(self.head_dim)
+            if mask is not None:
+                scores = scores + mask  # (bs, n_local_heads, slen, cache_len + slen)
+            scores = F.softmax(scores.float(), dim=-1).type_as(xq)
+            output = torch.matmul(scores, values)  # (bs, n_local_heads, slen, head_dim)
+            output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)
+        else:
+            xq = xq.view(bsz, self.n_local_heads, self.head_dim)
+            xk = xk.view(bsz, self.num_key_value_heads, self.head_dim)
+            xv = xv.view(bsz, self.num_key_value_heads, self.head_dim)
+
+            output = awq_inference_engine.single_query_attention(
+                xq,
+                xk,
+                xv,
+                self.cache_k,
+                self.cache_v,
+                None,
+                None,
+                start_pos,
+                self.head_dim,
+                self.rope_theta,
+                self.rope_scaling,
+                True,
+            )
+            output = output.reshape(bsz, 1, -1)
+
+        return self.o_proj(output)
+
+
+class QuantLlamaAttentionFusedFlash(nn.Module):
+    """Flash_attn_func from 'Flash{A}ttention-2: Faster Attention with Better Parallelism and Work Partitioning' paper"""
+
+    """This function is faster than the varlen one but only supports single-batch inference"""
+
+    def __init__(
+        self, hidden_size, num_heads, kv_max_seq_len, qkv_layer, o_proj, dev, args
+    ):
+        super().__init__()
+
+        self.args = args
+        self.n_local_heads = args.num_attention_heads
+        self.hidden_size = args.hidden_size
+        self.num_heads = args.num_attention_heads
+        self.head_dim = self.hidden_size // self.num_heads
+
+        self.num_key_value_heads = args.num_key_value_heads
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.max_position_embeddings = args.max_position_embeddings
+        self.rope_theta = args.rope_theta
+        self.rope_scaling = args.rope_scaling
+        if self.rope_scaling is None:
+            self.rope_scaling = 1.0
+        elif isinstance(self.rope_scaling, dict):
+            self.rope_scaling = self.rope_scaling.get("factor", 1.0)
+
+        self.qkv_proj = qkv_layer
+        self.o_proj = o_proj
+
+        self.kv_max_seq_len = kv_max_seq_len
+        # following fastertransformer definition
+        # For short seqlence, we use fused kernel to accelerate decoding.
+        if self.kv_max_seq_len <= 8192:
+            self.cache_v = (
+                torch.zeros(
+                    (
+                        max_batch_size,
+                        self.num_key_value_heads,
+                        # args.max_position_embeddings,
+                        self.kv_max_seq_len,
+                        self.head_dim,
+                    )
+                )
+                .to(dev)
+                .half()
+            )  # added to half
+            # 8: pack 8 fp16 in FT, if fp32 then use 4
+            self.cache_k = (
+                torch.zeros(
+                    (
+                        max_batch_size,
+                        self.num_key_value_heads,
+                        self.head_dim // 8,
+                        # args.max_position_embeddings,
+                        kv_max_seq_len,
+                        8,
+                    )
+                )
+                .to(dev)
+                .half()
+            )  # added to half
+            self.forward = self.short_forward
+        # For long sequence, we use flash attantion for both prefilling and decoding to avoid OOM.
+        else:
+            self.cache_v = (
+                torch.zeros(
+                    (
+                        max_batch_size,
+                        self.kv_max_seq_len,
+                        self.num_key_value_heads,
+                        self.head_dim,
+                    )
+                )
+                .to(dev)
+                .half()
+            )  # added to half
+            self.cache_k = (
+                torch.zeros(
+                    (
+                        max_batch_size,
+                        self.kv_max_seq_len,
+                        self.num_key_value_heads,
+                        self.head_dim,
+                    )
+                )
+                .to(dev)
+                .half()
+            )  # added to half
+            self.forward = self.long_forward
+
+    def short_forward(
+        self,
+        x: torch.Tensor,
+        start_pos: int,
+        freqs: torch.Tensor,
+        mask: Optional[torch.Tensor],
+        chunk_prefilling: bool = False,
+    ):
+        bsz, seqlen, _ = x.shape
+        xqkv = self.qkv_proj(x)
+        xqkv = xqkv.view(
+            bsz,
+            seqlen,
+            self.n_local_heads + self.num_key_value_heads * 2,
+            self.head_dim,
+        )
+        xq = xqkv[:, :, 0 : self.n_local_heads]
+        xk = xqkv[
+            :, :, self.n_local_heads : (self.n_local_heads + self.num_key_value_heads)
+        ]
+        xv = xqkv[:, :, -self.num_key_value_heads :]
+
+        if seqlen > 1:
+            xq = awq_inference_engine.fused_rope_with_pos_forward_func(xq, freqs, True)
+            xk = awq_inference_engine.fused_rope_with_pos_forward_func(xk, freqs, True)
+
+            self.cache_k = self.cache_k.to(xq)
+            self.cache_v = self.cache_v.to(xq)
+
+            values_store = xv.transpose(2, 1)
+            keys_store = (
+                xk.reshape(bsz, seqlen, self.num_key_value_heads, self.head_dim // 8, 8)
+                .permute(0, 2, 3, 1, 4)
+                .contiguous()
+            )
+
+            self.cache_v[:bsz, :, start_pos : start_pos + seqlen, :] = values_store
+            self.cache_k[:bsz, :, :, start_pos : start_pos + seqlen, :] = keys_store
+
+            if chunk_prefilling:
+                keys = self.cache_k[:, :, :, 0 : start_pos + seqlen, :]
+                keys = (
+                    keys.permute(0, 3, 1, 2, 4)
+                    .reshape(
+                        bsz, start_pos + seqlen, self.num_key_value_heads, self.head_dim
+                    )
+                    .contiguous()
+                )
+                values = self.cache_v[:, :, 0 : start_pos + seqlen, :]
+                values = (
+                    values.transpose(2, 1)
+                    .reshape(
+                        bsz, start_pos + seqlen, self.num_key_value_heads, self.head_dim
+                    )
+                    .contiguous()
+                )
+            else:
+                keys = xk
+                values = xv
+
+            output = flash_attn_func(
+                q=xq,
+                k=keys,
+                v=values,
+                causal=True,
+            )
+            output = output.contiguous().view(bsz, seqlen, -1)
+        else:
+            xq = xq.view(bsz, self.n_local_heads, self.head_dim)
+            xk = xk.view(bsz, self.num_key_value_heads, self.head_dim)
+            xv = xv.view(bsz, self.num_key_value_heads, self.head_dim)
+            output = awq_inference_engine.single_query_attention(
+                xq,
+                xk,
+                xv,
+                self.cache_k,
+                self.cache_v,
+                None,
+                None,
+                start_pos,
+                self.head_dim,
+                self.rope_theta,
+                self.rope_scaling,
+                True,
+            )
+            output = output.reshape(bsz, 1, -1)
+        return self.o_proj(output)
+
+    def long_forward(
+        self,
+        x: torch.Tensor,
+        start_pos: int,
+        freqs: torch.Tensor,
+        mask: Optional[torch.Tensor],
+        chunk_prefilling: bool = False,
+    ):
+        bsz, seqlen, _ = x.shape
+        xqkv = self.qkv_proj(x)
+        xqkv = xqkv.view(
+            bsz,
+            seqlen,
+            self.n_local_heads + self.num_key_value_heads * 2,
+            self.head_dim,
+        )
+        xq = xqkv[:, :, 0 : self.n_local_heads]
+        xk = xqkv[
+            :, :, self.n_local_heads : (self.n_local_heads + self.num_key_value_heads)
+        ]
+        xv = xqkv[:, :, -self.num_key_value_heads :]
+
+        xq = awq_inference_engine.fused_rope_with_pos_forward_func(xq, freqs, True)
+        xk = awq_inference_engine.fused_rope_with_pos_forward_func(xk, freqs, True)
+
+        self.cache_k = self.cache_k.to(xq)
+        self.cache_v = self.cache_v.to(xq)
+
+        self.cache_v[:bsz, start_pos : start_pos + seqlen] = xv
+        self.cache_k[:bsz, start_pos : start_pos + seqlen] = xk
+
+        keys = self.cache_k[:, 0 : start_pos + seqlen]
+        values = self.cache_v[:, 0 : start_pos + seqlen]
+
+        output = flash_attn_func(
+            q=xq,
+            k=keys,
+            v=values,
+            causal=True,
+        )
+        output = output.view(bsz, seqlen, -1)
+        return self.o_proj(output)
+
+
+def make_quant_attn(model, dev, flash_attn=True):
+    """
+    Replace all LlamaAttention modules with QuantLlamaAttention modules, fusing the q, k, v projections.
+    """
+    model = model.cpu()
+    for name, m in model.named_modules():
+        if not m.__class__.__name__ in [
+            "LlamaAttention",
+            "LlamaAttentionFused",
+            "Qwen2AttentionFused",
+        ]:
+            continue
+
+        q_proj = m.q_proj
+        k_proj = m.k_proj
+        v_proj = m.v_proj
+
+        qweights = torch.cat([q_proj.qweight, k_proj.qweight, v_proj.qweight], dim=0)
+        scaled_zeros = torch.cat(
+            [q_proj.scaled_zeros, k_proj.scaled_zeros, v_proj.scaled_zeros], dim=1
+        ).contiguous()
+        scales = torch.cat(
+            [q_proj.scales, k_proj.scales, v_proj.scales], dim=1
+        ).contiguous()
+        # g_idx = torch.cat([q_proj.g_idx, k_proj.g_idx, v_proj.g_idx], dim=0)
+        g_idx = None
+        bias = (
+            torch.cat([q_proj.bias, k_proj.bias, v_proj.bias], dim=0)
+            if q_proj.bias is not None
+            else None
+        )
+
+        qkv_layer = WQLinear(
+            q_proj.w_bit,
+            q_proj.group_size,
+            q_proj.in_features,
+            q_proj.out_features + k_proj.out_features + v_proj.out_features,
+            q_proj.bias is not None,
+            q_proj.qweight.device,
+        )
+        qkv_layer.qweight = qweights
+        qkv_layer.scaled_zeros = scaled_zeros
+        qkv_layer.scales = scales
+
+        qkv_layer.bias = bias
+        qkv_layer.split_k_iters = q_proj.split_k_iters
+        # We're dropping the rotary embedding layer m.rotary_emb here. We don't need it in the triton branch.
+        if isinstance(m, LlamaAttention):
+            attn = QuantLlamaAttention(
+                m.hidden_size, m.num_heads, qkv_layer, m.o_proj, dev
+            )
+        else:
+            if flash_attn:
+                attn = QuantLlamaAttentionFusedFlash(
+                    m.args.hidden_size,
+                    m.args.num_attention_heads,
+                    m.kv_max_seq_len,
+                    qkv_layer,
+                    m.o_proj,
+                    dev,
+                    m.args,
+                )
+            else:
+                attn = QuantLlamaAttentionFused(
+                    m.args.hidden_size,
+                    m.args.num_attention_heads,
+                    m.kv_max_seq_len,
+                    qkv_layer,
+                    m.o_proj,
+                    dev,
+                    m.args,
+                )
+        if "." in name:
+            parent_name = name.rsplit(".", 1)[0]
+            child_name = name[len(parent_name) + 1 :]
+            parent = model.get_submodule(parent_name)
+        else:
+            parent_name = ""
+            parent = model
+            child_name = name
+
+        # print(f"Replacing {name} with quant_attn; parent: {parent_name}, child's name: {child_name}")
+        setattr(parent, child_name, attn)
+        gc.collect()
+        torch.cuda.empty_cache()
+    model = model.to(dev)

llm-awq/tinychat/modules/fused_internencoder.py

@@ -0,0 +1,237 @@
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from einops import rearrange
+from timm.layers import DropPath
+from torch import nn
+from transformers.activations import ACT2FN
+from transformers.modeling_outputs import (BaseModelOutput,
+                                           BaseModelOutputWithPooling)
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import logging
+
+from awq.quantize import W8A8OF16LinearDynamicInputScale
+import awq_inference_engine
+
+from tinychat.models.internvl.internvit import (FlashAttention,
+                                                InternRMSNorm,
+                                                InternVisionEmbeddings,
+                                                InternAttention,
+                                                InternMLP,
+                                                InternVisionEncoderLayer,
+                                                InternVisionEncoder)
+from tinychat.models.internvl.configuration_internvl import InternVisionConfig
+
+try:
+    from flash_attn.bert_padding import pad_input, unpad_input
+    from flash_attn.flash_attn_interface import \
+        flash_attn_varlen_qkvpacked_func
+    has_flash_attn = True
+except:
+    print('FlashAttention2 is not installed.')
+    has_flash_attn = False
+
+logger = logging.get_logger(__name__)
+
+
+class QuantInternVisionEncoder(nn.Module):
+    """
+    Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a
+    [`InternEncoderLayer`].
+
+    Args:
+        config (`InternConfig`):
+            The corresponding vision configuration for the `InternEncoder`.
+    """
+
+    def __init__(self, module: InternVisionEncoder, bsz=64, seqlen=1024):
+        super().__init__()
+        self.config = module.config
+        # stochastic depth decay rule
+        self.layers = nn.ModuleList([QuantInternVisionEncoderLayer(layer, self.config) for layer in module.layers])
+        self.gradient_checkpointing = True
+        self.bsz = bsz
+        self.seqlen = seqlen
+
+    def forward(
+        self,
+        inputs_embeds,
+        attention_mask: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutput]:
+        r"""
+        Args:
+            inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+                Embedded representation of the inputs. Should be float, not int tokens.
+            output_hidden_states (`bool`, *optional*):
+                Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
+                for more detail.
+            return_dict (`bool`, *optional*):
+                Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+        """
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        encoder_states = () if output_hidden_states else None
+        hidden_states = inputs_embeds
+
+        for idx, encoder_layer in enumerate(self.layers):
+            if output_hidden_states:
+                encoder_states = encoder_states + (hidden_states,)
+            if self.gradient_checkpointing and self.training:
+                layer_outputs = torch.utils.checkpoint.checkpoint(
+                    encoder_layer,
+                    hidden_states)
+            else:
+                layer_outputs = encoder_layer(
+                    hidden_states,
+                )
+            hidden_states = layer_outputs
+
+        if output_hidden_states:
+            encoder_states = encoder_states + (hidden_states,)
+
+        if not return_dict:
+            return tuple(v for v in [hidden_states, encoder_states] if v is not None)
+        return BaseModelOutput(
+            last_hidden_state=hidden_states, hidden_states=encoder_states
+        )
+
+class QuantInternRMSNorm(nn.Module):
+    def __init__(self, module: nn.Module, use_per_token_quant=True):
+        super().__init__()
+        self.weight = nn.Parameter(module.weight.data, requires_grad=False)
+        self.bias = nn.Parameter(module.bias.data, requires_grad=False)
+        self.variance_epsilon = module.eps
+        self.use_per_token_quant = use_per_token_quant
+
+    def forward(self, hidden_states):
+        bsz, seqlen, hidden_size = hidden_states.shape
+        output = torch.empty((bsz * seqlen), hidden_size, device=hidden_states.device, dtype=torch.int8)
+        scale = torch.empty((bsz * seqlen), device=hidden_states.device, dtype=hidden_states.dtype)
+        awq_inference_engine.rms_norm_general(
+            output,
+            hidden_states,
+            self.weight,
+            self.bias,
+            scale,
+            self.variance_epsilon,
+            self.use_per_token_quant,
+        )
+        return output, scale
+
+class QuantInternAttention(nn.Module):
+    def __init__(self, module: InternAttention, config: InternVisionConfig, init_only=False):
+        super().__init__()
+        self.config = config
+        self.embed_dim = module.embed_dim
+        self.num_heads = module.num_heads
+        self.head_dim = self.embed_dim // self.num_heads
+        self.scale = module.scale
+        self.use_flash_attn = config.use_flash_attn
+
+        self.qkv = W8A8OF16LinearDynamicInputScale.from_linear(module.qkv, init_only=init_only)
+        self.proj = W8A8OF16LinearDynamicInputScale.from_linear(module.proj, init_only=init_only)
+
+        self.qk_normalization = module.qk_normalization
+        if self.qk_normalization:
+            self.q_norm = QuantInternRMSNorm(module.q_norm)
+            self.k_norm = QuantInternRMSNorm(module.k_norm)
+
+        if self.use_flash_attn:
+            from tinychat.models.internvl.internvit import FlashAttention
+            self.inner_attn = FlashAttention(attention_dropout=config.attention_dropout)
+
+    def forward(self, hidden_states: torch.Tensor, scale_in: torch.Tensor):
+        bsz, seqlen, hidden_size = hidden_states.shape
+
+        qkv_out = torch.empty(bsz * seqlen, 3 * hidden_size, dtype=torch.float16, device=hidden_states.device)
+        self.qkv(hidden_states.reshape(-1, hidden_size), scale_in, qkv_out)
+
+        qkv = rearrange(qkv_out.view(bsz, seqlen, -1), 'b s (three h d) -> b s three h d', three=3, h=self.num_heads)
+
+        if self.qk_normalization:
+            q, k, v = qkv.unbind(2)
+            q, _ = self.q_norm(q.flatten(-2, -1)); q = q.view_as(q)
+            k, _ = self.k_norm(k.flatten(-2, -1)); k = k.view_as(k)
+            qkv = torch.stack([q, k, v], dim=2)
+
+        attn_out, _ = self.inner_attn(qkv, need_weights=False, causal=False)
+        attn_out = rearrange(attn_out, 'b s h d -> (b s) (h d)')
+
+        quant_out = torch.empty_like(attn_out, dtype=torch.int8)
+        scale_proj_in = torch.empty(bsz * seqlen, device=hidden_states.device, dtype=torch.float16)
+        awq_inference_engine.invoke_quant(quant_out, attn_out, scale_proj_in)
+
+        proj_out = torch.empty_like(attn_out)
+        self.proj(quant_out, scale_proj_in, proj_out)
+
+        return proj_out
+
+class QuantInternMLP(nn.Module):
+    def __init__(self, module: InternMLP, config: InternVisionConfig):
+        super().__init__()
+        self.config = config
+        self.act = module.act
+        self.fc1 = W8A8OF16LinearDynamicInputScale.from_linear(module.fc1)
+        self.fc2 = W8A8OF16LinearDynamicInputScale.from_linear(module.fc2)
+
+    def forward(self, hidden_states: torch.Tensor, scale_in: torch.Tensor):
+        bsz, seqlen, hidden_size = hidden_states.shape
+        device = hidden_states.device
+        
+        fc1_out = torch.empty((bsz * seqlen), self.config.intermediate_size, dtype=torch.float16, device=device)
+        self.fc1(hidden_states.reshape(-1, hidden_size), scale_in, fc1_out)
+        
+        tmp = torch.empty(
+            ((bsz * seqlen) * self.config.intermediate_size),
+            device=device,
+            dtype=torch.float16,
+        )
+        act_out = torch.empty_like(fc1_out, dtype=torch.int8)
+        scale_act = torch.empty(bsz * seqlen, device=device, dtype=torch.float16)
+        awq_inference_engine.gelu_and_quant(act_out, fc1_out, scale_act, tmp)
+
+        fc2_out = torch.empty((bsz * seqlen), hidden_size, dtype=torch.float16, device=device)
+        self.fc2(act_out, scale_act, fc2_out)
+
+        return fc2_out
+
+class QuantInternVisionEncoderLayer(nn.Module):
+    def __init__(self, module: InternVisionEncoderLayer, config: InternVisionConfig):
+        super().__init__()
+        self.config = config
+        self.embed_dim = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+
+        self.attn = QuantInternAttention(module.attn, config)
+        self.mlp = QuantInternMLP(module.mlp, config)
+
+        self.norm1 = QuantInternRMSNorm(module.norm1)
+        self.norm2 = QuantInternRMSNorm(module.norm2)
+
+        self.ls1 = module.ls1
+        self.ls2 = module.ls2
+
+    def forward(self, hidden_states: torch.Tensor):
+        bsz, seqlen, hidden_size = hidden_states.shape
+
+        residual = hidden_states
+        norm1_out, scale1 = self.norm1(hidden_states)
+        attn_out = self.attn(norm1_out.view(bsz, seqlen, hidden_size), scale1)
+        hidden_states = residual + attn_out.view(bsz, seqlen, hidden_size) * self.ls1
+
+        residual = hidden_states
+        norm2_out, scale2 = self.norm2(hidden_states)
+        mlp_out = self.mlp(norm2_out.view(bsz, seqlen, hidden_size), scale2)
+        hidden_states = residual + mlp_out.view(bsz, seqlen, hidden_size) * self.ls2
+
+        return hidden_states
+
+

llm-awq/tinychat/modules/fused_mlp.py

@@ -0,0 +1,101 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.cuda.amp import custom_bwd, custom_fwd
+from transformers.models.llama.modeling_llama import LlamaMLP
+
+import awq_inference_engine
+
+
+class QuantLlamaMLP(nn.Module):
+    def __init__(
+        self,
+        gate_proj,
+        down_proj,
+        up_proj,
+    ):
+        super().__init__()
+        self.register_buffer("gate_proj_qweight", gate_proj.qweight)
+        self.register_buffer("gate_proj_scales", gate_proj.scales)
+        self.register_buffer("gate_proj_scaled_zeros", gate_proj.scaled_zeros)
+        self.register_buffer("up_proj_qweight", up_proj.qweight)
+        self.register_buffer("up_proj_scales", up_proj.scales)
+        self.register_buffer("up_proj_scaled_zeros", up_proj.scaled_zeros)
+
+        self.in_features = gate_proj.in_features
+        self.intermediate_size = gate_proj.out_features
+        self.out_features = down_proj.out_features
+        self.w_bit = gate_proj.w_bit
+        self.down_proj = down_proj
+        self.split_k_iters = down_proj.split_k_iters
+
+    def forward(self, x):
+        return self.down_proj(self.our_llama_mlp(x))
+
+    def our_llama_mlp(self, x):
+        # out_shape = x.shape[:-1] + (self.intermediate_size,)
+        # x = x.reshape(-1, x.shape[-1])
+        if x.numel() // x.shape[-1] < 8:
+            gate_output = awq_inference_engine.gemv_forward_cuda_new(
+                x,
+                self.gate_proj_qweight,
+                self.gate_proj_scales,
+                self.gate_proj_scaled_zeros,
+                x.numel() // x.shape[-1],
+                self.intermediate_size,
+                self.in_features,
+                self.down_proj.group_size,
+            )
+            gate_output = F.silu(gate_output)
+            up_output = awq_inference_engine.gemv_forward_cuda_new(
+                x,
+                self.up_proj_qweight,
+                self.up_proj_scales,
+                self.up_proj_scaled_zeros,
+                x.numel() // x.shape[-1],
+                self.intermediate_size,
+                self.in_features,
+                self.down_proj.group_size,
+            )
+        else:
+            # num_mn_tiles = (x.shape[0] // 32) * (self.intermediate_size // 128)
+            # cuda_Semaphores_gate = torch.empty(num_mn_tiles).int().to(x.device)
+            # cuda_Semaphores_up = torch.empty(num_mn_tiles).int().to(x.device)
+            gate_output = awq_inference_engine.gemm_forward_cuda_new(
+                x,
+                self.gate_proj_qweight,
+                self.gate_proj_scales,
+                self.gate_proj_scaled_zeros - 8 * self.gate_proj_scales,
+                # self.gate_cuda_semaphores
+            )
+            up_output = awq_inference_engine.gemm_forward_cuda_new(
+                x,
+                self.up_proj_qweight,
+                self.up_proj_scales,
+                self.up_proj_scaled_zeros - 8 * self.up_proj_scales,
+                # self.up_cuda_semaphores
+            )
+            gate_output = F.silu(gate_output)
+
+        c = gate_output * up_output
+        # c = c.reshape(out_shape)
+        return c
+
+
+def make_fused_mlp(m, parent_name=""):
+    if not hasattr(make_fused_mlp, "called"):
+        # print("[Warning] Calling a fake MLP fusion. But still faster than Huggingface Implimentation.")
+        make_fused_mlp.called = True
+    """
+    Replace all LlamaMLP modules with QuantLlamaMLP modules, which fuses many of the operations.
+    """
+    if m.__class__.__name__ in ["LlamaMLP"]:
+        return QuantLlamaMLP(m.gate_proj, m.down_proj, m.up_proj)
+
+    for name, child in m.named_children():
+        child = make_fused_mlp(child, parent_name=f"{parent_name}.{name}")
+
+        if isinstance(child, QuantLlamaMLP):
+            setattr(m, name, child)
+    return m

llm-awq/tinychat/modules/fused_norm.py

@@ -0,0 +1,46 @@
+import torch
+from torch import nn
+from transformers.models.llama.modeling_llama import LlamaRMSNorm
+import awq_inference_engine
+
+
+class FTLlamaRMSNorm(nn.Module):
+    def __init__(self, weight, eps=1e-6):
+        """
+        LlamaRMSNorm is equivalent to T5LayerNorm
+        """
+        super().__init__()
+        self.weight = weight
+        self.variance_epsilon = eps
+
+    def forward(self, x):
+        output = torch.empty_like(x)
+        awq_inference_engine.layernorm_forward_cuda(
+            x, self.weight, output, self.variance_epsilon
+        )
+        return output
+
+
+def make_quant_norm(model):
+    """
+    Replace all LlamaRMSNorm modules with FTLlamaRMSNorm modules
+    """
+
+    for name, m in model.named_modules():
+        if not isinstance(m, LlamaRMSNorm):
+            continue
+
+        norm = FTLlamaRMSNorm(m.weight, m.variance_epsilon)
+
+        if "." in name:
+            parent_name = name.rsplit(".", 1)[0]
+            child_name = name[len(parent_name) + 1 :]
+            parent = model.get_submodule(parent_name)
+        else:
+            parent_name = ""
+            parent = model
+            child_name = name
+
+        # print(f"Replacing {name} with quant_attn; parent: {parent_name}, child's name: {child_name}")
+
+        setattr(parent, child_name, norm)

llm-awq/tinychat/modules/fused_siglipdecoder.py

@@ -0,0 +1,282 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from awq.quantize import W8A8OF16LinearDynamicInputScale
+from llava.model.multimodal_encoder.siglip.modeling_siglip import (
+    SiglipMLP,
+    SiglipEncoder,
+    SiglipAttention,
+    SiglipEncoderLayer,
+)
+from tinychat.utils.input_metadata import ActivationBuffer
+from transformers.modeling_outputs import BaseModelOutput
+from typing import Optional, Tuple, Union
+from flash_attn import flash_attn_func
+import time
+
+CLIP_RANGE = 5
+
+
+import awq_inference_engine
+
+
+class QuantSiglipEncoder(nn.Module):
+    def __init__(self, module: SiglipEncoder, bsz=64, seqlen=1024):
+        super().__init__()
+        self.config = module.config
+        self.layers = [QuantSiglipEncoderLayer(layer) for layer in module.layers]
+        self.buffer = ActivationBuffer(module)
+        self.bsz = bsz
+        self.seqlen = seqlen
+        self.buffer.allocate_activation_buffer(self.bsz * self.seqlen)
+
+    # Ignore copy
+    def forward(
+        self,
+        inputs_embeds,
+        attention_mask: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,  # dummy
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutput]:
+        # TODO Find why this code is necessary
+        # torch.sum(inputs_embeds!=inputs_embeds)
+        bsz, seqlen, _ = inputs_embeds.shape
+        if self.bsz != bsz or self.seqlen != seqlen:
+            self.buffer.allocate_activation_buffer(bsz * seqlen)
+            self.bsz = bsz
+            self.seqlen = seqlen
+
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+
+        encoder_states = () if output_hidden_states else None
+
+        hidden_states = inputs_embeds
+        for i, encoder_layer in enumerate(self.layers):
+            if output_hidden_states:
+                encoder_states = encoder_states + (
+                    hidden_states.reshape(bsz, seqlen, -1),
+                )
+            hidden_states = encoder_layer(
+                hidden_states, self.buffer, attention_mask, bsz, seqlen
+            )
+
+        if output_hidden_states:
+            encoder_states = encoder_states + (hidden_states.reshape(bsz, seqlen, -1),)
+        if not return_dict:
+            return tuple(v for v in [hidden_states, encoder_states] if v is not None)
+        return BaseModelOutput(
+            last_hidden_state=hidden_states.reshape(bsz, seqlen, -1),
+            hidden_states=encoder_states,
+            attentions=None,
+        )
+
+
+class QuantSiglipMLP(nn.Module):
+    def __init__(self, siglipmlp, init_only=False):
+        super().__init__()
+        self.config = siglipmlp.config
+        self.activation_fn = siglipmlp.activation_fn
+        self.fc1 = W8A8OF16LinearDynamicInputScale.from_linear(
+            siglipmlp.fc1, init_only=init_only, fc1=False
+        )
+        self.fc2 = W8A8OF16LinearDynamicInputScale.from_linear(
+            siglipmlp.fc2, init_only=init_only
+        )
+        self.invoke_quant = self.invoke_quant_mlp
+
+    def invoke_quant_mlp(self, buffer, actfn_output):
+        awq_inference_engine.invoke_quant(
+            buffer.quantized_mlp_act_buffer,
+            actfn_output,
+            buffer.quantized_scale_buffer,
+        )
+
+    def forward(self, buffer: ActivationBuffer) -> torch.Tensor:
+        # INT8 in, FP16 out
+        self.fc1(
+            buffer.quantized_hidden_states_buffer,
+            buffer.quantized_scale_buffer,
+            buffer.fc1_buffer,
+        )
+        # Act & quantization
+        awq_inference_engine.gelu_and_quant(
+            buffer.quantized_mlp_act_buffer,
+            buffer.fc1_buffer,
+            buffer.quantized_scale_buffer,
+            buffer.tmp,
+        )
+        # INT8 in, FP16 out
+        self.fc2(
+            buffer.quantized_mlp_act_buffer,
+            buffer.quantized_scale_buffer,
+            buffer.in_out_fc2_act_buffer,
+        )
+
+
+class QuantSiglipFlashAttention2(nn.Module):
+    def __init__(
+        self,
+        module: SiglipAttention,
+        init_only=False,
+    ):
+        super().__init__()
+        self.config = module.config
+        self.embed_dim = module.embed_dim
+        self.num_heads = module.num_heads
+        self.head_dim = self.embed_dim // self.num_heads
+
+        self.qkv_proj = W8A8OF16LinearDynamicInputScale.from_qkv(
+            module.q_proj, module.k_proj, module.v_proj, init_only=init_only
+        )
+        self.out_proj = W8A8OF16LinearDynamicInputScale.from_linear(
+            module.out_proj, init_only=init_only
+        )
+        self.invoke_quant = self.invoke_quant_wo
+
+    def invoke_quant_wo(self, buffer, attn_output):
+        awq_inference_engine.invoke_quant(
+            buffer.quantized_hidden_states_buffer,
+            attn_output,
+            buffer.quantized_scale_buffer,
+        )
+
+    # Adapted from transformers.models.llama.modeling_llama.LlamaFlashAttention2.forward
+    def forward(
+        self, buffer: ActivationBuffer, bsz=64, seqlen=1024
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        # qkv
+        self.qkv_proj(
+            buffer.quantized_hidden_states_buffer,
+            buffer.quantized_scale_buffer,
+            buffer.qkv_proj_act_buffer,
+        )
+        q, k, v = buffer.qkv_proj_act_buffer.split(
+            [self.embed_dim, self.embed_dim, self.embed_dim], dim=-1
+        )
+        q = q.reshape(bsz, seqlen, self.num_heads, self.head_dim)
+        k = k.reshape(bsz, seqlen, self.num_heads, self.head_dim)
+        v = v.reshape(bsz, seqlen, self.num_heads, self.head_dim)
+        attn_output = flash_attn_func(q, k, v, softmax_scale=None, causal=False)
+        attn_output = attn_output.reshape(bsz * seqlen, -1)
+        # FP16 -> int8
+        self.invoke_quant(buffer, attn_output)
+        # INT8 in, FP16 out
+        self.out_proj(
+            buffer.quantized_hidden_states_buffer,
+            buffer.quantized_scale_buffer,
+            buffer.in_out_fc2_act_buffer,
+        )
+
+
+class QuantSiglipEncoderLayer(nn.Module):
+    def __init__(self, module: SiglipEncoderLayer):
+        super().__init__()
+        self.embed_dim = module.embed_dim
+        self.self_attn = QuantSiglipFlashAttention2(module.self_attn)
+        self.layer_norm1 = RMSNormGeneral(
+            module.layer_norm1.weight.data,
+            module.layer_norm1.bias.data,
+            module.layer_norm1.eps,
+            True,
+        ).cuda()
+        self.mlp = QuantSiglipMLP(module.mlp)
+        self.layer_norm2 = RMSNormGeneral(
+            module.layer_norm2.weight.data,
+            module.layer_norm2.bias.data,
+            module.layer_norm2.eps,
+            True,
+        ).cuda()
+        self.quant = self.invoke_quant_norm
+
+    def invoke_quant_norm(self, buffer, normfn_output):
+        awq_inference_engine.invoke_quant(
+            buffer.quantized_hidden_states_buffer,
+            normfn_output,
+            buffer.quantized_scale_buffer,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        buffer: ActivationBuffer,
+        attention_mask,
+        bsz,
+        seqlen,
+    ) -> Tuple[torch.FloatTensor]:
+        # Attention block
+        # FP16 in int8 out, layernorm & quantization
+        residual = hidden_states
+        self.layer_norm1(
+            hidden_states.reshape(-1, self.embed_dim),
+            buffer.quantized_hidden_states_buffer,
+            buffer.quantized_scale_buffer,
+        )
+
+        # INT8 -> FP16
+        self.self_attn(buffer, bsz, seqlen)
+        hidden_states = (
+            residual.reshape(-1, self.embed_dim) + buffer.in_out_fc2_act_buffer
+        )
+        # Fully Connected
+        residual = hidden_states
+        # FP16 in int8 out, layernorm & quantization
+        self.layer_norm2(
+            hidden_states.reshape(-1, self.embed_dim),
+            buffer.quantized_hidden_states_buffer,
+            buffer.quantized_scale_buffer,
+        )
+
+        # INT8 -> FP16
+        self.mlp(buffer)
+        hidden_states = (
+            residual.reshape(-1, self.embed_dim) + buffer.in_out_fc2_act_buffer
+        )
+        return hidden_states
+
+
+class RMSNormGeneral(nn.Module):
+    """Root mean square normalization (w/ per-token or per-tensor quant).
+
+    Computes x -> w * x / sqrt(E[x^2] + eps) where w is the learned weight.
+    Refer to https://arxiv.org/abs/1910.07467
+    """
+
+    def __init__(
+        self,
+        weight: torch.tensor,
+        bias: torch.tensor,
+        eps: float = 1e-6,
+        use_per_token_quant: bool = True,
+    ) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(weight, requires_grad=False)
+        self.bias = nn.Parameter(bias, requires_grad=False)
+        self.variance_epsilon = eps
+        self.use_per_token_quant = use_per_token_quant
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        quantized_hidden_states_buffer: torch.Tensor,
+        quantized_scale_buffer: torch.Tensor,
+        quantized_sum_buffer: torch.Tensor = None,
+    ) -> torch.Tensor:
+        # quantized_sum_buffer is not used, only to keep the consistency of the interface
+        awq_inference_engine.rms_norm_general(
+            quantized_hidden_states_buffer,
+            x,
+            self.weight.data,
+            self.bias.data,
+            quantized_scale_buffer,
+            self.variance_epsilon,
+            self.use_per_token_quant,
+        )

llm-awq/tinychat/modules/fused_vision_attn.py

@@ -0,0 +1,272 @@
+import math
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from typing import Optional, Tuple
+
+# from awq.quantize.qmodule import WQLinear
+# import awq_inference_engine
+# from tinychat.models.llama import apply_rotary_emb
+import gc
+
+import tinychat.utils.constants
+
+max_batch_size = tinychat.utils.constants.max_batch_size
+max_seq_len = tinychat.utils.constants.max_seq_len
+
+from transformers.activations import ACT2FN
+from transformers.models.clip.configuration_clip import (
+    CLIPConfig,
+    CLIPTextConfig,
+    CLIPVisionConfig,
+)
+from transformers.models.clip.modeling_clip import CLIPAttention
+
+
+class CLIPAttentionFused(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(
+        self, hidden_size, num_heads, qkv_proj, out_proj, dev, attention_dropout=0.0
+    ):
+        super().__init__()
+        self.embed_dim = hidden_size
+        self.num_heads = num_heads
+        self.head_dim = hidden_size // num_heads
+        self.scale = self.head_dim**-0.5
+        self.dropout = attention_dropout
+
+        if (self.head_dim * num_heads) != self.embed_dim:
+            raise ValueError(
+                f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
+                f" and `num_heads`: {num_heads})."
+            )
+        self.qkv_proj = qkv_proj
+        self.out_proj = out_proj
+
+    def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
+        return (
+            tensor.view(bsz, seq_len, self.num_heads, self.head_dim)
+            .transpose(1, 2)
+            .contiguous()
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        causal_attention_mask: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        """Input shape: Batch x Time x Channel"""
+
+        bsz, tgt_len, embed_dim = hidden_states.size()
+
+        qkv_states = self.qkv_proj(hidden_states)
+        qkv_states = qkv_states.view(bsz, tgt_len, 3, self.num_heads, self.head_dim)
+
+        # This updates the query and key states in-place, saving VRAM.
+        query_states, key_states, value_states = torch.split(qkv_states, 1, dim=2)
+        del qkv_states
+
+        proj_shape = (bsz * self.num_heads, -1, self.head_dim)
+
+        query_states = (
+            query_states.view(bsz, tgt_len, self.num_heads, self.head_dim)
+            .transpose(1, 2)
+            .view(*proj_shape)
+            * self.scale
+        )
+        key_states = (
+            key_states.view(bsz, tgt_len, self.num_heads, self.head_dim)
+            .transpose(1, 2)
+            .view(*proj_shape)
+        )
+        value_states = (
+            value_states.view(bsz, tgt_len, self.num_heads, self.head_dim)
+            .transpose(1, 2)
+            .view(*proj_shape)
+        )
+
+        src_len = key_states.size(1)
+        attn_weights = torch.bmm(query_states, key_states.transpose(1, 2))
+
+        if attn_weights.size() != (bsz * self.num_heads, tgt_len, src_len):
+            raise ValueError(
+                f"Attention weights should be of size {(bsz * self.num_heads, tgt_len, src_len)}, but is"
+                f" {attn_weights.size()}"
+            )
+
+        # apply the causal_attention_mask first
+        if causal_attention_mask is not None:
+            if causal_attention_mask.size() != (bsz, 1, tgt_len, src_len):
+                raise ValueError(
+                    f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is"
+                    f" {causal_attention_mask.size()}"
+                )
+            attn_weights = (
+                attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
+                + causal_attention_mask
+            )
+            attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
+
+        if attention_mask is not None:
+            if attention_mask.size() != (bsz, 1, tgt_len, src_len):
+                raise ValueError(
+                    f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is {attention_mask.size()}"
+                )
+            attn_weights = (
+                attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
+                + attention_mask
+            )
+            attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
+
+        attn_weights = nn.functional.softmax(attn_weights, dim=-1)
+
+        if output_attentions:
+            # this operation is a bit akward, but it's required to
+            # make sure that attn_weights keeps its gradient.
+            # In order to do so, attn_weights have to reshaped
+            # twice and have to be reused in the following
+            attn_weights_reshaped = attn_weights.view(
+                bsz, self.num_heads, tgt_len, src_len
+            )
+            attn_weights = attn_weights_reshaped.view(
+                bsz * self.num_heads, tgt_len, src_len
+            )
+        else:
+            attn_weights_reshaped = None
+
+        attn_probs = nn.functional.dropout(
+            attn_weights, p=self.dropout, training=self.training
+        )
+
+        attn_output = torch.bmm(attn_probs, value_states)
+
+        if attn_output.size() != (bsz * self.num_heads, tgt_len, self.head_dim):
+            raise ValueError(
+                f"`attn_output` should be of size {(bsz, self.num_heads, tgt_len, self.head_dim)}, but is"
+                f" {attn_output.size()}"
+            )
+
+        attn_output = attn_output.view(bsz, self.num_heads, tgt_len, self.head_dim)
+        attn_output = attn_output.transpose(1, 2)
+        attn_output = attn_output.reshape(bsz, tgt_len, embed_dim)
+
+        attn_output = self.out_proj(attn_output)
+
+        return attn_output, attn_weights_reshaped
+
+
+class CLIPMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.activation_fn = ACT2FN[config.hidden_act]
+        self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)
+        self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.fc1(hidden_states)
+        hidden_states = self.activation_fn(hidden_states)
+        hidden_states = self.fc2(hidden_states)
+        return hidden_states
+
+
+class CLIPEncoderLayer(nn.Module):
+    def __init__(self, config: CLIPConfig):
+        super().__init__()
+        self.embed_dim = config.hidden_size
+        self.self_attn = CLIPAttention(config)
+        self.layer_norm1 = nn.LayerNorm(self.embed_dim)
+        self.mlp = CLIPMLP(config)
+        self.layer_norm2 = nn.LayerNorm(self.embed_dim)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+        causal_attention_mask: torch.Tensor,
+        output_attentions: Optional[bool] = False,
+    ) -> Tuple[torch.FloatTensor]:
+        """
+        Args:
+            hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
+            attention_mask (`torch.FloatTensor`): attention mask of size
+                `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
+                `(config.encoder_attention_heads,)`.
+            output_attentions (`bool`, *optional*):
+                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+                returned tensors for more detail.
+        """
+        residual = hidden_states
+
+        hidden_states = self.layer_norm1(hidden_states)
+        hidden_states, attn_weights = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            causal_attention_mask=causal_attention_mask,
+            output_attentions=output_attentions,
+        )
+        hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.layer_norm2(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        outputs = (hidden_states,)
+
+        if output_attentions:
+            outputs += (attn_weights,)
+
+        return outputs
+
+
+def make_fused_vision_attn(model, dev):
+    """
+    Replace all LlamaAttention modules with QuantLlamaAttention modules, fusing the q, k, v projections.
+    """
+    model = model.cpu()
+    for name, m in model.named_modules():
+        if not m.__class__.__name__ in ["CLIPAttention", "CLIPAttentionFused"]:
+            continue
+
+        q_proj = m.q_proj
+        k_proj = m.k_proj
+        v_proj = m.v_proj
+
+        weights = torch.cat([q_proj.weight, k_proj.weight, v_proj.weight], dim=0)
+        bias = (
+            torch.cat([q_proj.bias, k_proj.bias, v_proj.bias], dim=0)
+            if q_proj.bias is not None
+            else None
+        )
+
+        qkv_layer = nn.Linear(
+            q_proj.in_features,
+            q_proj.out_features + k_proj.out_features + v_proj.out_features,
+            q_proj.bias is not None,
+            q_proj.weight.device,
+        )
+        qkv_layer.weight.data = weights
+
+        qkv_layer.bias.data = bias
+        if isinstance(m, CLIPAttention):
+            attn = CLIPAttentionFused(
+                m.embed_dim, m.num_heads, qkv_layer, m.out_proj, dev
+            )
+        if "." in name:
+            parent_name = name.rsplit(".", 1)[0]
+            child_name = name[len(parent_name) + 1 :]
+            parent = model.get_submodule(parent_name)
+        else:
+            parent_name = ""
+            parent = model
+            child_name = name
+
+        # print(f"Replacing {name} with quant_attn; parent: {parent_name}, child's name: {child_name}")
+        setattr(parent, child_name, attn)
+        gc.collect()
+        torch.cuda.empty_cache()
+    model = model.to(dev)

llm-awq/tinychat/nvila_benchmark.py

@@ -0,0 +1,163 @@
+import argparse
+
+from termcolor import colored
+
+import llava
+from llava import conversation as clib
+from llava.media import Image, Video
+import torch
+from awq.quantize import fake_quant
+from awq.quantize.quantizer import real_quantize_model_weight
+from transformers import AutoConfig
+import tinychat
+
+
+def skip(*args, **kwargs):
+    pass
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--model-path",
+        "-m",
+        type=str,
+        default="/home/yuming/workspace/qwen/models/nvila-internal-8b-v1",
+    )
+    parser.add_argument(
+        "--quant_path",
+        type=str,
+        default="/PATH/TO/QUANT",
+    )
+    # parser.add_argument("--model-path", "-m", type=str, default="Efficient-Large-Model/J65")
+    # parser.add_argument("--quant_path", type=str, default="/home/yuming/workspace/qwen/models/J65/llm/vila2-J65-w4-g128-awq-v2.pt")
+    parser.add_argument("--conv-mode", "-c", type=str, default="auto")
+    # parser.add_argument("--media", type=str, default="/home/yuming/workspace/space_woaudio.mp4")
+    parser.add_argument("--device", type=str, default="cuda:0")
+    parser.add_argument(
+        "--act_scale_path",
+        type=str,
+        default="/PATH/TO/SCALE",
+    )
+    # quantization options
+    parser.add_argument("--quant_llm", action="store_true")
+    parser.add_argument("--quant_VT", action="store_true")
+    # Four basic tasks
+    parser.add_argument("--video_caption", action="store_true")
+    parser.add_argument("--video_QA", action="store_true")
+    parser.add_argument("--image_caption", action="store_true")
+    parser.add_argument("--image_QA", action="store_true")
+
+    parser.add_argument(
+        "--all",
+        action="store_true",
+        help="Whether to quantize visiontower and llm, and test all 4 tasks",
+    )
+    parser.add_argument(
+        "--fakequant_VT",
+        action="store_true",
+        help="Use fake quant or real quant for VisionTower",
+    )
+    parser.add_argument(
+        "--all_task", action="store_true", help="Whether to test all 4 tasks"
+    )
+    parser.add_argument(
+        "--video_path", type=str, default="../figures/nvila_demo_video.mp4"
+    )
+    parser.add_argument("--image_path", type=str, default="../figures/vila-logo.jpg")
+    parser.add_argument("--max_seq_len", type=int, default=8192)
+    args = parser.parse_args()
+
+    torch.nn.init.kaiming_uniform_ = skip
+    torch.nn.init.kaiming_normal_ = skip
+    torch.nn.init.uniform_ = skip
+    torch.nn.init.normal_ = skip
+    import tinychat.utils.constants
+
+    tinychat.utils.constants.max_seq_len = args.max_seq_len
+    from transformers import modeling_utils
+
+    modeling_utils._init_weights = False
+
+    # Load model
+    from tinychat.models.nvila_qwen2 import NVILAQwen2
+
+    config = AutoConfig.from_pretrained(args.model_path)
+    config.resume_path = args.model_path
+    model = NVILAQwen2(config).half()
+    model.llm = model.llm.eval()
+    if args.quant_llm or args.all:
+        from tinychat.modules import (
+            make_quant_norm,
+            make_quant_attn,
+            make_fused_mlp,
+            make_fused_vision_attn,
+        )
+
+        real_quantize_model_weight(
+            model.llm,
+            w_bit=4,
+            q_config=dict(q_group_size=128, zero_point=True),
+            init_only=True,
+        )
+        make_quant_attn(model.llm, "cuda", True)
+        make_quant_norm(model.llm)
+        make_fused_mlp(model.llm)
+        model = model.to("cuda")
+    model = model.to(args.device)
+    if args.quant_VT or args.all:
+        from tinychat.modules import QuantSiglipEncoder
+
+        model.vision_tower.vision_tower.vision_model.encoder = QuantSiglipEncoder(
+            model.vision_tower.vision_tower.vision_model.encoder
+        )
+    model = model.cuda().eval()
+
+    if args.video_caption or args.all or args.all_task:
+        print("-" * 80)
+        print("Video_Caption")
+        # Set conversation mode
+        clib.default_conversation = clib.conv_templates[args.conv_mode].copy()
+        media = Video(args.video_path)
+        text = "Elaborate on the visual and narrative elements of the video in detail."  # + "1"+" 1"*3069
+        prompt = [media, text]
+        # Generate response
+        with torch.no_grad():
+            response = model.benchmark(prompt, args.quant_llm)
+    if args.video_QA or args.all or args.all_task:
+        print("-" * 80)
+        print("Video_QA")
+        # Set conversation mode
+        clib.default_conversation = clib.conv_templates[args.conv_mode].copy()
+        media = Video(args.video_path)
+        text = "What is the person in the video doing? Select the option that best describes their action: A. Folding paper B. Playing computer games C. Sleeping."  # + "1"+" 1"*3069
+        prompt = [media, text]
+        # Generate response
+        with torch.no_grad():
+            response = model.benchmark(prompt, args.quant_llm)
+    if args.image_caption or args.all or args.all_task:
+        print("-" * 80)
+        print("Image_Caption")
+        # Set conversation mode
+        clib.default_conversation = clib.conv_templates[args.conv_mode].copy()
+        media = Image(args.image_path)
+        text = "Describe the image in detail."
+        prompt = [media, text]
+        # Generate response
+        with torch.no_grad():
+            response = model.benchmark(prompt, args.quant_llm)
+    if args.image_QA or args.all or args.all_task:
+        print("-" * 80)
+        print("Image_QA")
+        # Set conversation mode
+        clib.default_conversation = clib.conv_templates[args.conv_mode].copy()
+        media = Image(args.image_path)
+        text = "What does the text in the image say? Choose the option that best matches: A. VILA B. AIIV C. ALIV."
+        prompt = [media, text]
+        # Generate response
+        with torch.no_grad():
+            response = model.benchmark(prompt, args.quant_llm)
+
+
+if __name__ == "__main__":
+    main()

llm-awq/tinychat/nvila_demo.py

@@ -0,0 +1,272 @@
+import argparse
+
+from termcolor import colored
+from huggingface_hub import hf_hub_download
+import os
+import llava
+from llava import conversation as clib
+from llava.media import Image, Video
+import torch
+from awq.quantize import fake_quant
+from transformers import AutoConfig
+from tinychat.utils.load_quant import load_awq_model
+from tinychat.utils.llava_image_processing import (
+    load_images,
+    vis_images,
+)
+
+
+def skip(*args, **kwargs):
+    pass
+
+
+from tinychat.utils.tune import (
+    device_warmup,
+    tune_all_wqlinears,
+    tune_llava_patch_embedding,
+)
+from tinychat.utils.prompt_templates import (
+    get_prompter,
+    get_stop_token_ids,
+    get_image_token,
+)
+from llava.utils.media import extract_media
+import tinychat.utils.constants
+from tinychat.stream_generators.NVILA_stream_gen import NVILAStreamGenerator
+from tinychat.utils.conversation_utils import gen_params, stream_output, TimeStats
+
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+
+def download_model_file(
+    repo_id: str = "Efficient-Large-Model/NVILA-AWQ",
+    filename: str = None,
+    local_dir: str = "./hf_cache",      
+    force_download: bool = False,         
+) -> str:
+    os.makedirs(local_dir, exist_ok=True)
+    local_path = os.path.join(local_dir, filename)
+    if force_download or not os.path.exists(local_path):
+        print(f"Downloading {filename} from {repo_id}...")
+        hf_hub_download(
+            repo_id=repo_id,
+            filename=filename,
+            local_dir=local_dir,
+            local_dir_use_symlinks=False,  
+            resume_download=True,       
+            force_download=force_download,
+        )
+        print(f"File saved to: {local_path}")
+    
+    return local_path
+
+
+
+def main(args):
+    # Accelerate model initialization
+    setattr(torch.nn.Linear, "reset_parameters", lambda self: None)
+    setattr(torch.nn.LayerNorm, "reset_parameters", lambda self: None)
+    torch.nn.init.kaiming_uniform_ = skip
+    torch.nn.init.kaiming_normal_ = skip
+    torch.nn.init.uniform_ = skip
+    torch.nn.init.normal_ = skip
+    tinychat.utils.constants.max_seq_len = args.max_seq_len
+
+    # Prepare model
+    from tinychat.models.nvila_qwen2 import NVILAQwen2
+    from transformers import AutoConfig
+    from tinychat.models.qwen2 import Qwen2ForCausalLM
+
+    config = AutoConfig.from_pretrained(args.model_path)
+    config.resume_path = args.model_path
+    if args.quant_llm or args.all:
+        model = NVILAQwen2(config, False).half()
+    else:
+        model = NVILAQwen2(config, True).half()
+
+    if args.smooth_VT or args.all:
+        from awq.quantize import smooth_lm
+        args.act_scale_path=download_model_file(filename=args.act_scale_path)
+        act_scales = torch.load(args.act_scale_path)
+        smooth_lm(model.vision_tower, act_scales, 0.3)
+    if args.quant_llm or args.all:
+        from tinychat.modules import (
+            make_quant_norm,
+            make_quant_attn,
+            make_fused_mlp,
+            make_fused_vision_attn,
+        )
+        args.quant_path=download_model_file(filename=args.quant_path)
+        model.llm = Qwen2ForCausalLM(model.llm_cfg).half()
+        model.llm = load_awq_model(model.llm, args.quant_path, 4, 128, args.device)
+        make_quant_attn(model.llm, args.device, True)
+        make_quant_norm(model.llm)
+        model.llm.cpu()
+        model.llm.resize_token_embeddings(len(model.tokenizer))
+
+    if args.quant_VT or args.all:
+        from tinychat.modules import QuantSiglipEncoder
+
+        if args.fakequant_VT:
+            fake_quant(model.vision_tower.vision_tower.vision_model.encoder)
+        else:
+            model.vision_tower.vision_tower.vision_model.encoder = QuantSiglipEncoder(
+                model.vision_tower.vision_tower.vision_model.encoder
+            )
+    model = model.cuda().eval()
+    device_warmup(args.device)
+    tune_llava_patch_embedding(model.vision_tower, device=args.device)
+
+    # Pre-prepare media
+    prompt = []
+    media_files = []
+    if args.media is not None:
+        for media in args.media or []:
+            if any(media.endswith(ext) for ext in [".jpg", ".jpeg", ".png"]):
+                media = Image(media)
+                media_files.append(media)
+                media_prompt = "<image>"
+            elif any(media.endswith(ext) for ext in [".mp4", ".mkv", ".webm"]):
+                media = Video(media)
+                media_files.append(media)
+                media_prompt = "<vila/video>"
+            else:
+                raise ValueError(f"Unsupported media type: {media}")
+            prompt.append(media)
+    media_num = len(media_files)
+    if args.vis_image:
+        print("=" * 50)
+        print("Input Image:")
+        vis_images(args.media)
+    conversation = [{"from": "human", "value": prompt}]
+    media, media_cfg = model.prepare_media(conversation)
+    # Prepare streaming
+    stream_generator = NVILAStreamGenerator
+    # Prepare prompt
+    if args.max_seq_len <= 1024:
+        short_prompt = True
+    else:
+        short_prompt = False
+    model_prompter = get_prompter(
+        args.model_type, args.model_path, short_prompt, args.empty_prompt
+    )
+    stop_token_ids = get_stop_token_ids(args.model_type, args.model_path)
+    count = 0
+
+    if args.empty_prompt:
+        input_indicator = "Input: "
+        output_indicator = "Generated: "
+    else:
+        input_indicator = "USER: "
+        output_indicator = "ASSISTANT: "
+
+    count = 0
+    model.eval()
+    time_stats = TimeStats()
+    start_pos = 0
+    while True:
+        # Get input from the user
+        print("=" * 50)
+        input_prompt = input(input_indicator)
+        print("-" * 50)
+        if input_prompt == "":
+            print("EXIT...")
+            time_stats.show()
+            break
+        if count == 0:  # Insert media here
+            if args.media is not None:
+                if media_prompt in input_prompt:
+                    input_prompt = input_prompt
+                else:
+                    input_prompt = media_prompt * media_num + input_prompt
+            model_prompter.insert_prompt(input_prompt)
+        else:
+            model_prompter.insert_prompt(input_prompt)
+            if args.chunk_prefilling:
+                media = None
+                media_cfg = None
+        output_stream = stream_generator(
+            model,
+            gen_params,
+            model_prompter.model_input,
+            media,
+            media_cfg,
+            start_pos,
+            device=args.device,
+            stop_token_ids=stop_token_ids,
+            chunk_prefilling=args.chunk_prefilling,
+            quant_llm=args.quant_llm or args.all,
+        )
+        print(output_indicator, end="", flush=True)
+        if count == 0:
+            outputs, total_tokens = stream_output(output_stream, time_stats)
+        else:
+            outputs, total_tokens = stream_output(output_stream)
+        if args.chunk_prefilling:
+            start_pos += total_tokens
+        if (
+            args.single_round is not True and args.max_seq_len > 512
+        ):  # Only memorize previous conversations when kv_cache_size > 512
+            model_prompter.update_template(outputs, args.chunk_prefilling)
+        count += 1
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--model_type", type=str, default="LLaMa", help="type of the model"
+    )
+    parser.add_argument(
+        "--model-path", type=str, default="/data/llm/checkpoints/llava/llava-v1.5-7b"
+    )
+    parser.add_argument(
+        "--quant_path",
+        type=str,
+        default="/data/llm/checkpoints/llava/llava-v1.5-7b-w4-g128-awq.pt",
+    )
+    parser.add_argument(
+        "--act_scale_path",
+        type=str,
+        default="/PATH/TO/SCALE",
+    )
+    parser.add_argument(
+        "--media", type=str, nargs="+", help="Multi-modal input (Video or image path)"
+    )
+    parser.add_argument("--device", type=str, default="cuda:0")
+    parser.add_argument("--max_seq_len", type=int, default=2048)
+    parser.add_argument(
+        "--single_round",
+        action="store_true",
+        help="whether to memorize previous conversations",
+    )
+    parser.add_argument(
+        "--vis-image",
+        action="store_true",
+        help="whether to visualize the image while chatting",
+    )
+    parser.add_argument(
+        "--empty-prompt",
+        action="store_true",
+        help="whether to use empty prompt template",
+    )
+    parser.add_argument(
+        "--flash_attn",
+        action="store_true",
+        help="whether to use flash attention",
+    )
+    parser.add_argument(
+        "--chunk_prefilling",
+        action="store_true",
+        help="If used, in context stage, the history tokens will not be recalculated, greatly speeding up the calculation",
+    )
+    # smooth and quantization options
+    parser.add_argument("--quant_llm", action="store_true")
+    parser.add_argument("--quant_VT", action="store_true")
+    parser.add_argument("--smooth_VT", action="store_true")
+    parser.add_argument("--all", action="store_true")
+    parser.add_argument(
+        "--fakequant_VT",
+        action="store_true",
+        help="Use fake quant or real quant for VisionTower",
+    )
+    args = parser.parse_args()
+    main(args)

llm-awq/tinychat/offline-weight-repacker.py

@@ -0,0 +1,157 @@
+import argparse
+import torch
+import numpy as np
+from typing import List
+from collections import OrderedDict
+
+
+def qweight_unpack(qweight):
+    assert qweight.dtype == torch.int32
+    n = qweight.shape[0]
+    k = qweight.shape[1] * 8
+    unpacked_qweight = torch.zeros((n, k), dtype=torch.int32, device=qweight.device)
+    mask = 0x0000000F
+    for kk in range(k):
+        ele_offset = kk // 8
+        bit_offset = (kk % 8) * 4
+        unpacked_qweight[:, kk] = (qweight[:, ele_offset] >> bit_offset) & mask
+
+    return unpacked_qweight
+
+
+def packing_v2_from_unpacked(unpacked_qweight, interleave, kstride):
+    # unpacked_qweight: [N, K]
+    N = unpacked_qweight.shape[0]
+    K = unpacked_qweight.shape[1]
+
+    Packed_Kernel = unpacked_qweight.cpu().numpy().reshape(N, K // 32, 32)
+    # np.arange(32).reshape(4, 4, 2).transpose(1, 0, 2) => [0, 1, 8, 9, 16, 17, 24, 25, ...]
+    Packed_Kernel = Packed_Kernel.reshape(N, K // 32, 4, 4, 2).transpose(0, 1, 3, 2, 4)
+    Packed_Kernel = Packed_Kernel.reshape(N, K // 32, 32)
+
+    # reorder each 8 weights for fast dequantization
+    # [0, 1, 2, 3, 4, 5, 6, 7] => [0, 2, 4, 6, 1, 3, 5, 7]
+    Packed_Kernel = Packed_Kernel.reshape(N, K // 32, 4, 8)
+    Packed_Kernel = Packed_Kernel.reshape(N, K // 32, 4, 4, 2).transpose(0, 1, 2, 4, 3)
+    Packed_Kernel = Packed_Kernel.reshape(N, K)
+
+    # interleaving every four rows
+    Packed_Kernel = Packed_Kernel.reshape(
+        N // interleave, interleave, K // kstride, kstride
+    )
+    # N // 4, K // 64, 4, 64
+    Packed_Kernel = Packed_Kernel.transpose(0, 2, 1, 3)
+    Packed_Kernel = Packed_Kernel.reshape(
+        N // interleave, K // kstride, kstride, interleave
+    )
+    # Packing -> (N // 4, K // 64, 64)
+    Packed_Kernel = (
+        Packed_Kernel[..., 0]
+        | (Packed_Kernel[..., 1] << 4)
+        | (Packed_Kernel[..., 2] << 8)
+        | (Packed_Kernel[..., 3] << 12)
+    )
+    # reshape to (N // 4, K), FP16 format
+    Packed_Kernel = Packed_Kernel.reshape(N // interleave, K)
+    qweight_v2 = (
+        torch.tensor(Packed_Kernel.astype("int16"))
+        .to(unpacked_qweight.device)
+        .contiguous()
+    )
+    return qweight_v2
+
+
+def multiply_scale_qzero_negative(scales, qzeros, zp_shift=-8):
+    pack_size = 8
+    k_groups = scales.shape[1]
+    scaled_zeros = torch.zeros_like(scales)
+    for group_idx in range(k_groups):
+        zero_idx = group_idx // pack_size
+        zero_offset = group_idx % pack_size
+        zero = qzeros[:, zero_idx] >> (4 * zero_offset) & 0x0000000F
+        scaled_zeros[:, group_idx] = scales[:, group_idx] * zero.to(scales.dtype)
+    return -(scaled_zeros + (zp_shift * scales))
+
+
+def qweight_pack_v1_to_v2(qweight, interleave, kstride):
+    unpacked_qweight = qweight_unpack(qweight)
+    qweight_v2 = packing_v2_from_unpacked(unpacked_qweight, interleave, kstride)
+    return qweight_v2
+
+
+def ckpt_check():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--input1", type=str, default="./vicuna-7b-w4-g128-awq-v2-1.pt")
+    parser.add_argument("--input2", type=str, default="./vicuna-7b-w4-g128-awq-v2-2.pt")
+    args = parser.parse_args()
+
+    model_dict1 = torch.load(args.input1)
+    model_dict2 = torch.load(args.input2)
+
+    keys = model_dict1.keys()
+    for key in keys:
+        param = model_dict1[key]
+        assert type(param) == torch.Tensor
+        if (
+            "qweight" in key
+            or "scales" in key
+            or "qzeros" in key
+            or "scaled_zeros" in key
+        ):
+            print("=" * 50)
+            print(key)
+            # print(model_dict1[key])
+            # print(model_dict2[key])
+            diff = torch.max(torch.abs(model_dict2[key] - model_dict1[key]))
+            print(diff)
+            assert diff < 1e-6
+            print("=" * 50)
+
+
+def offline_repacker():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--input", type=str, default="./vicuna-7b-w4-g128-awq.pt")
+    parser.add_argument("--output", type=str, default="./vicuna-7b-w4-g128-awq-v2.pt")
+    args = parser.parse_args()
+
+    model_dict = torch.load(args.input)
+    model_dict_v2 = OrderedDict()
+
+    keys = model_dict.keys()
+    for key in keys:
+        param = model_dict[key]
+        assert type(param) == torch.Tensor
+        if "qweight" in key:
+            print("repacking:", key)
+            qweight = param
+            qweight_v2 = qweight_pack_v1_to_v2(qweight, interleave=4, kstride=64)
+            model_dict_v2[key] = qweight_v2
+        elif "scales" in key:
+            print("repacking:", key)
+            scales = param
+            # print(scales.shape)
+            scales_v2 = scales.transpose(1, 0).contiguous()
+            model_dict_v2[key] = scales_v2
+
+            # deal with corresponding zero points
+            zeros_key = key.replace("scales", "qzeros")
+            print("repacking:", zeros_key)
+
+            zeros_key_v2 = key.replace("scales", "scaled_zeros")
+            qzeros = model_dict[zeros_key]
+            scaled_zeros_v2 = multiply_scale_qzero_negative(scales, qzeros, zp_shift=0)
+            # K // G, N
+            scaled_zeros_v2 = scaled_zeros_v2.transpose(1, 0).contiguous()
+            model_dict_v2[zeros_key_v2] = scaled_zeros_v2
+        elif "qzeros" in key:
+            pass
+        else:
+            print("copying:", key)
+            model_dict_v2[key] = param
+
+    torch.save(model_dict_v2, args.output)
+
+
+if __name__ == "__main__":
+    offline_repacker()
+    # ckpt_check()

llm-awq/tinychat/scripts/internvl_demo.sh

@@ -0,0 +1,18 @@
+MODEL_PATH=PATH_TO_INTERNVL
+MODEL_NAME=InternVL3-8B
+
+# run AWQ search
+python -m awq.entry --model_path $MODEL_PATH \
+    --w_bit 4 --q_group_size 128 \
+    --run_awq --dump_awq awq_cache/$MODEL_NAME-w4-g128.pt
+
+# generate real quantized weights (w4)
+python -m awq.entry --model_path $MODEL_PATH \
+    --w_bit 4 --q_group_size 128 --load_awq awq_cache/$MODEL_NAME-w4-g128.pt \
+    --q_backend real --dump_quant quant_cache/$MODEL_NAME-w4-128-awq.pt
+
+# Run the TinyChat demo:
+PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True python internvl_demo.py --model-path $MODEL_PATH \
+    --quant_path quant_cache/$MODEL_NAME-w4-128-awq-v2.pt \
+    --media ../figures/vila-logo.jpg --max_seq_len 4096 --chunk \
+    --model_type internvl3 --quant_VT --quant_llm

llm-awq/tinychat/scripts/llama2_demo.sh

@@ -0,0 +1,31 @@
+MODEL_PATH=/data/llm/checkpoints/llama2-hf
+MODEL_NAME=llama-2-7b-chat
+
+# # Perform AWQ search and save search results (we already did it for you):
+# mkdir -p awq_cache
+# python -m awq.entry --model_path $MODEL_PATH/$MODEL_NAME \
+#     --w_bit 4 --q_group_size 128 \
+#     --run_awq --dump_awq awq_cache/llama-2-7b-chat-w4-g128.pt
+
+# Generate real quantized weights (INT4):
+mkdir -p quant_cache
+python -m awq.entry --model_path $MODEL_PATH/$MODEL_NAME \
+    --w_bit 4 --q_group_size 128 \
+    --load_awq awq_cache/llama-2-7b-chat-w4-g128.pt \
+    --q_backend real --dump_quant quant_cache/llama-2-7b-chat-w4-g128-awq.pt
+
+# Run the TinyChat demo:
+python demo.py --model_type llama \
+    --model_path $MODEL_PATH/$MODEL_NAME \
+    --q_group_size 128 --load_quant quant_cache/llama-2-7b-chat-w4-g128-awq.pt \
+    --precision W4A16
+
+# Split checkpoint into shards for mem-efficient loading:
+python split_ckpt.py --input_path quant_cache/llama-2-7b-chat-w4-g128-awq.pt \
+    --output_path quant_cache/llama-2-7b-chat-w4-g128-awq
+
+# Run the TinyChat demo in mem_efficient_load mode:
+python demo.py --model_type llama \
+    --model_path $MODEL_PATH/$MODEL_NAME \
+    --q_group_size 128 --load_quant quant_cache/llama-2-7b-chat-w4-g128-awq \
+    --precision W4A16 --mem_efficient_load --flash --chunk_prefilling 

llm-awq/tinychat/scripts/nvila_demo.sh

@@ -0,0 +1,22 @@
+MODEL_PATH=PATH_TO_NVILA
+MODEL_NAME=NVILA-8B
+
+# run AWQ search
+python -m awq.entry --model_path $MODEL_PATH \
+    --smooth_scale --media_path https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2-VL/space_woaudio.mp4  \
+    --act_scale_path awq_cache/$MODEL_NAME-smooth-scale.pt --vila-20 \
+    --w_bit 4 --q_group_size 128 \
+    --run_awq --dump_awq awq_cache/$MODEL_NAME.pt
+
+# generate real quantized weights (w4)
+python -m awq.entry --model_path $MODEL_PATH/llm \
+    --w_bit 4 --q_group_size 128  \
+    --load_awq awq_cache/$MODEL_NAME.pt \
+    --q_backend real --dump_quant quant_cache/$MODEL_NAME-w4-g128-awq.pt --vila-20
+
+# Run the TinyChat demo:
+python nvila_demo.py --model-path $MODEL_PATH    \
+       --quant_path quant_cache/$MODEL_NAME-w4-g128-awq.pt      \
+       --media ../figures/nvila-logo.jpg     \
+       --act_scale_path awq_cache/$MODEL_NAME-smooth-scale.pt    \
+       --all --chunk --model_type nvila --vis_image

llm-awq/tinychat/serve/README.md

@@ -0,0 +1,27 @@
+## Gradio demo: VILA with TinyChat
+
+We provide scripts for building your own gradio server to run VILA models with TinyChat. Please run the following commands to launch the server.
+
+#### Launch a controller
+```bash
+python -m tinychat.serve.controller --host 0.0.0.0 --port 10000
+```
+
+#### Launch gradio web server.
+```bash
+python -m tinychat.serve.gradio_web_server --controller http://localhost:10000 --model-list-mode reload --share --auto-pad-image-token
+```
+After launching this script, the web interface will be served on your machine and you can access it with a public URL (or localhost URL).
+
+#### Launch a model worker
+
+```bash
+python -m tinychat.serve.model_worker_new --host 0.0.0.0 --controller http://localhost:10000 --port 40000 --worker http://localhost:40000 --model-path <path-to-fp16-hf-model> --quant-path <path-to-awq-checkpoint>
+# Please change tinychat.serve.model_worker_new to tinychat.serve.model_worker if you want to serve VILA rather than VILA-1.5
+```
+
+Note: You can launch multiple model workers onto the same web server. And please remember to specify different ports for each model worker.
+
+### Acknowlegement
+
+This demo is inspired by [LLaVA](https://github.com/haotian-liu/LLaVA). We thank LLaVA for providing an elegant way to build the Gradio Web UI.

llm-awq/tinychat/serve/controller.py

@@ -0,0 +1,325 @@
+#    Modified from https://github.com/haotian-liu/LLaVA
+#    Copyright 2023 Haotian Liu
+#
+#    Licensed under the Apache License, Version 2.0 (the "License");
+#    you may not use this file except in compliance with the License.
+#    You may obtain a copy of the License at
+#
+#        http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS,
+#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#    See the License for the specific language governing permissions and
+#    limitations under the License.
+
+"""
+A controller manages distributed workers.
+It sends worker addresses to clients.
+"""
+import argparse
+import asyncio
+import dataclasses
+from enum import Enum, auto
+import json
+import logging
+import time
+from typing import List, Union
+import threading
+
+from fastapi import FastAPI, Request
+from fastapi.responses import StreamingResponse
+import numpy as np
+import requests
+import uvicorn
+
+from tinychat.utils.constants import CONTROLLER_HEART_BEAT_EXPIRATION
+from tinychat.utils.log_utils import build_logger, server_error_msg
+
+
+logger = build_logger("controller", "controller.log")
+
+
+class DispatchMethod(Enum):
+    LOTTERY = auto()
+    SHORTEST_QUEUE = auto()
+
+    @classmethod
+    def from_str(cls, name):
+        if name == "lottery":
+            return cls.LOTTERY
+        elif name == "shortest_queue":
+            return cls.SHORTEST_QUEUE
+        else:
+            raise ValueError(f"Invalid dispatch method")
+
+
+@dataclasses.dataclass
+class WorkerInfo:
+    model_names: List[str]
+    speed: int
+    queue_length: int
+    check_heart_beat: bool
+    last_heart_beat: str
+
+
+def heart_beat_controller(controller):
+    while True:
+        time.sleep(CONTROLLER_HEART_BEAT_EXPIRATION)
+        controller.remove_stable_workers_by_expiration()
+
+
+class Controller:
+    def __init__(self, dispatch_method: str):
+        # Dict[str -> WorkerInfo]
+        self.worker_info = {}
+        self.dispatch_method = DispatchMethod.from_str(dispatch_method)
+
+        self.heart_beat_thread = threading.Thread(
+            target=heart_beat_controller, args=(self,)
+        )
+        self.heart_beat_thread.start()
+
+        logger.info("Init controller")
+
+    def register_worker(
+        self, worker_name: str, check_heart_beat: bool, worker_status: dict
+    ):
+        if worker_name not in self.worker_info:
+            logger.info(f"Register a new worker: {worker_name}")
+        else:
+            logger.info(f"Register an existing worker: {worker_name}")
+
+        if not worker_status:
+            worker_status = self.get_worker_status(worker_name)
+        if not worker_status:
+            return False
+
+        self.worker_info[worker_name] = WorkerInfo(
+            worker_status["model_names"],
+            worker_status["speed"],
+            worker_status["queue_length"],
+            check_heart_beat,
+            time.time(),
+        )
+
+        logger.info(f"Register done: {worker_name}, {worker_status}")
+        return True
+
+    def get_worker_status(self, worker_name: str):
+        try:
+            r = requests.post(worker_name + "/worker_get_status", timeout=5)
+        except requests.exceptions.RequestException as e:
+            logger.error(f"Get status fails: {worker_name}, {e}")
+            return None
+
+        if r.status_code != 200:
+            logger.error(f"Get status fails: {worker_name}, {r}")
+            return None
+
+        return r.json()
+
+    def remove_worker(self, worker_name: str):
+        del self.worker_info[worker_name]
+
+    def refresh_all_workers(self):
+        old_info = dict(self.worker_info)
+        self.worker_info = {}
+
+        for w_name, w_info in old_info.items():
+            if not self.register_worker(w_name, w_info.check_heart_beat, None):
+                logger.info(f"Remove stale worker: {w_name}")
+
+    def list_models(self):
+        model_names = set()
+
+        for w_name, w_info in self.worker_info.items():
+            model_names.update(w_info.model_names)
+
+        return list(model_names)
+
+    def get_worker_address(self, model_name: str):
+        if self.dispatch_method == DispatchMethod.LOTTERY:
+            worker_names = []
+            worker_speeds = []
+            for w_name, w_info in self.worker_info.items():
+                if model_name in w_info.model_names:
+                    worker_names.append(w_name)
+                    worker_speeds.append(w_info.speed)
+            worker_speeds = np.array(worker_speeds, dtype=np.float32)
+            norm = np.sum(worker_speeds)
+            if norm < 1e-4:
+                return ""
+            worker_speeds = worker_speeds / norm
+            if True:  # Directly return address
+                pt = np.random.choice(np.arange(len(worker_names)), p=worker_speeds)
+                worker_name = worker_names[pt]
+                return worker_name
+
+            # Check status before returning
+            while True:
+                pt = np.random.choice(np.arange(len(worker_names)), p=worker_speeds)
+                worker_name = worker_names[pt]
+
+                if self.get_worker_status(worker_name):
+                    break
+                else:
+                    self.remove_worker(worker_name)
+                    worker_speeds[pt] = 0
+                    norm = np.sum(worker_speeds)
+                    if norm < 1e-4:
+                        return ""
+                    worker_speeds = worker_speeds / norm
+                    continue
+            return worker_name
+        elif self.dispatch_method == DispatchMethod.SHORTEST_QUEUE:
+            worker_names = []
+            worker_qlen = []
+            for w_name, w_info in self.worker_info.items():
+                if model_name in w_info.model_names:
+                    worker_names.append(w_name)
+                    worker_qlen.append(w_info.queue_length / w_info.speed)
+            if len(worker_names) == 0:
+                return ""
+            min_index = np.argmin(worker_qlen)
+            w_name = worker_names[min_index]
+            self.worker_info[w_name].queue_length += 1
+            logger.info(
+                f"names: {worker_names}, queue_lens: {worker_qlen}, ret: {w_name}"
+            )
+            return w_name
+        else:
+            raise ValueError(f"Invalid dispatch method: {self.dispatch_method}")
+
+    def receive_heart_beat(self, worker_name: str, queue_length: int):
+        if worker_name not in self.worker_info:
+            logger.info(f"Receive unknown heart beat. {worker_name}")
+            return False
+
+        self.worker_info[worker_name].queue_length = queue_length
+        self.worker_info[worker_name].last_heart_beat = time.time()
+        logger.info(f"Receive heart beat. {worker_name}")
+        return True
+
+    def remove_stable_workers_by_expiration(self):
+        expire = time.time() - CONTROLLER_HEART_BEAT_EXPIRATION
+        to_delete = []
+        for worker_name, w_info in self.worker_info.items():
+            if w_info.check_heart_beat and w_info.last_heart_beat < expire:
+                to_delete.append(worker_name)
+
+        for worker_name in to_delete:
+            self.remove_worker(worker_name)
+
+    def worker_api_generate_stream(self, params):
+        worker_addr = self.get_worker_address(params["model"])
+        if not worker_addr:
+            logger.info(f"no worker: {params['model']}")
+            ret = {
+                "text": server_error_msg,
+                "error_code": 2,
+            }
+            yield json.dumps(ret).encode() + b"\0"
+
+        try:
+            response = requests.post(
+                worker_addr + "/worker_generate_stream",
+                json=params,
+                stream=True,
+                timeout=5,
+            )
+            for chunk in response.iter_lines(decode_unicode=False, delimiter=b"\0"):
+                if chunk:
+                    yield chunk + b"\0"
+        except requests.exceptions.RequestException as e:
+            logger.info(f"worker timeout: {worker_addr}")
+            ret = {
+                "text": server_error_msg,
+                "error_code": 3,
+            }
+            yield json.dumps(ret).encode() + b"\0"
+
+    # Let the controller act as a worker to achieve hierarchical
+    # management. This can be used to connect isolated sub networks.
+    def worker_api_get_status(self):
+        model_names = set()
+        speed = 0
+        queue_length = 0
+
+        for w_name in self.worker_info:
+            worker_status = self.get_worker_status(w_name)
+            if worker_status is not None:
+                model_names.update(worker_status["model_names"])
+                speed += worker_status["speed"]
+                queue_length += worker_status["queue_length"]
+
+        return {
+            "model_names": list(model_names),
+            "speed": speed,
+            "queue_length": queue_length,
+        }
+
+
+app = FastAPI()
+
+
+@app.post("/register_worker")
+async def register_worker(request: Request):
+    data = await request.json()
+    controller.register_worker(
+        data["worker_name"], data["check_heart_beat"], data.get("worker_status", None)
+    )
+
+
+@app.post("/refresh_all_workers")
+async def refresh_all_workers():
+    models = controller.refresh_all_workers()
+
+
+@app.post("/list_models")
+async def list_models():
+    models = controller.list_models()
+    return {"models": models}
+
+
+@app.post("/get_worker_address")
+async def get_worker_address(request: Request):
+    data = await request.json()
+    addr = controller.get_worker_address(data["model"])
+    return {"address": addr}
+
+
+@app.post("/receive_heart_beat")
+async def receive_heart_beat(request: Request):
+    data = await request.json()
+    exist = controller.receive_heart_beat(data["worker_name"], data["queue_length"])
+    return {"exist": exist}
+
+
+@app.post("/worker_generate_stream")
+async def worker_api_generate_stream(request: Request):
+    params = await request.json()
+    generator = controller.worker_api_generate_stream(params)
+    return StreamingResponse(generator)
+
+
+@app.post("/worker_get_status")
+async def worker_api_get_status(request: Request):
+    return controller.worker_api_get_status()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--host", type=str, default="localhost")
+    parser.add_argument("--port", type=int, default=21001)
+    parser.add_argument(
+        "--dispatch-method",
+        type=str,
+        choices=["lottery", "shortest_queue"],
+        default="shortest_queue",
+    )
+    args = parser.parse_args()
+    logger.info(f"args: {args}")
+
+    controller = Controller(args.dispatch_method)
+    uvicorn.run(app, host=args.host, port=args.port, log_level="info")

llm-awq/tinychat/serve/gradio_web_server.py

@@ -0,0 +1,1200 @@
+#    Modified from https://github.com/haotian-liu/LLaVA
+#    Copyright 2023 Haotian Liu
+#
+#    Licensed under the Apache License, Version 2.0 (the "License");
+#    you may not use this file except in compliance with the License.
+#    You may obtain a copy of the License at
+#
+#        http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS,
+#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#    See the License for the specific language governing permissions and
+#    limitations under the License.
+
+import argparse
+import datetime
+import json
+import os
+import time
+
+import gradio as gr
+import requests
+
+LOGDIR = "."
+
+from tinychat.serve.llava_conv import (
+    default_conversation,
+    conv_templates,
+    get_conversation,
+    SeparatorStyle,
+)
+from tinychat.utils.log_utils import (
+    build_logger,
+    server_error_msg,
+    violates_moderation,
+    moderation_msg,
+)
+import hashlib
+
+IMAGE_BOX_NUM = 3
+BUTTON_LIST_LEN = 2
+
+logger = build_logger("gradio_web_server", "gradio_web_server.log")
+
+headers = {"User-Agent": "TinyChat AWQ Chatbot"}
+
+no_change_btn = gr.Button.update()
+enable_btn = gr.Button.update(interactive=True)
+disable_btn = gr.Button.update(interactive=False)
+
+from tinychat.utils.constants import (
+    LLAVA_DEFAULT_IMAGE_TOKEN,
+    LLAVA_DEFAULT_IM_TOKEN_PLACE_HOLDER,
+    AUTO_FILL_IM_TOKEN_HOLDER,
+)
+
+# IMAGE_TOKEN_VIS = "**[IMAGE]**"
+IMAGE_TOKEN_VIS = "**\<image\>**"
+
+priority = {
+    "vicuna-13b": "aaaaaaa",
+    "koala-13b": "aaaaaab",
+}
+
+
+def get_conv_log_filename():
+    t = datetime.datetime.now()
+    name = os.path.join(LOGDIR, f"{t.year}-{t.month:02d}-{t.day:02d}-conv.json")
+    return name
+
+
+def get_model_list():
+    ret = requests.post(args.controller_url + "/refresh_all_workers")
+    assert ret.status_code == 200
+    ret = requests.post(args.controller_url + "/list_models")
+    models = ret.json()["models"]
+    models.sort(key=lambda x: priority.get(x, x))
+    logger.info(f"Models: {models}")
+    return models
+
+
+get_window_url_params = """
+function() {
+    const params = new URLSearchParams(window.location.search);
+    url_params = Object.fromEntries(params);
+    console.log(url_params);
+    return url_params;
+    }
+"""
+
+
+def load_demo(url_params, prompt_style_btn, request: gr.Request):
+    logger.info(f"load_demo. ip: {request.client.host}. params: {url_params}")
+
+    dropdown_update = gr.Dropdown.update(visible=True)
+    if "model" in url_params:
+        model = url_params["model"]
+        if model in models:
+            dropdown_update = gr.Dropdown.update(value=model, visible=True)
+    state = get_conversation(prompt_style_btn)
+    # state = default_conversation.copy()
+    return state, dropdown_update
+
+
+def load_demo_refresh_model_list(prompt_style_btn, request: gr.Request):
+    logger.info(f"load_demo. ip: {request.client.host}")
+    models = get_model_list()
+    state = get_conversation(prompt_style_btn)
+    # state = default_conversation.copy()
+    dropdown_update = gr.Dropdown.update(
+        choices=models, value=models[0] if len(models) > 0 else ""
+    )
+    return state, dropdown_update
+
+
+# def vote_last_response(state, vote_type, model_selector, request: gr.Request):
+#     with open(get_conv_log_filename(), "a") as fout:
+#         data = {
+#             "tstamp": round(time.time(), 4),
+#             "type": vote_type,
+#             "model": model_selector,
+#             "state": state.dict(),
+#             "ip": request.client.host,
+#         }
+#         fout.write(json.dumps(data) + "\n")
+
+
+# def upvote_last_response(state, model_selector, request: gr.Request):
+#     logger.info(f"upvote. ip: {request.client.host}")
+#     vote_last_response(state, "upvote", model_selector, request)
+#     return ("",) + (disable_btn,) * 3
+
+
+# def downvote_last_response(state, model_selector, request: gr.Request):
+#     logger.info(f"downvote. ip: {request.client.host}")
+#     vote_last_response(state, "downvote", model_selector, request)
+#     return ("",) + (disable_btn,) * 3
+
+
+# def flag_last_response(state, model_selector, request: gr.Request):
+#     logger.info(f"flag. ip: {request.client.host}")
+#     vote_last_response(state, "flag", model_selector, request)
+#     return ("",) + (disable_btn,) * 3
+
+
+def regenerate(state, image_process_mode, request: gr.Request):
+    logger.info(f"regenerate. ip: {request.client.host}")
+    state.messages[-1][-1] = None
+    prev_human_msg = state.messages[-2]
+    if type(prev_human_msg[1]) in (tuple, list):
+        prev_human_msg[1] = (*prev_human_msg[1][:2], image_process_mode)
+    state.skip_next = False
+    return (state, state.to_gradio_chatbot(), "") + (disable_btn,) * BUTTON_LIST_LEN
+
+
+def change_prompt_style(state, prompt_style_btn, request: gr.Request):
+    if state.version != prompt_style_btn:
+        state = get_conversation(prompt_style_btn)
+    return state
+
+
+def clear_history(prompt_style_btn, request: gr.Request):
+    logger.info(f"clear_history. ip: {request.client.host}")
+    state = get_conversation(prompt_style_btn)
+    return (
+        (state, state.to_gradio_chatbot(), "")
+        + (None,) * IMAGE_BOX_NUM
+        + (None,)  # Videobox
+        + (disable_btn,) * BUTTON_LIST_LEN
+    )
+
+
+def clear_text_history(state, prompt_style_btn, request: gr.Request):
+    state = get_conversation(prompt_style_btn)
+    return (state, state.to_gradio_chatbot())
+
+
+def clear_after_click_example_1_video(videobox, textbox):
+    imagebox = None
+    imagebox_2 = None
+    imagebox_3 = None
+    state = get_conversation("default")
+    prompt_style_btn = "default"
+    return (state, imagebox, imagebox_2, imagebox_3, videobox, prompt_style_btn)
+
+
+def clear_after_click_example_1_image(imagebox, textbox):
+    imagebox_2 = None
+    imagebox_3 = None
+    videobox = None
+    state = get_conversation("default")
+    prompt_style_btn = "default"
+    return (state, imagebox, imagebox_2, imagebox_3, videobox, prompt_style_btn)
+
+
+def clear_after_click_example_2_image(imagebox, imagebox_2, textbox):
+    imagebox_3 = None
+    videobox = None
+    state = get_conversation("default")
+    prompt_style_btn = "default"
+    return (state, imagebox, imagebox_2, imagebox_3, videobox, prompt_style_btn)
+
+
+def clear_after_click_example_3_image(imagebox, imagebox_2, imagebox_3, textbox):
+    videobox = None
+    state = get_conversation("default")
+    prompt_style_btn = "default"
+    return (state, imagebox, imagebox_2, imagebox_3, videobox, prompt_style_btn)
+
+
+def clear_after_click_example_3_image_icl(imagebox, imagebox_2, imagebox_3, textbox):
+    videobox = None
+    state = get_conversation("no-sys")
+    prompt_style_btn = "no-sys"
+    return (state, imagebox, imagebox_2, imagebox_3, videobox, prompt_style_btn)
+
+
+def add_images(
+    state,
+    imagebox,
+    imagebox_2,
+    imagebox_3,
+    videobox,
+    image_process_mode,
+    request: gr.Request,
+):
+    if state.image_loaded:
+        # return (state,) + (None,) * IMAGE_BOX_NUM
+        return state
+
+    def extract_frames(video_path):
+        import cv2
+        from PIL import Image
+
+        vidcap = cv2.VideoCapture(video_path)
+        fps = vidcap.get(cv2.CAP_PROP_FPS)
+        frame_count = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))
+        duration = frame_count / fps
+
+        frame_interval = frame_count // 8
+        print(
+            "duration:", duration, "frames:", frame_count, "intervals:", frame_interval
+        )
+        # frame_interval = 10
+
+        def get_frame(max_frames):
+            # frame_id = int(fps * stamp)
+            # vidcap.set(cv2.CAP_PROP_POS_FRAMES, frame_id)
+            # ret, frame = vidcap.read()
+            images = []
+            count = 0
+            success = True
+            while success:
+                success, frame = vidcap.read()
+                if count % frame_interval == 0:
+                    img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+                    im_pil = Image.fromarray(img)
+                    images.append(im_pil)
+                    if len(images) == max_frames:
+                        return images
+
+                count += 1
+            # assert ret, "videocap.read fails!"
+            # img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            # im_pil = Image.fromarray(img)
+            # print(f"loading {stamp} success")
+            return images
+
+        # return [get_frame(0), get_frame(stamp1), get_frame(stamp2)]
+        # img = get_frame(0)
+        # img1 = get_frame(frame_interval * 1)
+        # return [img, img1, img, img1, img, img1,]
+        return get_frame(8)
+
+    frames = [
+        None,
+    ]
+    if videobox is not None:
+        frames = extract_frames(videobox)
+        # add frames as regular images
+        logger.info(f"Got videobox: {videobox}.")
+
+    logger.info(f"add_image. ip: {request.client.host}.")
+    image_list = [imagebox, imagebox_2, imagebox_3, *frames]
+    logger.info(f"image_list: {image_list}")
+
+    im_count = 0
+    for image in image_list:
+        if image is not None:
+            im_count += 1
+    for image in image_list:
+        if image is not None:
+            if args.auto_pad_image_token or im_count == 1:
+                text = (AUTO_FILL_IM_TOKEN_HOLDER, image, image_process_mode)
+            else:
+                text = ("", image, image_process_mode)
+            state.append_message(None, text)
+            state.append_message(
+                None, None
+            )  # in order to match the input-output pair for textbox outputs
+            # state.append_message(state.roles[0], text)
+            # state.append_message(state.roles[1], None)
+    # state.skip_next = False
+    logger.info(f"im_count {im_count}. ip: {request.client.host}.")
+    state.image_loaded = True
+    # return (state,) + (None,) * IMAGE_BOX_NUM
+    return state
+
+
+def add_text_only(state, text, request: gr.Request):
+    logger.info(f"add_text_only. ip: {request.client.host}. len: {len(text)}")
+
+    if args.moderate:
+        flagged = violates_moderation(text)
+        if flagged:
+            state.skip_next = True
+            return (state, moderation_msg) + (no_change_btn,) * BUTTON_LIST_LEN
+
+    # This is 1536 characters, rather than tokens
+    text = text[:1536]  # Hard cut-off
+    state.append_message(state.roles[0], text)
+    state.append_message(state.roles[1], None)
+    state.skip_next = False
+    return (state, "") + (disable_btn,) * BUTTON_LIST_LEN
+
+
+def add_text(
+    state, text, image, image_process_mode, prompt_style_btn, request: gr.Request
+):
+    logger.info(f"add_text. ip: {request.client.host}. len: {len(text)}")
+    if len(text) <= 0 and image is None:
+        state.skip_next = True
+        return (state, state.to_gradio_chatbot(), "", None) + (
+            no_change_btn,
+        ) * BUTTON_LIST_LEN
+    if args.moderate:
+        flagged = violates_moderation(text)
+        if flagged:
+            state.skip_next = True
+            return (state, state.to_gradio_chatbot(), moderation_msg, None) + (
+                no_change_btn,
+            ) * BUTTON_LIST_LEN
+
+    text = text[:1536]  # Hard cut-off
+    if image is not None:
+        text = text[:1200]  # Hard cut-off for images
+        if "<image>" not in text:
+            # text = '<Image><image></Image>' + text
+            text = text + "\n<image>"
+        text = (text, image, image_process_mode)
+        if len(state.get_images(return_pil=True)) > 0:
+            state = get_conversation(prompt_style_btn)
+            # state = default_conversation.copy()
+    state.append_message(state.roles[0], text)
+    state.append_message(state.roles[1], None)
+    state.skip_next = False
+    return (state, state.to_gradio_chatbot(), "", None) + (
+        disable_btn,
+    ) * BUTTON_LIST_LEN
+
+
+def http_bot(
+    state,
+    model_selector,
+    temperature,
+    top_p,
+    max_new_tokens,
+    prompt_style_btn,
+    request: gr.Request,
+):
+    logger.info(f"http_bot. ip: {request.client.host}")
+    start_tstamp = time.time()
+    model_name = model_selector
+
+    if state.skip_next:
+        # This generate call is skipped due to invalid inputs
+        yield (state, state.to_gradio_chatbot()) + (no_change_btn,) * BUTTON_LIST_LEN
+        return
+    if len(state.messages) == state.offset + 2:
+        # First round of conversation
+        if "llava" in model_name.lower():
+            if "llama-2" in model_name.lower():
+                template_name = "llava_llama_2"
+            elif "v1" in model_name.lower():
+                if "mmtag" in model_name.lower():
+                    template_name = "v1_mmtag"
+                elif (
+                    "plain" in model_name.lower()
+                    and "finetune" not in model_name.lower()
+                ):
+                    template_name = "v1_mmtag"
+                else:
+                    template_name = "llava_v1"
+            elif "mpt" in model_name.lower():
+                template_name = "mpt"
+            else:
+                if "mmtag" in model_name.lower():
+                    template_name = "v0_mmtag"
+                elif (
+                    "plain" in model_name.lower()
+                    and "finetune" not in model_name.lower()
+                ):
+                    template_name = "v0_mmtag"
+                else:
+                    template_name = "llava_v0"
+        elif "mpt" in model_name:
+            template_name = "mpt_text"
+        elif "llama-2" in model_name:
+            template_name = "llama_2"
+        else:
+            template_name = "vicuna_v1"
+        if prompt_style_btn == "no-sys":
+            new_state = get_conversation(prompt_style_btn)
+        else:
+            new_state = conv_templates[template_name].copy()
+        new_state.append_message(new_state.roles[0], state.messages[-2][1])
+        new_state.append_message(new_state.roles[1], None)
+        state = new_state
+
+    # Query worker address
+    controller_url = args.controller_url
+    ret = requests.post(
+        controller_url + "/get_worker_address", json={"model": model_name}
+    )
+    worker_addr = ret.json()["address"]
+    logger.info(f"model_name: {model_name}, worker_addr: {worker_addr}")
+
+    # No available worker
+    if worker_addr == "":
+        state.messages[-1][-1] = server_error_msg
+        yield (
+            state,
+            state.to_gradio_chatbot(),
+            # disable_btn,
+            # disable_btn,
+            # disable_btn,
+            enable_btn,
+            enable_btn,
+        )
+        return
+
+    # Construct prompt
+    prompt = state.get_prompt()
+
+    all_images = state.get_images(return_pil=True)
+    all_image_hash = [hashlib.md5(image.tobytes()).hexdigest() for image in all_images]
+    for image, hash in zip(all_images, all_image_hash):
+        t = datetime.datetime.now()
+        filename = os.path.join(
+            LOGDIR, "serve_images", f"{t.year}-{t.month:02d}-{t.day:02d}", f"{hash}.jpg"
+        )
+        if not os.path.isfile(filename):
+            os.makedirs(os.path.dirname(filename), exist_ok=True)
+            image.save(filename)
+
+    # Make requests
+    pload = {
+        "model": model_name,
+        "prompt": prompt,
+        "temperature": float(temperature),
+        "top_p": float(top_p),
+        "max_new_tokens": min(int(max_new_tokens), 1536),
+        "stop": (
+            state.sep
+            if state.sep_style in [SeparatorStyle.SINGLE, SeparatorStyle.MPT]
+            else state.sep2
+        ),
+        "images": f"List of {len(state.get_images())} images: {all_image_hash}",
+    }
+
+    image_num = len(state.get_images())
+    if image_num == 0:
+        state.messages[-1][
+            -1
+        ] = "**NO INPUT IMAGE RECEIVED BY THE SERVER. PLEASE CHECK YOUR INTERNET CONNECTION AND REFRESH THE PAGE.**"
+        yield (
+            state,
+            state.to_gradio_chatbot(),
+            # disable_btn,
+            # disable_btn,
+            # disable_btn,
+            enable_btn,
+            enable_btn,
+        )
+        return
+
+    count_auto_im_token = prompt.count(AUTO_FILL_IM_TOKEN_HOLDER)
+    count_manual_im_token = prompt.count(LLAVA_DEFAULT_IM_TOKEN_PLACE_HOLDER)
+    if (count_auto_im_token == image_num) and (
+        count_manual_im_token == 0
+    ):  # Use default system prompt
+        prompt = prompt.replace(AUTO_FILL_IM_TOKEN_HOLDER, LLAVA_DEFAULT_IMAGE_TOKEN)
+    elif (count_auto_im_token == image_num) and (
+        count_manual_im_token == image_num
+    ):  # Use <image> token inserted by user
+        prompt = prompt.replace(AUTO_FILL_IM_TOKEN_HOLDER, "")
+        prompt = prompt.replace(
+            LLAVA_DEFAULT_IM_TOKEN_PLACE_HOLDER, LLAVA_DEFAULT_IMAGE_TOKEN
+        )
+    elif (count_auto_im_token == 0) and (count_manual_im_token == image_num):
+        prompt = prompt.replace(
+            LLAVA_DEFAULT_IM_TOKEN_PLACE_HOLDER, LLAVA_DEFAULT_IMAGE_TOKEN
+        )
+    else:
+        state.messages[-1][
+            -1
+        ] = "**IMAGE NUM MISMATCHES IMAGE TOKEN PLACEHOLDER. PLEASE CHECK YOUR INPUT AND REFRESH THE PAGE.**"
+        yield (
+            state,
+            state.to_gradio_chatbot(),
+            # disable_btn,
+            # disable_btn,
+            # disable_btn,
+            enable_btn,
+            enable_btn,
+        )
+        return
+
+    pload["prompt"] = prompt
+    logger.info(f"==== request ====\n{pload}")
+    pload["images"] = state.get_images()
+
+    state.messages[-1][-1] = "▌"
+    ret = state.to_gradio_chatbot()
+    ret[0][0] = ret[0][0].replace(LLAVA_DEFAULT_IM_TOKEN_PLACE_HOLDER, IMAGE_TOKEN_VIS)
+    yield (state, ret) + (disable_btn,) * BUTTON_LIST_LEN
+
+    try:
+        # Stream output
+        response = requests.post(
+            worker_addr + "/worker_generate_stream",
+            headers=headers,
+            json=pload,
+            stream=True,
+            timeout=10,
+        )
+        for chunk in response.iter_lines(decode_unicode=False, delimiter=b"\0"):
+            if chunk:
+                data = json.loads(chunk.decode())
+                if data["error_code"] == 0:
+                    output = data["text"][len(prompt) :].strip()
+                    state.messages[-1][-1] = output + "▌"
+                    ret = state.to_gradio_chatbot()
+                    ret[0][0] = ret[0][0].replace(
+                        LLAVA_DEFAULT_IM_TOKEN_PLACE_HOLDER, IMAGE_TOKEN_VIS
+                    )
+                    yield (state, ret) + (disable_btn,) * BUTTON_LIST_LEN
+                else:
+                    output = data["text"] + f" (error_code: {data['error_code']})"
+                    state.messages[-1][-1] = output
+                    ret = state.to_gradio_chatbot()
+                    ret[0][0] = ret[0][0].replace(
+                        LLAVA_DEFAULT_IM_TOKEN_PLACE_HOLDER, IMAGE_TOKEN_VIS
+                    )
+                    yield (state, ret) + (
+                        # disable_btn,
+                        # disable_btn,
+                        # disable_btn,
+                        enable_btn,
+                        enable_btn,
+                    )
+                    return
+                time.sleep(0.03)
+    except requests.exceptions.RequestException as e:
+        state.messages[-1][-1] = server_error_msg
+        yield (state, state.to_gradio_chatbot()) + (
+            # disable_btn,
+            # disable_btn,
+            # disable_btn,
+            enable_btn,
+            enable_btn,
+        )
+        return
+
+    state.messages[-1][-1] = state.messages[-1][-1][:-1]
+    ret = state.to_gradio_chatbot()
+    ret[0][0] = ret[0][0].replace(LLAVA_DEFAULT_IM_TOKEN_PLACE_HOLDER, IMAGE_TOKEN_VIS)
+    yield (state, ret) + (enable_btn,) * BUTTON_LIST_LEN
+
+    finish_tstamp = time.time()
+    logger.info(f"{output}")
+
+    with open(get_conv_log_filename(), "a") as fout:
+        data = {
+            "tstamp": round(finish_tstamp, 4),
+            "type": "chat",
+            "model": model_name,
+            "start": round(start_tstamp, 4),
+            "finish": round(finish_tstamp, 4),
+            "state": state.dict(),
+            "images": all_image_hash,
+            "ip": request.client.host,
+        }
+        fout.write(json.dumps(data) + "\n")
+
+
+title_markdown = """
+# VILA: On Pre-training for Visual Language Models
+[\[Paper\]](https://arxiv.org/abs/2312.07533)  [\[Github\]](https://github.com/NVlabs/VILA)
+### Powered by [TinyChat](https://github.com/mit-han-lab/llm-awq/tree/main/tinychat) with 4-bit [AWQ](https://arxiv.org/abs/2306.00978).
+"""
+
+tos_markdown = """
+### Terms of Use
+By using this service, users are required to agree to the following terms:
+The service is a research preview intended for non-commercial use only. It only provides limited safety measures and may generate offensive content. It must not be used for any illegal, harmful, violent, racist, or sexual purposes. The service may collect user dialogue data for future research.
+Please click the "Flag" button if you get any inappropriate answer! We will collect those to keep improving our moderator.
+For an optimal experience, please use desktop computers for this demo, as mobile devices may compromise its quality.
+"""
+
+
+learn_more_markdown = """
+### License
+The service is a research preview intended for non-commercial use only, subject to the model [License](https://github.com/facebookresearch/llama/blob/main/MODEL_CARD.md) of LLaMA, [Terms of Use](https://openai.com/policies/terms-of-use) of the data generated by OpenAI, and [Privacy Practices](https://chrome.google.com/webstore/detail/sharegpt-share-your-chatg/daiacboceoaocpibfodeljbdfacokfjb) of ShareGPT. Please contact us if you find any potential violation.
+"""
+
+ack_markdown = """
+### Acknowledgement
+This demo is inspired by [LLaVA](https://github.com/haotian-liu/LLaVA). We thank LLaVA for providing an elegant way to build the Gradio Web UI.          
+"""
+
+block_css = """
+
+#buttons button {
+    min-width: min(120px,100%);
+}
+
+"""
+
+
+def build_demo(embed_mode):
+    textbox = gr.Textbox(
+        show_label=False, placeholder="Enter text and press ENTER", container=False
+    )
+    with gr.Blocks(
+        title="VILA on TinyChat", theme=gr.themes.Default(), css=block_css
+    ) as demo:
+        state = gr.State()
+
+        if not embed_mode:
+            gr.Markdown(title_markdown)
+
+        with gr.Row():
+            with gr.Column(scale=8):
+                with gr.Row():
+                    imagebox = gr.Image(type="pil")
+                    imagebox_2 = gr.Image(type="pil")
+                    imagebox_3 = gr.Image(type="pil")
+                    videobox = gr.Video(label="1 video = 8 frames")
+                image_process_mode = gr.Radio(
+                    ["Crop", "Resize", "Pad", "Default"],
+                    value="Default",
+                    label="Preprocess for non-square image",
+                    visible=False,
+                )
+                # imagebox_out = gr.Image(height=150)
+                with gr.Row():
+                    with gr.Column(scale=5):
+                        textbox.render()
+                    with gr.Column(scale=1, min_width=100):
+                        submit_btn = gr.Button(value="Send", variant="primary")
+                    with gr.Column(scale=1, min_width=100):
+                        clear_btn = gr.Button(
+                            value="🗑️  Clear", variant="primary", interactive=False
+                        )
+                    with gr.Column(scale=1, min_width=100):
+                        regenerate_btn = gr.Button(
+                            value="🔄  Retry", variant="primary", interactive=False
+                        )
+                with gr.Row():
+                    gr.Markdown(
+                        "### *** Before changing the current images, uploading new images or switching the prompt style, please click the clear button."
+                    )
+                chatbot = gr.Chatbot(
+                    elem_id="chatbot", label="TinyChat Assistant", height=550
+                )
+
+            with gr.Column(scale=4):
+                with gr.Row(equal_height=True):
+                    with gr.Column(scale=1, min_width=50):
+                        model_selector = gr.Dropdown(
+                            choices=models,
+                            value=models[0] if len(models) > 0 else "",
+                            label="Model",
+                            interactive=True,
+                            show_label=True,
+                            container=False,
+                        )
+                    with gr.Column(scale=1, min_width=50):
+                        prompt_style_btn = gr.Radio(
+                            ["default", "no-sys"],
+                            label="Prompt style",
+                            value="default",
+                            interactive=True,
+                        )
+
+                # with gr.Row():
+                # with gr.Column(scale=1, min_width=50):
+                #     im_submit_btn = gr.Button(value="Add image", variant="primary")
+                # with gr.Column(scale=1, min_width=50):
+                #     submit_btn_1 = gr.Button(value="Send", variant="primary")
+
+                cur_dir = os.path.dirname(os.path.abspath(__file__))
+                with gr.Row(equal_height=True):
+                    gr.Examples(
+                        examples=[
+                            [
+                                f"{cur_dir}/examples/video/qZDF__7LNKc.4.mp4",
+                                "Elaborate on the visual and narrative elements of the video in detail.",
+                            ],
+                        ],
+                        label="Video Example",
+                        inputs=[videobox, textbox],
+                        fn=clear_after_click_example_1_video,
+                        outputs=[
+                            state,
+                            imagebox,
+                            imagebox_2,
+                            imagebox_3,
+                            videobox,
+                            prompt_style_btn,
+                        ],
+                        run_on_click=True,
+                    )
+                with gr.Row(equal_height=True):
+                    with gr.Column(scale=1, min_width=50):
+                        gr.Examples(
+                            examples=[
+                                [
+                                    f"{cur_dir}/examples/pedestrain.png",
+                                    "<image> What is the person in the center of the image doing?",
+                                ],
+                            ],
+                            label="Image Example 1",
+                            inputs=[imagebox, textbox],
+                            fn=clear_after_click_example_1_image,
+                            outputs=[
+                                state,
+                                imagebox,
+                                imagebox_2,
+                                imagebox_3,
+                                videobox,
+                                prompt_style_btn,
+                            ],
+                            run_on_click=True,
+                        )
+                    with gr.Column(scale=1, min_width=50):
+                        gr.Examples(
+                            examples=[
+                                [
+                                    f"{cur_dir}/examples/car_repair.png",
+                                    "<image> What is the brand of the silver car in the image?",
+                                ],
+                            ],
+                            label="Image Example 2",
+                            inputs=[imagebox, textbox],
+                            fn=clear_after_click_example_1_image,
+                            outputs=[
+                                state,
+                                imagebox,
+                                imagebox_2,
+                                imagebox_3,
+                                videobox,
+                                prompt_style_btn,
+                            ],
+                            run_on_click=True,
+                        )
+                with gr.Row(equal_height=True):
+                    with gr.Column(scale=1, min_width=50):
+                        gr.Examples(
+                            examples=[
+                                [
+                                    f"{cur_dir}/examples/CPR.jpg",
+                                    "<image> What are the people doing in this image?",
+                                ],
+                            ],
+                            label="Image Example 3",
+                            inputs=[imagebox, textbox],
+                            fn=clear_after_click_example_1_image,
+                            outputs=[
+                                state,
+                                imagebox,
+                                imagebox_2,
+                                imagebox_3,
+                                videobox,
+                                prompt_style_btn,
+                            ],
+                            run_on_click=True,
+                        )
+                    with gr.Column(scale=1, min_width=50):
+                        gr.Examples(
+                            examples=[
+                                [
+                                    f"{cur_dir}/examples/Wall_fissure.png",
+                                    "<image> What are the likely service needed for this building?",
+                                ],
+                            ],
+                            label="Image Example 4",
+                            inputs=[imagebox, textbox],
+                            fn=clear_after_click_example_1_image,
+                            outputs=[
+                                state,
+                                imagebox,
+                                imagebox_2,
+                                imagebox_3,
+                                videobox,
+                                prompt_style_btn,
+                            ],
+                            run_on_click=True,
+                        )
+
+                with gr.Row(equal_height=True):
+                    with gr.Column(scale=1, min_width=50):
+                        gr.Examples(
+                            examples=[
+                                [
+                                    f"{cur_dir}/examples/animal_blocking.png",
+                                    "<image> What is unusual in this image?",
+                                ],
+                            ],
+                            label="Image Example 5",
+                            inputs=[imagebox, textbox],
+                            fn=clear_after_click_example_1_image,
+                            outputs=[
+                                state,
+                                imagebox,
+                                imagebox_2,
+                                imagebox_3,
+                                videobox,
+                                prompt_style_btn,
+                            ],
+                            run_on_click=True,
+                        )
+                    with gr.Column(scale=1, min_width=50):
+                        gr.Examples(
+                            examples=[
+                                [
+                                    f"{cur_dir}/examples/windmill_people.png",
+                                    "<image> Can you describe what is happening?",
+                                ],
+                            ],
+                            label="Image Example 6",
+                            inputs=[imagebox, textbox],
+                            fn=clear_after_click_example_1_image,
+                            outputs=[
+                                state,
+                                imagebox,
+                                imagebox_2,
+                                imagebox_3,
+                                videobox,
+                                prompt_style_btn,
+                            ],
+                            run_on_click=True,
+                        )
+
+                gr.Examples(
+                    examples=[
+                        [
+                            f"{cur_dir}/examples/climate_change/climate_change_1.png",
+                            f"{cur_dir}/examples/climate_change/climate_change_2.png",
+                            "<image> <image> What is the implication of temperature based on this image?",
+                        ],
+                    ],
+                    inputs=[imagebox, imagebox_2, textbox],
+                    label="Multi-image Example 1",
+                    fn=clear_after_click_example_2_image,
+                    outputs=[
+                        state,
+                        imagebox,
+                        imagebox_2,
+                        imagebox_3,
+                        videobox,
+                        prompt_style_btn,
+                    ],
+                    run_on_click=True,
+                )
+
+                gr.Examples(
+                    examples=[
+                        [
+                            f"{cur_dir}/examples/palms/palm1.png",
+                            f"{cur_dir}/examples/palms/palm2.png",
+                            f"{cur_dir}/examples/palms/palm3.png",
+                            "8:15am: <image> 12:45pm: <image> 16:00pm: <image> When did I have lunch and what did I eat for lunch?",
+                        ],
+                    ],
+                    inputs=[imagebox, imagebox_2, imagebox_3, textbox],
+                    label="Multi-image Example 2",
+                    fn=clear_after_click_example_3_image,
+                    outputs=[
+                        state,
+                        imagebox,
+                        imagebox_2,
+                        imagebox_3,
+                        videobox,
+                        prompt_style_btn,
+                    ],
+                    run_on_click=True,
+                )
+
+                gr.Examples(
+                    examples=[
+                        [
+                            f"{cur_dir}/examples/golf/Golfman1.png",
+                            f"{cur_dir}/examples/golf/Golfman2.png",
+                            f"{cur_dir}/examples/golf/Golfman3.png",
+                            "<image> <image> <image> What happens to the man after hitting the ball? And why does it happen?",
+                        ],
+                    ],
+                    inputs=[imagebox, imagebox_2, imagebox_3, textbox],
+                    label="Multi-image Example 3",
+                    fn=clear_after_click_example_3_image,
+                    outputs=[
+                        state,
+                        imagebox,
+                        imagebox_2,
+                        imagebox_3,
+                        videobox,
+                        prompt_style_btn,
+                    ],
+                    run_on_click=True,
+                )
+
+                gr.Examples(
+                    examples=[
+                        [
+                            f"{cur_dir}/examples/icl-logo/google.webp",
+                            f"{cur_dir}/examples/icl-logo/apple.jpg",
+                            f"{cur_dir}/examples/icl-logo/nvidia.png",
+                            "<image> is famous for its search engine. <image> is famous for Mac and iPhone. <image> ",
+                        ],
+                    ],
+                    inputs=[imagebox, imagebox_2, imagebox_3, textbox],
+                    label="In-context Learning Example 1 (Please switch the prompt style to 'no-sys')",
+                    fn=clear_after_click_example_3_image_icl,
+                    outputs=[
+                        state,
+                        imagebox,
+                        imagebox_2,
+                        imagebox_3,
+                        videobox,
+                        prompt_style_btn,
+                    ],
+                    run_on_click=True,
+                )
+
+                gr.Examples(
+                    examples=[
+                        [
+                            f"{cur_dir}/examples/icl-building/csail_building.jpeg",
+                            f"{cur_dir}/examples/icl-building/Toronto_Tower.jpeg",
+                            f"{cur_dir}/examples/icl-building/Golden_State_Bridge.jpeg",
+                            "<image> Boston. <image> Toronto. <image> ",
+                        ],
+                    ],
+                    inputs=[imagebox, imagebox_2, imagebox_3, textbox],
+                    label="In-context Learning Example 2 (Please switch the prompt style to 'no-sys')",
+                    fn=clear_after_click_example_3_image_icl,
+                    outputs=[
+                        state,
+                        imagebox,
+                        imagebox_2,
+                        imagebox_3,
+                        videobox,
+                        prompt_style_btn,
+                    ],
+                    run_on_click=True,
+                )
+
+                gr.Examples(
+                    examples=[
+                        [
+                            f"{cur_dir}/examples/arts/sunflowers.jpg",
+                            f"{cur_dir}/examples/arts/the_persistence_of_memory.png",
+                            f"{cur_dir}/examples/arts/impression_sunrise.png",
+                            "<image> Vincent Van Gogh. <image> Salvador Dalí. <image>",
+                        ],
+                    ],
+                    inputs=[imagebox, imagebox_2, imagebox_3, textbox],
+                    label="In-context Learning Example 3 (Please switch the prompt style to 'no-sys')",
+                    fn=clear_after_click_example_3_image_icl,
+                    outputs=[
+                        state,
+                        imagebox,
+                        imagebox_2,
+                        imagebox_3,
+                        videobox,
+                        prompt_style_btn,
+                    ],
+                    run_on_click=True,
+                )
+
+                with gr.Accordion("Parameters", open=False) as parameter_row:
+                    temperature = gr.Slider(
+                        minimum=0.0,
+                        maximum=1.0,
+                        value=0.2,
+                        step=0.1,
+                        interactive=True,
+                        label="Temperature",
+                    )
+                    top_p = gr.Slider(
+                        minimum=0.0,
+                        maximum=1.0,
+                        value=1.0,
+                        step=0.1,
+                        interactive=True,
+                        label="Top P",
+                    )
+                    max_output_tokens = gr.Slider(
+                        minimum=0,
+                        maximum=1024,
+                        value=512,
+                        step=64,
+                        interactive=True,
+                        label="Max output tokens",
+                    )
+
+                # with gr.Row(elem_id="buttons") as button_row:
+                # upvote_btn = gr.Button(value="👍  Upvote", interactive=False)
+                # downvote_btn = gr.Button(value="👎  Downvote", interactive=False)
+                # flag_btn = gr.Button(value="⚠️  Flag", interactive=False)
+                # stop_btn = gr.Button(value="⏹️  Stop Generation", interactive=False)
+                # regenerate_btn = gr.Button(value="🔄  Regenerate", interactive=False)
+
+        if not embed_mode:
+            gr.Markdown(tos_markdown)
+            gr.Markdown(learn_more_markdown)
+            gr.Markdown(ack_markdown)
+        url_params = gr.JSON(visible=False)
+
+        # Register listeners
+        # btn_list = [upvote_btn, downvote_btn, flag_btn, regenerate_btn, clear_btn]
+        btn_list = [regenerate_btn, clear_btn]
+
+        # upvote_btn.click(
+        #     upvote_last_response,
+        #     [state, model_selector],
+        #     [textbox, upvote_btn, downvote_btn, flag_btn],
+        #     queue=False,
+        # )
+        # downvote_btn.click(
+        #     downvote_last_response,
+        #     [state, model_selector],
+        #     [textbox, upvote_btn, downvote_btn, flag_btn],
+        #     queue=False,
+        # )
+        # flag_btn.click(
+        #     flag_last_response,
+        #     [state, model_selector],
+        #     [textbox, upvote_btn, downvote_btn, flag_btn],
+        #     queue=False,
+        # )
+
+        regenerate_btn.click(
+            regenerate,
+            [state, image_process_mode],
+            [state, chatbot, textbox] + btn_list,
+            queue=False,
+        ).then(
+            http_bot,
+            [
+                state,
+                model_selector,
+                temperature,
+                top_p,
+                max_output_tokens,
+                prompt_style_btn,
+            ],
+            [state, chatbot] + btn_list,
+        )
+
+        prompt_style_btn.change(
+            change_prompt_style, [state, prompt_style_btn], [state], queue=False
+        )
+
+        clear_btn.click(
+            clear_history,
+            [prompt_style_btn],
+            [state, chatbot, textbox, imagebox, imagebox_2, imagebox_3, videobox]
+            + btn_list,
+            queue=False,
+        )
+
+        # textbox.submit(
+        #     add_text,
+        #     [state, textbox, imagebox, image_process_mode],
+        #     [state, chatbot, textbox, imagebox] + btn_list,
+        #     queue=False
+        # ).then(
+        #     http_bot,
+        #     [state, model_selector, temperature, top_p, max_output_tokens],
+        #     [state, chatbot] + btn_list
+        # )
+
+        # im_submit_btn.click(
+        #     mirror,
+        #     inputs=[imagebox],
+        #     outputs=[imagebox_out]
+        # ).then(
+        #     add_image,
+        #     [state, imagebox, image_process_mode],
+        #     [state, imagebox] + btn_list,
+        #     queue=False
+        # )
+
+        textbox.submit(
+            clear_text_history, [state, prompt_style_btn], [state, chatbot], queue=False
+        ).then(
+            add_images,
+            [state, imagebox, imagebox_2, imagebox_3, videobox, image_process_mode],
+            [state],
+            queue=False,
+        ).then(
+            add_text_only, [state, textbox], [state, textbox] + btn_list, queue=False
+        ).then(
+            http_bot,
+            [
+                state,
+                model_selector,
+                temperature,
+                top_p,
+                max_output_tokens,
+                prompt_style_btn,
+            ],
+            [state, chatbot] + btn_list,
+        )
+
+        submit_btn.click(
+            clear_text_history, [state, prompt_style_btn], [state, chatbot], queue=False
+        ).then(
+            add_images,
+            [state, imagebox, imagebox_2, imagebox_3, videobox, image_process_mode],
+            [state],
+            queue=False,
+        ).then(
+            add_text_only, [state, textbox], [state, textbox] + btn_list, queue=False
+        ).then(
+            http_bot,
+            [
+                state,
+                model_selector,
+                temperature,
+                top_p,
+                max_output_tokens,
+                prompt_style_btn,
+            ],
+            [state, chatbot] + btn_list,
+        )
+
+        if args.model_list_mode == "once":
+            demo.load(
+                load_demo,
+                [url_params, prompt_style_btn],
+                [state, model_selector],
+                _js=get_window_url_params,
+                queue=False,
+            )
+        elif args.model_list_mode == "reload":
+            demo.load(
+                load_demo_refresh_model_list,
+                [prompt_style_btn],
+                [state, model_selector],
+                queue=False,
+            )
+        else:
+            raise ValueError(f"Unknown model list mode: {args.model_list_mode}")
+
+    return demo
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--host", type=str, default="0.0.0.0")
+    parser.add_argument("--port", type=int)
+    parser.add_argument("--controller-url", type=str, default="http://localhost:21001")
+    parser.add_argument("--concurrency-count", type=int, default=10)
+    parser.add_argument(
+        "--model-list-mode", type=str, default="once", choices=["once", "reload"]
+    )
+    parser.add_argument("--share", action="store_true")
+    parser.add_argument("--moderate", action="store_true")
+    parser.add_argument("--embed", action="store_true")
+    parser.add_argument(
+        "--auto-pad-image-token",
+        action="store_true",
+        help="Automatically pad <image> token to the before of the prompt if no user inputs.",
+    )
+    # NOTE: For single image input, we still auto pad <image> token even if the --auto-pad-image-token is False
+    args = parser.parse_args()
+    logger.info(f"args: {args}")
+
+    models = get_model_list()
+
+    logger.info(args)
+    demo = build_demo(args.embed)
+    demo.queue(concurrency_count=args.concurrency_count, api_open=False).launch(
+        server_name=args.host, server_port=args.port, share=args.share
+    )

llm-awq/tinychat/serve/llava_conv.py

@@ -0,0 +1,454 @@
+#    Modified from https://github.com/haotian-liu/LLaVA
+#    Copyright 2023 Haotian Liu
+#
+#    Licensed under the Apache License, Version 2.0 (the "License");
+#    you may not use this file except in compliance with the License.
+#    You may obtain a copy of the License at
+#
+#        http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS,
+#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#    See the License for the specific language governing permissions and
+#    limitations under the License.
+
+import dataclasses
+from enum import auto, Enum
+from typing import List, Tuple
+
+from tinychat.utils.constants import AUTO_FILL_IM_TOKEN_HOLDER
+
+
+class SeparatorStyle(Enum):
+    """Different separator style."""
+
+    SINGLE = auto()
+    TWO = auto()
+    MPT = auto()
+    PLAIN = auto()
+    LLAMA_2 = auto()
+
+
+@dataclasses.dataclass
+class Conversation:
+    """A class that keeps all conversation history."""
+
+    system: str
+    roles: List[str]
+    messages: List[List[str]]
+    offset: int
+    sep_style: SeparatorStyle = SeparatorStyle.SINGLE
+    sep: str = "###"
+    sep2: str = None
+    version: str = "Unknown"
+
+    skip_next: bool = False
+    image_loaded: bool = False
+
+    def get_prompt(self):
+        messages = self.messages
+        if len(messages) > 0 and type(messages[0][1]) is tuple:
+            messages = self.messages.copy()
+            init_role, init_msg = messages[0].copy()
+            # init_msg = init_msg[0].replace(AUTO_FILL_IM_TOKEN_HOLDER, "").strip()
+            init_msg = init_msg[0].replace("", "").strip()
+            if "mmtag" in self.version:
+                messages[0] = (init_role, init_msg)
+                messages.insert(0, (self.roles[0], "<Image><image></Image>"))
+                messages.insert(1, (self.roles[1], "Received."))
+            else:
+                # messages[0] = (init_role, AUTO_FILL_IM_TOKEN_HOLDER + init_msg)
+                messages[0] = (init_role, "" + init_msg)
+
+        if self.sep_style == SeparatorStyle.SINGLE:
+            ret = self.system + self.sep
+            for role, message in messages:
+                if role:
+                    if message:
+                        if type(message) is tuple:
+                            message, _, _ = message
+                        ret += role + ": " + message + self.sep
+                    else:
+                        ret += role + ":"
+                else:
+                    if message:
+                        if type(message) is tuple:
+                            message, _, _ = message
+                        ret += message
+        elif self.sep_style == SeparatorStyle.TWO:
+            seps = [self.sep, self.sep2]
+            ret = self.system + seps[0]
+            for i, (role, message) in enumerate(messages):
+                if role:
+                    if message:
+                        if type(message) is tuple:
+                            message, _, _ = message
+                        ret += role + ": " + message + seps[i % 2]
+                    else:
+                        ret += role + ":"
+                else:
+                    if message:
+                        if type(message) is tuple:
+                            message, _, _ = message
+                        ret += message + seps[i % 2]
+        elif self.sep_style == SeparatorStyle.MPT:
+            ret = self.system + self.sep
+            for role, message in messages:
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += role + message + self.sep
+                else:
+                    ret += role
+        elif self.sep_style == SeparatorStyle.LLAMA_2:
+            wrap_sys = lambda msg: f"<<SYS>>\n{msg}\n<</SYS>>\n\n"
+            wrap_inst = lambda msg: f"[INST] {msg} [/INST]"
+            ret = ""
+
+            for i, (role, message) in enumerate(messages):
+                if i == 0:
+                    assert message, "first message should not be none"
+                    assert role == self.roles[0], "first message should come from user"
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    if i == 0:
+                        message = wrap_sys(self.system) + message
+                    if i % 2 == 0:
+                        message = wrap_inst(message)
+                        ret += self.sep + message
+                    else:
+                        ret += " " + message + " " + self.sep2
+                else:
+                    ret += ""
+            ret = ret.lstrip(self.sep)
+        elif self.sep_style == SeparatorStyle.PLAIN:
+            seps = [self.sep, self.sep2]
+            ret = self.system
+            for i, (role, message) in enumerate(messages):
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += message + seps[i % 2]
+                else:
+                    ret += ""
+        else:
+            raise ValueError(f"Invalid style: {self.sep_style}")
+
+        return ret
+
+    def append_message(self, role, message):
+        self.messages.append([role, message])
+
+    def get_images(self, return_pil=False):
+        images = []
+        for i, (role, msg) in enumerate(self.messages[self.offset :]):
+            if i % 2 == 0:
+                if type(msg) is tuple:
+                    import base64
+                    from io import BytesIO
+                    from PIL import Image
+
+                    msg, image, image_process_mode = msg
+                    if image_process_mode == "Pad":
+
+                        def expand2square(pil_img, background_color=(122, 116, 104)):
+                            width, height = pil_img.size
+                            if width == height:
+                                return pil_img
+                            elif width > height:
+                                result = Image.new(
+                                    pil_img.mode, (width, width), background_color
+                                )
+                                result.paste(pil_img, (0, (width - height) // 2))
+                                return result
+                            else:
+                                result = Image.new(
+                                    pil_img.mode, (height, height), background_color
+                                )
+                                result.paste(pil_img, ((height - width) // 2, 0))
+                                return result
+
+                        image = expand2square(image)
+                    elif image_process_mode in ["Default", "Crop"]:
+                        pass
+                    elif image_process_mode == "Resize":
+                        image = image.resize((336, 336))
+                    else:
+                        raise ValueError(
+                            f"Invalid image_process_mode: {image_process_mode}"
+                        )
+                    max_hw, min_hw = max(image.size), min(image.size)
+                    aspect_ratio = max_hw / min_hw
+                    max_len, min_len = 800, 400
+                    shortest_edge = int(min(max_len / aspect_ratio, min_len, min_hw))
+                    longest_edge = int(shortest_edge * aspect_ratio)
+                    W, H = image.size
+                    if longest_edge != max(image.size):
+                        if H > W:
+                            H, W = longest_edge, shortest_edge
+                        else:
+                            H, W = shortest_edge, longest_edge
+                        image = image.resize((W, H))
+                    if return_pil:
+                        images.append(image)
+                    else:
+                        buffered = BytesIO()
+                        image.save(buffered, format="PNG")
+                        img_b64_str = base64.b64encode(buffered.getvalue()).decode()
+                        images.append(img_b64_str)
+        return images
+
+    def to_gradio_chatbot(self):
+        ret = []
+        # count the figures to skip visualizing them in the text box.
+        cur_num_fig = 0
+        for i, (role, msg) in enumerate(self.messages[self.offset :]):
+            if i % 2 == 0:
+                if type(msg) is tuple:
+                    # Skip the visualization of image in the chatbox
+                    cur_num_fig += 1
+                    continue
+                    # import base64
+                    # from io import BytesIO
+                    # msg, image, image_process_mode = msg
+                    # max_hw, min_hw = max(image.size), min(image.size)
+                    # aspect_ratio = max_hw / min_hw
+                    # max_len, min_len = 800, 400
+                    # shortest_edge = int(min(max_len / aspect_ratio, min_len, min_hw))
+                    # longest_edge = int(shortest_edge * aspect_ratio)
+                    # W, H = image.size
+                    # if H > W:
+                    #     H, W = longest_edge, shortest_edge
+                    # else:
+                    #     H, W = shortest_edge, longest_edge
+                    # image = image.resize((W, H))
+                    # buffered = BytesIO()
+                    # image.save(buffered, format="JPEG")
+                    # img_b64_str = base64.b64encode(buffered.getvalue()).decode()
+                    # img_str = f'<img src="data:image/png;base64,{img_b64_str}" alt="user upload image" />'
+                    # msg = img_str + msg.replace('<image>', '').strip()
+                    # ret.append([msg[0], None])
+                else:
+                    ret.append([msg, None])
+            else:
+                if cur_num_fig > 0:
+                    cur_num_fig -= 1
+                    continue
+                ret[-1][-1] = msg
+        return ret
+
+    def copy(self):
+        return Conversation(
+            system=self.system,
+            roles=self.roles,
+            messages=[[x, y] for x, y in self.messages],
+            offset=self.offset,
+            sep_style=self.sep_style,
+            sep=self.sep,
+            sep2=self.sep2,
+            version=self.version,
+            image_loaded=self.image_loaded,
+        )
+
+    def dict(self):
+        if len(self.get_images()) > 0:
+            return {
+                "system": self.system,
+                "roles": self.roles,
+                "messages": [
+                    [x, y[0] if type(y) is tuple else y] for x, y in self.messages
+                ],
+                "offset": self.offset,
+                "sep": self.sep,
+                "sep2": self.sep2,
+            }
+        return {
+            "system": self.system,
+            "roles": self.roles,
+            "messages": self.messages,
+            "offset": self.offset,
+            "sep": self.sep,
+            "sep2": self.sep2,
+        }
+
+
+conv_vicuna_v0 = Conversation(
+    system="A chat between a curious human and an artificial intelligence assistant. "
+    "The assistant gives helpful, detailed, and polite answers to the human's questions.",
+    roles=("Human", "Assistant"),
+    messages=(
+        (
+            "Human",
+            "What are the key differences between renewable and non-renewable energy sources?",
+        ),
+        (
+            "Assistant",
+            "Renewable energy sources are those that can be replenished naturally in a relatively "
+            "short amount of time, such as solar, wind, hydro, geothermal, and biomass. "
+            "Non-renewable energy sources, on the other hand, are finite and will eventually be "
+            "depleted, such as coal, oil, and natural gas. Here are some key differences between "
+            "renewable and non-renewable energy sources:\n"
+            "1. Availability: Renewable energy sources are virtually inexhaustible, while non-renewable "
+            "energy sources are finite and will eventually run out.\n"
+            "2. Environmental impact: Renewable energy sources have a much lower environmental impact "
+            "than non-renewable sources, which can lead to air and water pollution, greenhouse gas emissions, "
+            "and other negative effects.\n"
+            "3. Cost: Renewable energy sources can be more expensive to initially set up, but they typically "
+            "have lower operational costs than non-renewable sources.\n"
+            "4. Reliability: Renewable energy sources are often more reliable and can be used in more remote "
+            "locations than non-renewable sources.\n"
+            "5. Flexibility: Renewable energy sources are often more flexible and can be adapted to different "
+            "situations and needs, while non-renewable sources are more rigid and inflexible.\n"
+            "6. Sustainability: Renewable energy sources are more sustainable over the long term, while "
+            "non-renewable sources are not, and their depletion can lead to economic and social instability.\n",
+        ),
+    ),
+    offset=2,
+    sep_style=SeparatorStyle.SINGLE,
+    sep="###",
+)
+
+empty_conv = Conversation(
+    system="",
+    roles=("", ""),
+    version="no-sys",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.TWO,
+    sep="",
+    sep2="</s>",
+)
+
+conv_vicuna_v1 = Conversation(
+    system="A chat between a curious user and an artificial intelligence assistant. "
+    "The assistant gives helpful, detailed, and polite answers to the user's questions.",
+    roles=("USER", "ASSISTANT"),
+    version="default",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.TWO,
+    sep=" ",
+    sep2="</s>",
+)
+
+conv_llama_2 = Conversation(
+    system="""You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe.  Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.
+
+If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.""",
+    roles=("USER", "ASSISTANT"),
+    version="llama_v2",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.LLAMA_2,
+    sep="<s>",
+    sep2="</s>",
+)
+
+conv_llava_llama_2 = Conversation(
+    system="You are a helpful language and vision assistant. "
+    "You are able to understand the visual content that the user provides, "
+    "and assist the user with a variety of tasks using natural language.",
+    roles=("USER", "ASSISTANT"),
+    version="llama_v2",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.LLAMA_2,
+    sep="<s>",
+    sep2="</s>",
+)
+
+conv_mpt = Conversation(
+    system="""<|im_start|>system
+A conversation between a user and an LLM-based AI assistant. The assistant gives helpful and honest answers.""",
+    roles=("<|im_start|>user\n", "<|im_start|>assistant\n"),
+    version="mpt",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.MPT,
+    sep="<|im_end|>",
+)
+
+conv_llava_plain = Conversation(
+    system="",
+    roles=("", ""),
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.PLAIN,
+    sep="\n",
+)
+
+conv_llava_v0 = Conversation(
+    system="A chat between a curious human and an artificial intelligence assistant. "
+    "The assistant gives helpful, detailed, and polite answers to the human's questions.",
+    roles=("Human", "Assistant"),
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.SINGLE,
+    sep="###",
+)
+
+conv_llava_v0_mmtag = Conversation(
+    system="A chat between a curious user and an artificial intelligence assistant. "
+    "The assistant is able to understand the visual content that the user provides, and assist the user with a variety of tasks using natural language."
+    "The visual content will be provided with the following format: <Image>visual content</Image>.",
+    roles=("Human", "Assistant"),
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.SINGLE,
+    sep="###",
+    version="v0_mmtag",
+)
+
+conv_llava_v1 = Conversation(
+    system="A chat between a curious human and an artificial intelligence assistant. "
+    "The assistant gives helpful, detailed, and polite answers to the human's questions.",
+    roles=("USER", "ASSISTANT"),
+    version="v1",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.TWO,
+    sep=" ",
+    sep2="</s>",
+)
+
+conv_llava_v1_mmtag = Conversation(
+    system="A chat between a curious user and an artificial intelligence assistant. "
+    "The assistant is able to understand the visual content that the user provides, and assist the user with a variety of tasks using natural language."
+    "The visual content will be provided with the following format: <Image>visual content</Image>.",
+    roles=("USER", "ASSISTANT"),
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.TWO,
+    sep=" ",
+    sep2="</s>",
+    version="v1_mmtag",
+)
+
+default_conversation = conv_vicuna_v1
+conv_templates = {
+    "no-sys": empty_conv,
+    "default": conv_vicuna_v1,
+    "v0": conv_vicuna_v0,
+    "v1": conv_vicuna_v1,
+    "vicuna_v1": conv_vicuna_v1,
+    "llama_2": conv_llama_2,
+    "plain": conv_llava_plain,
+    "v0_plain": conv_llava_plain,
+    "llava_v0": conv_llava_v0,
+    "v0_mmtag": conv_llava_v0_mmtag,
+    "llava_v1": conv_llava_v1,
+    "v1_mmtag": conv_llava_v1_mmtag,
+    "llava_llama_2": conv_llava_llama_2,
+    "mpt": conv_mpt,
+}
+
+
+def get_conversation(version: str):
+    conv = conv_templates.get(version, conv_vicuna_v1).copy()
+    return conv
+
+
+if __name__ == "__main__":
+    print(default_conversation.get_prompt())

llm-awq/tinychat/serve/model_worker.py

@@ -0,0 +1,433 @@
+#    Modified from https://github.com/haotian-liu/LLaVA
+#    Copyright 2023 Haotian Liu
+#
+#    Licensed under the Apache License, Version 2.0 (the "License");
+#    you may not use this file except in compliance with the License.
+#    You may obtain a copy of the License at
+#
+#        http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS,
+#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#    See the License for the specific language governing permissions and
+#    limitations under the License.
+
+"""
+A model worker executes the model.
+"""
+import argparse
+import asyncio
+import json
+import time
+import threading
+import uuid
+
+from fastapi import FastAPI, Request, BackgroundTasks
+from fastapi.responses import StreamingResponse
+import requests
+import torch
+import uvicorn
+from functools import partial
+from tqdm import tqdm
+
+import tinychat.utils.constants
+from tinychat.utils.constants import (
+    WORKER_HEART_BEAT_INTERVAL,
+    LLAVA_DEFAULT_IMAGE_TOKEN_IDX,
+    LLAVA_DEFAULT_IMAGE_TOKEN,
+    LLAVA_DEFAULT_IM_START_TOKEN,
+    LLAVA_DEFAULT_IM_END_TOKEN,
+)
+from tinychat.utils.log_utils import (
+    build_logger,
+    server_error_msg,
+    pretty_print_semaphore,
+)
+from tinychat.stream_generators.llava_stream_gen import tokenizer_image_token
+from tinychat.utils.llava_image_processing import process_images, load_image_from_base64
+from tinychat.models.llava_llama import LlavaLlamaForCausalLM
+from tinychat.stream_generators.llava_stream_gen import LlavaStreamGenerator
+from tinychat.utils.prompt_templates import (
+    get_prompter,
+    get_stop_token_ids,
+    get_image_token,
+)
+from tinychat.utils.conversation_utils import gen_params
+
+from transformers import AutoConfig, AutoTokenizer
+from accelerate import load_checkpoint_and_dispatch
+
+# import os
+# os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+
+GB = 1 << 30
+
+worker_id = str(uuid.uuid4())[:6]
+logger = build_logger("model_worker", f"model_worker_{worker_id}.log")
+global_counter = 0
+
+model_semaphore = None
+
+
+def heart_beat_worker(controller):
+    while True:
+        time.sleep(WORKER_HEART_BEAT_INTERVAL)
+        controller.send_heart_beat()
+
+
+def skip(*args, **kwargs):
+    pass
+
+
+class ModelWorker:
+    def __init__(
+        self,
+        controller_addr,
+        worker_addr,
+        worker_id,
+        no_register,
+        model_type,
+        model_path,
+        model_name,
+        quant_path,
+        precision,
+        device,
+    ):
+        self.controller_addr = controller_addr
+        self.worker_addr = worker_addr
+        self.worker_id = worker_id
+        self.model_type = model_type
+        self.model_path = model_path
+        if model_path.endswith("/"):
+            model_path = model_path[:-1]
+        if model_name is None:
+            model_paths = model_path.split("/")
+            if model_paths[-1].startswith("checkpoint-"):
+                self.model_name = model_paths[-2] + "_" + model_paths[-1]
+            else:
+                self.model_name = model_paths[-1]
+        else:
+            self.model_name = model_name
+        if precision == "W4A16":
+            self.model_name = self.model_name + "-4bit-AWQ"
+        self.device = device
+
+        # Load TinyChat model
+        logger.info(f"Loading the model {self.model_name} on worker {worker_id} ...")
+
+        setattr(torch.nn.Linear, "reset_parameters", lambda self: None)
+        setattr(torch.nn.LayerNorm, "reset_parameters", lambda self: None)
+        torch.nn.init.kaiming_uniform_ = skip
+        torch.nn.init.kaiming_normal_ = skip
+        torch.nn.init.uniform_ = skip
+        torch.nn.init.normal_ = skip
+
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
+        tinychat.utils.constants.LLAVA_DEFAULT_IMAGE_PATCH_TOKEN_IDX = (
+            self.tokenizer.convert_tokens_to_ids(
+                [tinychat.utils.constants.LLAVA_DEFAULT_IMAGE_PATCH_TOKEN]
+            )[0]
+        )
+        config = AutoConfig.from_pretrained(args.model_path, trust_remote_code=True)
+        config.min_max_range_path = args.model_path + "/emb_min_max.pt"
+        model = LlavaLlamaForCausalLM(config, args.device).half()
+        vision_tower = model.get_model().vision_tower
+        if not vision_tower.is_loaded:
+            vision_tower.load_model()
+        vision_tower = vision_tower.half()
+        self.image_processor = vision_tower.image_processor
+
+        if precision == "W16A16":
+            pbar = tqdm(range(1))
+            pbar.set_description("Loading checkpoint shards")
+            for i in pbar:
+                model = load_checkpoint_and_dispatch(
+                    model,
+                    model_path,
+                    no_split_module_classes=[
+                        "OPTDecoderLayer",
+                        "LlamaDecoderLayer",
+                        "BloomBlock",
+                        "MPTBlock",
+                        "DecoderLayer",
+                        "CLIPEncoderLayer",
+                    ],
+                ).to(device)
+        elif precision == "W4A16":
+            from tinychat.utils.load_quant import load_awq_model
+
+            model = load_awq_model(model, quant_path, 4, 128, device)
+            from tinychat.modules import (
+                make_quant_norm,
+                make_quant_attn,
+            )
+
+            make_quant_attn(model, device)
+            make_quant_norm(model)
+            model = model.to(device)
+        else:
+            raise NotImplementedError(f"Precision {precision} is not supported.")
+
+        self.model = model
+        self.is_multimodal = (
+            "llava" in self.model_name.lower() or "vila" in self.model_name.lower()
+        )
+
+        if not no_register:
+            self.register_to_controller()
+            self.heart_beat_thread = threading.Thread(
+                target=heart_beat_worker, args=(self,)
+            )
+            self.heart_beat_thread.start()
+
+    def register_to_controller(self):
+        logger.info("Register to controller")
+
+        url = self.controller_addr + "/register_worker"
+        data = {
+            "worker_name": self.worker_addr,
+            "check_heart_beat": True,
+            "worker_status": self.get_status(),
+        }
+        r = requests.post(url, json=data)
+        assert r.status_code == 200
+
+    def send_heart_beat(self):
+        logger.info(
+            f"Send heart beat. Models: {[self.model_name]}. "
+            f"Semaphore: {pretty_print_semaphore(model_semaphore)}. "
+            f"global_counter: {global_counter}"
+        )
+
+        url = self.controller_addr + "/receive_heart_beat"
+
+        while True:
+            try:
+                ret = requests.post(
+                    url,
+                    json={
+                        "worker_name": self.worker_addr,
+                        "queue_length": self.get_queue_length(),
+                    },
+                    timeout=5,
+                )
+                exist = ret.json()["exist"]
+                break
+            except requests.exceptions.RequestException as e:
+                logger.error(f"heart beat error: {e}")
+            time.sleep(5)
+
+        if not exist:
+            self.register_to_controller()
+
+    def get_queue_length(self):
+        if model_semaphore is None:
+            return 0
+        else:
+            return (
+                args.limit_model_concurrency
+                - model_semaphore._value
+                + (
+                    len(model_semaphore._waiters)
+                    if model_semaphore._waiters is not None
+                    else 0
+                )
+            )
+
+    def get_status(self):
+        return {
+            "model_names": [self.model_name],
+            "speed": 1,
+            "queue_length": self.get_queue_length(),
+        }
+
+    @torch.inference_mode()
+    def generate_stream(self, params):
+        tokenizer, model, image_processor = (
+            self.tokenizer,
+            self.model,
+            self.image_processor,
+        )
+
+        prompt = params["prompt"]
+        ori_prompt = prompt
+        images = params.get("images", None)
+        if images is not None and len(images) > 0 and self.is_multimodal:
+            if len(images) > 0:
+                if len(images) != prompt.count(LLAVA_DEFAULT_IMAGE_TOKEN):
+                    raise ValueError(
+                        "Number of images does not match number of <image> tokens in prompt"
+                    )
+
+                images = [load_image_from_base64(image) for image in images]
+                images = process_images(images, image_processor, model.config)
+
+                if type(images) is list:
+                    images = [
+                        image.to(model.device, dtype=torch.float16) for image in images
+                    ]
+                else:
+                    images = images.to(model.device, dtype=torch.float16)
+
+                replace_token = LLAVA_DEFAULT_IMAGE_TOKEN
+                if getattr(model.config, "mm_use_im_start_end", False):
+                    replace_token = (
+                        LLAVA_DEFAULT_IM_START_TOKEN
+                        + replace_token
+                        + LLAVA_DEFAULT_IM_END_TOKEN
+                    )
+                prompt = prompt.replace(LLAVA_DEFAULT_IMAGE_TOKEN, replace_token)
+            else:
+                images = None
+        else:
+            images = None
+
+        gen_params.temp = float(params.get("temperature", 1.0))
+        gen_params.top_p = float(params.get("top_p", 1.0))
+        gen_params.n_predict = min(int(params.get("max_new_tokens", 256)), 1024)
+
+        stream_generator = LlavaStreamGenerator
+        stop_token_ids = get_stop_token_ids(self.model_type, self.model_path)
+        image_token = get_image_token(model, self.model_path)
+        image_token_holder = (
+            tinychat.utils.constants.LLAVA_DEFAULT_IM_TOKEN_PLACE_HOLDER
+        )
+        prompt = prompt.replace(image_token_holder, image_token)
+
+        # print("=" * 50)
+        # print(prompt)
+        # print('=' * 50)
+        output_stream = stream_generator(
+            model,
+            tokenizer,
+            prompt,
+            gen_params,
+            device=model.device,
+            stop_token_ids=stop_token_ids,
+            image_tensor=images,
+        )
+
+        generated_text = ori_prompt
+        pre = 0
+        for outputs in output_stream:
+            output_text = outputs["text"]
+            output_text = output_text.strip().split(" ")
+            now = len(output_text) - 1
+            if now > pre:
+                generated_text += " ".join(output_text[pre:now]) + " "
+                yield json.dumps(
+                    {"text": generated_text, "error_code": 0}
+                ).encode() + b"\0"
+                pre = now
+        generated_text += " ".join(output_text[pre:])
+        yield json.dumps({"text": generated_text, "error_code": 0}).encode() + b"\0"
+
+    def generate_stream_gate(self, params):
+        try:
+            for x in self.generate_stream(params):
+                yield x
+        except ValueError as e:
+            print("Caught ValueError:", e)
+            ret = {
+                "text": server_error_msg,
+                "error_code": 1,
+            }
+            yield json.dumps(ret).encode() + b"\0"
+        except torch.cuda.CudaError as e:
+            print("Caught torch.cuda.CudaError:", e)
+            ret = {
+                "text": server_error_msg,
+                "error_code": 1,
+            }
+            yield json.dumps(ret).encode() + b"\0"
+        except Exception as e:
+            print("Caught Unknown Error", e)
+            ret = {
+                "text": server_error_msg,
+                "error_code": 1,
+            }
+            yield json.dumps(ret).encode() + b"\0"
+
+
+app = FastAPI()
+
+
+def release_model_semaphore(fn=None):
+    model_semaphore.release()
+    if fn is not None:
+        fn()
+
+
+@app.post("/worker_generate_stream")
+async def generate_stream(request: Request):
+    global model_semaphore, global_counter
+    global_counter += 1
+    params = await request.json()
+
+    if model_semaphore is None:
+        model_semaphore = asyncio.Semaphore(args.limit_model_concurrency)
+    await model_semaphore.acquire()
+    worker.send_heart_beat()
+    generator = worker.generate_stream_gate(params)
+    background_tasks = BackgroundTasks()
+    background_tasks.add_task(
+        partial(release_model_semaphore, fn=worker.send_heart_beat)
+    )
+    return StreamingResponse(generator, background=background_tasks)
+
+
+@app.post("/worker_get_status")
+async def get_status(request: Request):
+    return worker.get_status()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--host", type=str, default="localhost")
+    parser.add_argument("--port", type=int, default=21002)
+    parser.add_argument("--worker-address", type=str, default="http://localhost:21002")
+    parser.add_argument(
+        "--controller-address", type=str, default="http://localhost:21001"
+    )
+    parser.add_argument(
+        "--model-type",
+        type=str,
+        default="LLaMa",
+        help="type of the (base) language model",
+    )
+    parser.add_argument("--model-path", type=str, default="facebook/opt-350m")
+    parser.add_argument("--model-name", type=str)
+    parser.add_argument("--quant-path", type=str, default=None)
+    parser.add_argument("--precision", type=str, default="W4A16")
+    parser.add_argument("--device", type=str, default="cuda")
+    parser.add_argument(
+        "--multi-modal",
+        action="store_true",
+        help="Multimodal mode is automatically detected with model name, please make sure `llava` is included in the model path.",
+    )
+    parser.add_argument("--limit-model-concurrency", type=int, default=5)
+    parser.add_argument("--stream-interval", type=int, default=1)
+    parser.add_argument("--no-register", action="store_true")
+
+    args = parser.parse_args()
+    logger.info(f"args: {args}")
+
+    if args.multi_modal:
+        logger.warning(
+            "Multimodal mode is automatically detected with model name, please make sure `llava` is included in the model path."
+        )
+
+    worker = ModelWorker(
+        args.controller_address,
+        args.worker_address,
+        worker_id,
+        args.no_register,
+        args.model_type,
+        args.model_path,
+        args.model_name,
+        args.quant_path,
+        args.precision,
+        args.device,
+    )
+    uvicorn.run(app, host=args.host, port=args.port, log_level="info")