Add files using upload-large-folder tool

.python_history

@@ -0,0 +1,18 @@
+a=[4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 2]
+len(a)
+a= "   1   "
+int(a)
+b=int(a)
+b
+type(b)
+type(b)
+type(b)
+b
+3584*3584*2+3584+51282
+3584*18944*3
+203685888/25744978
+4096*4096*2+ 4096*1024*2
+4096*14336*3
+176160768/41943040
+a=[29, 23, 24, 30, 18, 28, 26, 20, 16, 27, 25, 17, 19, 21]
+len(a)

llm-awq/awq.egg-info/PKG-INFO

@@ -0,0 +1,319 @@
+Metadata-Version: 2.4
+Name: awq
+Version: 0.1.0
+Summary: An efficient and accurate low-bit weight quantization(INT3/4) method for LLMs.
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: Apache Software License
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: accelerate==0.34.2
+Requires-Dist: sentencepiece
+Requires-Dist: tokenizers>=0.12.1
+Requires-Dist: torch==2.3.0
+Requires-Dist: torchvision==0.18.0
+Requires-Dist: transformers==4.46.0
+Requires-Dist: lm_eval==0.3.0
+Requires-Dist: texttable
+Requires-Dist: toml
+Requires-Dist: attributedict
+Requires-Dist: protobuf
+Requires-Dist: gradio==3.35.2
+Requires-Dist: gradio_client==0.2.9
+Requires-Dist: fastapi
+Requires-Dist: uvicorn
+Requires-Dist: pydantic==1.10.19
+Dynamic: license-file
+
+# AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration 
+[[Paper](https://arxiv.org/abs/2306.00978)][[Website](https://hanlab.mit.edu/projects/awq)]
+
+**Efficient and accurate** low-bit weight quantization (INT3/4) for LLMs, supporting **instruction-tuned** models and **multi-modal** LMs.
+
+![overview](figures/overview.png)
+
+The current release supports: 
+
+- AWQ search for accurate quantization. 
+- Pre-computed AWQ model zoo for LLMs (Llama-1/2/3, OPT, CodeLlama, StarCoder, Vicuna, VILA, LLaVA; load to generate quantized weights).
+- Memory-efficient 4-bit Linear in PyTorch.
+- Efficient CUDA kernel implementation for fast inference (support context and decoding stage).
+- Examples on 4-bit inference of an instruction-tuned model (Vicuna) and **multi-modal LM** (VILA).
+- Chunk prefilling for faster prefilling in multi-round Q&A setting.
+- State-of-the-art prefilling speed of LLMs/VLMs on edge devices: [TinyChat 2.0](./tinychat).
+
+**Thanks to AWQ, TinyChat can deliver more efficient responses with LLM/VLM chatbots through 4-bit inference.**
+
+* 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](./tinychat/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](./tinychat/figures/orin_example_new.gif)
+
+
+**TinyChat also supports inference with vision language models (e.g., VILA, LLaVA). 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](./tinychat/figures/4090_nvila_single.gif)
+
+* TinyChat with NVILA-8B on RTX 4090 (multi-image inputs):
+
+![TinyChat with NVILA on 4090 multiple images](./tinychat/figures/4090_nvila_multi.gif)
+
+<!-- Check out [TinyChat](tinychat), which delievers **30 tokens/second** inference performance (**3.2x faster** than FP16) for the **Llama2** chatbot on the resource-constrained NVIDIA Jetson Orin!  -->
+
+* 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
+
+Check out [TinyChat](tinychat), which offers a turn-key solution for **on-device inference** of LLMs and VLMs on **resource-constrained edge platforms**. With TinyChat, it is now possible to efficiently run **large** models on **small** and **low-power** devices even without Internet connection!
+
+
+## News
+- [2025/04] 🔥 AWQ now supports DeepSeek-R1-Distilled models. Try our example [here](https://github.com/mit-han-lab/llm-awq/blob/main/scripts/DeepSeek_R1_Distill_example.sh)!
+- [2025/02] AWQ now supports BF16 precision. See example [here](https://github.com/mit-han-lab/llm-awq/blob/main/scripts/qwen_example.sh).
+- [2024/10] 🔥⚡ Explore advancements in [TinyChat 2.0](./tinychat), the latest version with significant advancements in prefilling speed of Edge LLMs and VLMs, **1.5-1.7x** faster than the previous version of TinyChat. Please refer to the [README](./tinychat/README.md) and [blog](https://hanlab.mit.edu/blog/tinychat20) for more details. 
+- [2024/05] 🏆 AWQ receives the **Best Paper Award** at **MLSys 2024**. 🎉 
+- [2024/05] 🔥 The **VILA-1.5** model family which features **video understanding** is now supported in AWQ and TinyChat. Check out out online demo powered by TinyChat [here](https://vila.hanlab.ai). Example is [here](scripts/vila15_example.sh).
+- [2024/05] 🔥 [AMD](https://community.amd.com/t5/ai/reduce-memory-footprint-and-improve-performance-running-llms-on/ba-p/686157) adopts AWQ to improve LLM serving efficiency.
+- [2024/04] 🔥 We released AWQ and TinyChat support for The **Llama-3** model family! Check out our example [here](scripts/llama3_example.sh).
+- [2024/02] 🔥 AWQ has been accepted to **MLSys 2024**!
+- [2024/02] 🔥 We supported [VILA Vision Languague Models](https://arxiv.org/abs/2312.07533) in AWQ & TinyChat! Check our latest demos with multi-image inputs!
+- [2024/02] 🔥 We released new version of quantized GEMM/GEMV kernels in [**TinyChat**](tinychat), leading to **38 tokens/second** inference speed on NVIDIA Jetson Orin!
+- [2024/01] 🔥 AWQ has been integrated by [Google Vertex AI](https://console.cloud.google.com/vertex-ai/publishers/meta/model-garden/llama-2-quantized)!
+- [2023/11] 🔥 AWQ has been integrated by [Amazon Sagemaker Containers](https://aws.amazon.com/blogs/machine-learning/boost-inference-performance-for-llms-with-new-amazon-sagemaker-containers/)!
+- [2023/11] 🔥 We added AWQ support and pre-computed search results for CodeLlama, StarCoder, StableCode models. Checkout our model zoo [here](https://huggingface.co/datasets/mit-han-lab/awq-model-zoo)!
+- [2023/11] 🔥 AWQ is now integrated natively in Hugging Face transformers through `from_pretrained`. You can either load quantized models from the Hub or your own HF quantized models.
+- [2023/10] AWQ is integrated into NVIDIA [TensorRT-LLM](https://github.com/NVIDIA/TensorRT-LLM/)
+- [2023/09] AWQ is integrated into [Intel Neural Compressor](https://github.com/intel/neural-compressor), [FastChat](https://github.com/lm-sys/FastChat/blob/main/docs/awq.md), [vLLM](https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/quantization/awq.py), [HuggingFace TGI](https://github.com/huggingface/text-generation-inference/pull/1054), and [LMDeploy](https://github.com/InternLM/lmdeploy). 
+- [2023/09] ⚡ Check out our latest [**TinyChat**](tinychat), which is ~2x faster than the first release on Orin!
+- [2023/09] ⚡ Check out [**AutoAWQ**](https://github.com/casper-hansen/AutoAWQ), a third-party implementation to make AWQ easier to expand to new models, improve inference speed, and integrate into Huggingface.
+- [2023/07] 🔥 We released **TinyChat**, an efficient and lightweight chatbot interface based on AWQ. TinyChat enables efficient LLM inference on both cloud and edge GPUs. Llama-2-chat models are supported! Check out our implementation [here](tinychat).
+- [2023/07] 🔥 We added AWQ support and pre-computed search results for Llama-2 models (7B & 13B). Checkout our model zoo [here](https://huggingface.co/datasets/mit-han-lab/awq-model-zoo)!
+- [2023/07] We extended the support for more LLM models including MPT, Falcon, and BLOOM. 
+
+## Contents
+
+- [AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration](#awq-activation-aware-weight-quantization-for-llm-compression-and-acceleration)
+  - [News](#news)
+  - [Contents](#contents)
+  - [Helpful Links](#helpful-links)
+  - [Install](#install)
+  - [AWQ Model Zoo](#awq-model-zoo)
+  - [Examples](#examples)
+  - [Usage](#usage)
+  - [Results on Visual Language Models](#results-on-visual-language-models)
+  - [Reference](#reference)
+  - [Related Projects](#related-projects)
+
+## Helpful Links
+
+- [VILA online demo](vila.hanlab.ai): Visual Language Models efficiently supported by AWQ & TinyChat.
+- [LLM on the Edge](https://github.com/mit-han-lab/llm-awq/tree/nv_laptop?tab=readme-ov-file#install): AWQ and TinyChat support edge GPUs such as NVIDIA Jetson Orin.
+- [VLMs on Laptop](https://github.com/mit-han-lab/llm-awq/tree/nv_laptop?tab=readme-ov-file#run-vila-on-laptop): Follow the instructions to deploy VLMs on NVIDIA Laptops with TinyChat.
+- [Gradio Server](https://github.com/mit-han-lab/llm-awq/tree/nv_laptop/tinychat/serve#gradio-demo-vila-with-tinychat): Try to build your own VLM online demo with AWQ and TinyChat!
+- [QServe](https://github.com/mit-han-lab/qserve): 🔥 **[New]** Efficient and accurate serving system for large-scale LLM inference.
+  
+## Install
+
+1. Clone this repository and navigate to AWQ folder
+```
+git clone https://github.com/mit-han-lab/llm-awq
+cd llm-awq
+```
+
+2. Install Package
+```
+conda create -n awq python=3.10 -y
+conda activate awq
+pip install --upgrade pip  # enable PEP 660 support
+pip install -e .
+```
+
+* For **edge devices** like Orin, before running the commands above, please:
+
+    1. Modify [pyproject.toml](pyproject.toml) by commenting out [this line](https://github.com/mit-han-lab/llm-awq/blob/3fce69061682fdd528824e5da3d03a8a8b545f2a/pyproject.toml#L17).
+    2. Manually install precompiled PyTorch binaries (>=2.0.0) from [NVIDIA](https://forums.developer.nvidia.com/t/pytorch-for-jetson/72048). You also need to install torchvision from this website when running NVILA.
+    3. Set the appropriate Python version for conda environment (e.g., `conda create -n awq python=3.8 -y` for JetPack 5).
+  
+3. Install efficient W4A16 (4-bit weight, 16-bit activation) CUDA kernel and optimized FP16 kernels (e.g. layernorm, positional encodings).
+```
+cd awq/kernels
+python setup.py install
+```
+
+4. Install Flash Attention
+```
+pip install flash-attn --no-build-isolation
+```
+
+We recommend starting an interactive python CLI interface and run `import flash_attn` to check whether FlashAttention-2 is installed successfully. If not, we recommend downloading pre-built wheels from [here](https://github.com/Dao-AILab/flash-attention/releases/tag/v2.5.8). Please notice:
+
+- PyTorch version needs to exactly match with the version specified in the `.whl` name;
+- Check out both `cxx11abiTRUE` and `cxx11abiFALSE` wheels if one of them does not work;
+- It's recommended to match CUDA version specified in the `.whl` filename, but minor mismatches (e.g. 12.1 vs 12.2, or even 11.8 vs 12.2) usually do not matter.
+
+
+5. [Optional] In order to run AWQ and TinyChat with NVILA model family, please install VILA:
+
+```bash
+git clone https://github.com/NVlabs/VILA.git
+cd VILA
+pip install -e .
+```
+
+## AWQ Model Zoo
+
+We provide pre-computed AWQ search results for multiple model families, including LLaMA, OPT, Vicuna, 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
+```
+
+The detailed support list:
+
+| Models | Sizes                       | INT4-g128 | INT3-g128 |
+| ------ | --------------------------- | --------- | --------- |
+| [DeepSeek-R1-Distill](/scripts/DeepSeek_R1_Distill_example.sh) | 1.5B/7B/8B |  ✅         |         |
+| [Qwen-2.5](/scripts/qwen_example.sh) | 7B/72B            | ✅         |          |
+| [NVILA](/scripts/nvila_example.sh) | 3B/8B |  ✅         |         |
+| [VILA-1.5](/scripts/vila15_example.sh)  | 3B/8B/13B/40B  | ✅         | ✅        |
+| [Llama3](/scripts/llama_example.sh)  | 8B/70B  | ✅         | ✅        |
+| [VILA](/scripts/vila_example.sh)    | 7B/13B                     | ✅         |           |
+| [Llama2](/scripts/llama_example.sh)  | 7B/13B/70B  | ✅         | ✅        |
+| [LLaMA](/scripts/llama2_example.sh)  | 7B/13B/30B/65B              | ✅         | ✅        |
+| [OPT](/scripts/opt_example.sh)    | 125m/1.3B/2.7B/6.7B/13B/30B | ✅         | ✅        |
+| [CodeLlama](/scripts/codellama_example.sh) | 7B/13B/34B               | ✅         | ✅        |
+| [StarCoder](/scripts/starcoder_example.sh) | 15.5B                    | ✅         | ✅        |
+| [Vicuna-v1.1](/scripts/vicuna_example.sh) | 7B/13B                 | ✅         |           |
+| [LLaVA-v0](/scripts/llava_example.sh) | 13B                       | ✅         |           |
+
+Note: We only list models that we have prepare the [AWQ searching results](https://huggingface.co/datasets/mit-han-lab/awq-model-zoo/tree/main) in the table above. AWQ also supports models such as LLaVA-v1.5 7B, and you may need to run the [AWQ search](#usage) on your own to quantize these models. For our latest VLM NVILA, quantized weights are available [here](https://huggingface.co/Efficient-Large-Model/NVILA-AWQ).
+
+## Examples
+
+AWQ can be easily applied to various LMs thanks to its good generalization, including instruction-tuned models and multi-modal LMs. It provides an easy-to-use tool to reduce the serving cost of LLMs.
+
+Here we provide two examples of AWQ application: Vicuna-7B (chatbot) and LLaVA-13B (visual reasoning) under `./examples` directory. AWQ can easily reduce the GPU memory of model serving and speed up token generation. It provides accurate quantization, providing reasoning outputs. You should be able to observe **memory savings** when running the models with 4-bit weights. 
+
+Note that we perform AWQ using only textual calibration data, depsite we are running on multi-modal input. Please refer to `./examples` for details.
+
+![overview](figures/example_vis.jpg)
+
+## Usage
+
+We provide several sample script to run AWQ (please refer to `./scripts`). We use Llama3-8B as an example.
+
+1. Perform AWQ search and save search results (we already did it for you):
+```bash
+python -m awq.entry --model_path /PATH/TO/LLAMA3/llama3-8b \
+    --w_bit 4 --q_group_size 128 \
+    --run_awq --dump_awq awq_cache/llama3-8b-w4-g128.pt
+```
+
+2. Evaluate the AWQ quantized model on WikiText-2 (simulated pseudo quantization)
+```bash
+python -m awq.entry --model_path /PATH/TO/LLAMA3/llama3-8b \
+    --tasks wikitext \
+    --w_bit 4 --q_group_size 128 \
+    --load_awq awq_cache/llama3-8b-w4-g128.pt \
+    --q_backend fake
+```
+
+3. Generate real quantized weights (INT4)
+```bash
+mkdir quant_cache
+python -m awq.entry --model_path /PATH/TO/LLAMA3/llama3-8b \
+    --w_bit 4 --q_group_size 128 \
+    --load_awq awq_cache/llama3-8b-w4-g128.pt \
+    --q_backend real --dump_quant quant_cache/llama3-8b-w4-g128-awq.pt
+```
+
+4. Load and evaluate the real quantized model (now you can see smaller gpu memory usage)
+```bash
+python -m awq.entry --model_path /PATH/TO/LLAMA3/llama3-8b \
+    --tasks wikitext \
+    --w_bit 4 --q_group_size 128 \
+    --load_quant quant_cache/llama3-8b-w4-g128-awq.pt
+```
+## Results on Visual Language Models
+
+AWQ also seamlessly supports large multi-modal models (LMMs). Please refer to [TinyChat](./tinychat/README.md) for more details.
+
+
+<!-- AWQ also seamlessly supports large multi-modal models (LMMs). We demonstrate the results on the recent [VILA-1.5](https://github.com/Efficient-Large-Model/VILA) model family. -->
+
+<!-- 
+| 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 | 
+
+
+## Inference speed ( Token/sec )
+
+| $~~~~~~$               | Precision | A100  | 4090  | Orin |
+| ---------------------- | --------- | ----- | ----- | ---- |
+| VILA1.5-3B           | fp16      | 104.6 | 137.6 | 25.4 |
+| VILA1.5-3B-AWQ       | int4      | 182.8 | 215.5 | 42.5 |
+| VILA1.5-3B-S2        | fp16      | 104.3 | 137.2 | 24.6 |
+| VILA1.5-3B-S2-AWQ    | int4      | 180.2 | 219.3 | 40.1 |
+| Llama-3-VILA1.5-8B     | fp16      | 74.9  | 57.4  | 10.2 |
+| Llama-3-VILA1.5-8B-AWQ | int4      | 168.9 | 150.2 | 28.7 |
+| VILA1.5-13B            | fp16      | 50.9  | OOM   | 6.1  |
+| VILA1.5-13B-AWQ        | int4      | 115.9 | 105.7 | 20.6 |
+| VILA1.5-40B            | fp16      | OOM  | OOM   | --  |
+| VILA1.5-40B-AWQ        | int4      | 57.0 | OOM | -- | -->
+
+
+## Reference
+
+If you find AWQ useful or relevant to your research, please kindly cite our paper:
+
+```
+@inproceedings{lin2023awq,
+  title={AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration},
+  author={Lin, Ji and Tang, Jiaming and Tang, Haotian and Yang, Shang and Chen, Wei-Ming and Wang, Wei-Chen and Xiao, Guangxuan and Dang, Xingyu and Gan, Chuang and Han, Song},
+  booktitle={MLSys},
+  year={2024}
+}
+```
+
+## Related Projects
+
+[SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models](https://github.com/mit-han-lab/smoothquant)
+
+[GPTQ: Accurate Post-training Compression for Generative Pretrained Transformers](https://arxiv.org/abs/2210.17323)
+
+[Vicuna and FastChat](https://github.com/lm-sys/FastChat#readme)
+
+[LLaVA: Large Language and Vision Assistant](https://github.com/haotian-liu/LLaVA)
+
+[VILA: On Pre-training for Visual Language Models](https://github.com/Efficient-Large-Model/VILA)
+

llm-awq/awq.egg-info/SOURCES.txt

@@ -0,0 +1,81 @@
+LICENSE
+README.md
+pyproject.toml
+awq/entry.py
+awq.egg-info/PKG-INFO
+awq.egg-info/SOURCES.txt
+awq.egg-info/dependency_links.txt
+awq.egg-info/requires.txt
+awq.egg-info/top_level.txt
+awq/kernels/setup.py
+awq/kernels/csrc/attention/setup.py
+awq/quantize/__init__.py
+awq/quantize/auto_clip.py
+awq/quantize/auto_scale.py
+awq/quantize/pre_quant.py
+awq/quantize/qmodule.py
+awq/quantize/quantizer.py
+awq/quantize/smooth.py
+awq/quantize/w8a8_linear.py
+awq/utils/__init__.py
+awq/utils/calib_data.py
+awq/utils/lm_eval_adaptor.py
+awq/utils/module.py
+awq/utils/parallel.py
+awq/utils/utils.py
+tinychat/benchmark.py
+tinychat/demo.py
+tinychat/internvl_benchmark.py
+tinychat/internvl_demo.py
+tinychat/nvila_benchmark.py
+tinychat/nvila_demo.py
+tinychat/offline-weight-repacker.py
+tinychat/split_ckpt.py
+tinychat/vila10_demo.py
+tinychat/vila15_demo.py
+tinychat/models/__init__.py
+tinychat/models/falcon.py
+tinychat/models/internvl3.py
+tinychat/models/llama.py
+tinychat/models/llava_llama.py
+tinychat/models/mpt.py
+tinychat/models/nvila_qwen2.py
+tinychat/models/qwen2.py
+tinychat/models/vila_llama.py
+tinychat/models/internvl/configuration_internvl.py
+tinychat/models/internvl/conversation.py
+tinychat/models/internvl/internvit.py
+tinychat/models/internvl/media.py
+tinychat/models/llava_base/llava_arch.py
+tinychat/models/llava_base/multimodal_encoder/builder.py
+tinychat/models/llava_base/multimodal_encoder/clip_encoder.py
+tinychat/models/llava_base/multimodal_projector/builder.py
+tinychat/models/nvila/builder.py
+tinychat/models/nvila/configuration_llava.py
+tinychat/models/nvila/llava_arch.py
+tinychat/modules/__init__.py
+tinychat/modules/fused_attn.py
+tinychat/modules/fused_internencoder.py
+tinychat/modules/fused_mlp.py
+tinychat/modules/fused_norm.py
+tinychat/modules/fused_siglipdecoder.py
+tinychat/modules/fused_vision_attn.py
+tinychat/serve/controller.py
+tinychat/serve/gradio_web_server.py
+tinychat/serve/llava_conv.py
+tinychat/serve/model_worker.py
+tinychat/serve/model_worker_new.py
+tinychat/stream_generators/NVILA_stream_gen.py
+tinychat/stream_generators/__init__.py
+tinychat/stream_generators/internvl_stream_gen.py
+tinychat/stream_generators/llava_stream_gen.py
+tinychat/stream_generators/stream_gen.py
+tinychat/utils/__init__.py
+tinychat/utils/constants.py
+tinychat/utils/conversation_utils.py
+tinychat/utils/input_metadata.py
+tinychat/utils/llava_image_processing.py
+tinychat/utils/load_quant.py
+tinychat/utils/log_utils.py
+tinychat/utils/prompt_templates.py
+tinychat/utils/tune.py

llm-awq/awq.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

llm-awq/awq.egg-info/requires.txt

@@ -0,0 +1,16 @@
+accelerate==0.34.2
+sentencepiece
+tokenizers>=0.12.1
+torch==2.3.0
+torchvision==0.18.0
+transformers==4.46.0
+lm_eval==0.3.0
+texttable
+toml
+attributedict
+protobuf
+gradio==3.35.2
+gradio_client==0.2.9
+fastapi
+uvicorn
+pydantic==1.10.19

llm-awq/awq/kernels/csrc/w8a8/reduction_utils.cuh

@@ -0,0 +1,170 @@
+/*
+ * Adapted from https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/reduce_kernel_utils.cuh
+ * Copyright (c) 2023, The vLLM team.
+ * Copyright (c) 2020-2023, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * 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.
+ */
+#pragma once
+#define FINAL_MASK 0xffffffff
+
+
+namespace vllm {
+
+template<typename T>
+__inline__ __device__ T warpReduceSum(T val) {
+#pragma unroll
+  for (int mask = 16; mask > 0; mask >>= 1)
+    val += __shfl_xor_sync(0xffffffff, val, mask, 32);
+  return val;
+}
+
+template <typename T, int NUM>
+__inline__ __device__ T warpReduceSumV2(T* val)
+{
+#pragma unroll
+    for (int i = 0; i < NUM; i++)
+    {
+#pragma unroll
+        for (int mask = 16; mask > 0; mask >>= 1)
+            val[i] += __shfl_xor_sync(FINAL_MASK, val[i], mask, 32);
+    }
+    return (T) (0.0f);
+}
+
+/* Calculate the sum of all elements in a block */
+template<typename T>
+__inline__ __device__ T blockReduceSum(T val) {
+  static __shared__ T shared[32];
+  int lane = threadIdx.x & 0x1f;
+  int wid = threadIdx.x >> 5;
+
+  val = warpReduceSum<T>(val);
+
+  if (lane == 0)
+    shared[wid] = val;
+
+  __syncthreads();
+
+  // Modify from blockDim.x << 5 to blockDim.x / 32. to prevent
+  // blockDim.x is not divided by 32
+  val = (threadIdx.x < (blockDim.x / 32.f)) ? shared[lane] : (T)(0.0f);
+  val = warpReduceSum<T>(val);
+  return val;
+}
+
+/* Calculate the sum of all elements in a block */
+template<typename T>
+__inline__ __device__ T blockAllReduceSum(T val) {
+  static __shared__ T shared[32];
+  int lane = threadIdx.x & 0x1f;
+  int wid = threadIdx.x >> 5;
+
+  val = warpReduceSum<T>(val);
+
+  if (lane == 0)
+    shared[wid] = val;
+
+  __syncthreads();
+
+  // Modify from blockDim.x << 5 to blockDim.x / 32. to prevent
+  // blockDim.x is not divided by 32
+  val = (lane < (blockDim.x / 32.f)) ? shared[lane] : (T)(0.0f);
+  val = warpReduceSum<T>(val);
+  return val;
+}
+
+template <typename T, int NUM>
+__inline__ __device__ T blockReduceSumV2(T* val)
+{
+    static __shared__ T shared[NUM][33];
+    int lane = threadIdx.x & 0x1f;
+    int wid = threadIdx.x >> 5;
+
+    warpReduceSumV2<T, NUM>(val);
+
+    if (lane == 0)
+    {
+#pragma unroll
+        for (int i = 0; i < NUM; i++)
+        {
+            shared[i][wid] = val[i];
+        }
+    }
+
+    __syncthreads();
+
+    bool is_mask = threadIdx.x < (blockDim.x / 32.f);
+#pragma unroll
+    for (int i = 0; i < NUM; i++)
+    {
+        val[i] = is_mask ? shared[i][lane] : (T) (0.0f);
+    }
+    warpReduceSumV2<T, NUM>(val);
+    return (T) 0.0f;
+}
+
+template<typename T>
+__inline__ __device__ T warpReduceMax(T val)
+{
+#pragma unroll
+    for (int mask = 16; mask > 0; mask >>= 1)
+        val = max(val, __shfl_xor_sync(0xffffffff, val, mask, 32));
+    return val;
+}
+/* Calculate the maximum of all elements in a block */
+template<typename T>
+__inline__ __device__ T blockReduceMax(T val)
+{
+    static __shared__ T shared[32];
+    int                 lane = threadIdx.x & 0x1f;  // in-warp idx
+    int                 wid  = threadIdx.x >> 5;    // warp idx
+    val = warpReduceMax(val);  // get maxx in each warp
+    if (lane == 0)  // record in-warp maxx by warp Idx
+        shared[wid] = val;
+    __syncthreads();
+    // Modify from blockDim.x << 5 to blockDim.x / 32. to prevent
+    // blockDim.x is not divided by 32
+    val = (threadIdx.x < (blockDim.x / 32.f)) ? shared[lane] : -1e20f;
+    val = warpReduceMax(val);
+    return val;
+}
+
+/* Calculate the maximum of all elements in a block */
+template <typename T>
+__inline__ __device__ T blockAllReduceMax(T val)
+{
+    static __shared__ T shared[32];
+    int lane = threadIdx.x & 0x1f; // in-warp idx
+    int wid = threadIdx.x >> 5;    // warp idx
+
+    val = warpReduceMax(val);      // get maxx in each warp
+
+    if (lane == 0)                 // record in-warp maxx by warp Idx
+        shared[wid] = val;
+
+    __syncthreads();
+
+    // Modify from blockDim.x << 5 to blockDim.x / 32. to prevent
+    // blockDim.x is not divided by 32
+    val = (lane < (blockDim.x / 32.f)) ? shared[lane] : -1e20f;
+    val = warpReduceMax(val);
+
+    return val;
+}
+
+
+
+
+
+} // namespace vllm

llm-awq/awq/utils/calib_data.py

@@ -0,0 +1,32 @@
+import torch
+from datasets import load_dataset
+
+
+def get_calib_dataset(data="pileval", tokenizer=None, n_samples=512, block_size=512):
+    if data == "pileval":
+        dataset = load_dataset("mit-han-lab/pile-val-backup", split="validation")
+    else:
+        raise NotImplementedError
+    dataset = dataset.shuffle(seed=42)
+    samples = []
+    n_run = 0
+    for data in dataset:
+        line = data["text"]
+        line = line.strip()
+        line_encoded = tokenizer.encode(line)
+        if len(line_encoded) > 512:
+            continue
+        sample = torch.tensor([line_encoded])
+        if sample.numel() == 0:
+            continue
+        samples.append(sample)
+        n_run += 1
+        if n_run == n_samples:
+            break
+    # now concatenate all samples and split according to block size
+    cat_samples = torch.cat(samples, dim=1)
+    n_split = cat_samples.shape[1] // block_size
+    print(f" * Split into {n_split} blocks")
+    return [
+        cat_samples[:, i * block_size : (i + 1) * block_size] for i in range(n_split)
+    ]

llm-awq/examples/README.md

@@ -0,0 +1,10 @@
+# AWQ Examples
+
+Here we provide two AWQ examples, applying to:
+- [Vicuna-7B](https://github.com/lm-sys/FastChat), a chatbot with instruction-tuning
+- [LLaVA-13B](https://github.com/lm-sys/FastChat), a visual LM for multi-modal applications like visual reasoning.
+- [A simple conversion script](https://github.com/mit-han-lab/llm-awq/tree/main/examples/convert_to_hf.py) to convert llm-awq weights into HF format. 
+
+Here are some example output from the two demos. You should able to observe memory saving when running the demos in 4-bit. Please check the notebooks for details. 
+
+![overview](../figures/example_vis.jpg)

llm-awq/examples/chat_demo.ipynb

@@ -0,0 +1,272 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# AWQ on Vicuna"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "In this notebook, we use Vicuna model to demonstrate the performance of AWQ on instruction-tuned models. We implement AWQ real-INT4 inference kernels, which are wrapped as Pytorch modules and can be easily used by existing models. We also provide a simple example to show how to use AWQ to quantize a model and save/load the quantized model checkpoint."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "In order to run this notebook, you need to install the following packages:\n",
+    "- [AWQ](https://github.com/mit-han-lab/llm-awq)\n",
+    "- [Pytorch](https://pytorch.org/)\n",
+    "- [Accelerate](https://github.com/huggingface/accelerate)\n",
+    "- [Transformers](https://github.com/huggingface/transformers)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "from accelerate import init_empty_weights, load_checkpoint_and_dispatch\n",
+    "from awq.quantize.quantizer import real_quantize_model_weight\n",
+    "from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig\n",
+    "from tinychat.demo import gen_params, stream_output\n",
+    "from tinychat.stream_generators import StreamGenerator\n",
+    "from tinychat.modules import make_quant_norm, make_quant_attn, make_fused_mlp\n",
+    "from tinychat.utils.prompt_templates import get_prompter\n",
+    "import os\n",
+    "# This demo only support single GPU for now\n",
+    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\""
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Please get the Vicuna model from [FastChat](https://github.com/lm-sys/FastChat) and run the following command to generate a quantized model checkpoint first.\n",
+    "\n",
+    "```bash\n",
+    "mkdir quant_cache\n",
+    "python -m awq.entry --model_path [vicuna-7b_model_path] \\\n",
+    "    --w_bit 4 --q_group_size 128 \\\n",
+    "    --load_awq awq_cache/vicuna-7b-w4-g128.pt \\\n",
+    "    --q_backend real --dump_quant quant_cache/vicuna-7b-w4-g128-awq.pt\n",
+    "```"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# model_path = \"\" # the path of vicuna-7b model\n",
+    "# load_quant_path = \"quant_cache/vicuna-7b-w4-g128-awq.pt\"\n",
+    "model_path = \"/data/llm/checkpoints/vicuna-hf/vicuna-7b\"\n",
+    "load_quant_path = \"/data/llm/checkpoints/vicuna-hf/vicuna-7b-awq-w4g128.pt\""
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We first load a empty model and replace all the linear layers with WQLinear layers. Then we load the quantized weights from the checkpoint. "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "8b79a82b73ab4d9191ba54f5d0f8cb86",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "real weight quantization...(init only): 100%|███████████████████| 32/32 [00:11<00:00,  2.69it/s]\n",
+      "The model weights are not tied. Please use the `tie_weights` method before using the `infer_auto_device` function.\n",
+      "The model weights are not tied. Please use the `tie_weights` method before using the `infer_auto_device` function.\n"
+     ]
+    }
+   ],
+   "source": [
+    "config = AutoConfig.from_pretrained(model_path)\n",
+    "tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)\n",
+    "with init_empty_weights():\n",
+    "    model = AutoModelForCausalLM.from_pretrained(model_path, config=config,\n",
+    "                                                    torch_dtype=torch.float16)\n",
+    "q_config = {\"zero_point\": True, \"q_group_size\": 128}\n",
+    "real_quantize_model_weight(\n",
+    "    model, w_bit=4, q_config=q_config, init_only=True)\n",
+    "\n",
+    "model = load_checkpoint_and_dispatch(\n",
+    "    model, load_quant_path,\n",
+    "    device_map=\"auto\",\n",
+    "    no_split_module_classes=[\"LlamaDecoderLayer\"]\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[Warning] Calling a fake MLP fusion. But still faster than Huggingface Implimentation.\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "LlamaForCausalLM(\n",
+       "  (model): LlamaModel(\n",
+       "    (embed_tokens): Embedding(32000, 4096, padding_idx=0)\n",
+       "    (layers): ModuleList(\n",
+       "      (0-31): 32 x LlamaDecoderLayer(\n",
+       "        (self_attn): QuantLlamaAttention(\n",
+       "          (qkv_proj): WQLinear(in_features=4096, out_features=12288, bias=False, w_bit=4, group_size=128)\n",
+       "          (o_proj): WQLinear(in_features=4096, out_features=4096, bias=False, w_bit=4, group_size=128)\n",
+       "          (rotary_emb): QuantLlamaRotaryEmbedding()\n",
+       "        )\n",
+       "        (mlp): QuantLlamaMLP(\n",
+       "          (down_proj): WQLinear(in_features=11008, out_features=4096, bias=False, w_bit=4, group_size=128)\n",
+       "        )\n",
+       "        (input_layernorm): FTLlamaRMSNorm()\n",
+       "        (post_attention_layernorm): FTLlamaRMSNorm()\n",
+       "      )\n",
+       "    )\n",
+       "    (norm): FTLlamaRMSNorm()\n",
+       "  )\n",
+       "  (lm_head): Linear(in_features=4096, out_features=32000, bias=False)\n",
+       ")"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "make_quant_attn(model, \"cuda:0\")\n",
+    "make_quant_norm(model)\n",
+    "make_fused_mlp(model)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdin",
+     "output_type": "stream",
+     "text": [
+      "USER:  Show me some attractions in Boston.\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "ASSISTANT: 1. Boston Public Library\n",
+      "2. Fenway Park\n",
+      "3. Harvard Square\n",
+      "4. Boston Common\n",
+      "5. Freedom Trail\n",
+      "6. Museum of Fine Arts\n",
+      "7. Isabella Stewart Gardner Museum\n",
+      "8. Paul Revere House\n",
+      "9. New England Aquarium\n",
+      "10. Museum of Science\n",
+      "==================================================\n",
+      "Speed of Inference\n",
+      "--------------------------------------------------\n",
+      "Context Stage    : 7.18 ms/token\n",
+      "Generation Stage : 9.49 ms/token\n",
+      "Average Speed    : 8.53 ms/token\n",
+      "==================================================\n"
+     ]
+    },
+    {
+     "name": "stdin",
+     "output_type": "stream",
+     "text": [
+      "USER:  \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "EXIT...\n"
+     ]
+    }
+   ],
+   "source": [
+    "model_prompter = get_prompter(\"llama\", model_path)\n",
+    "stream_generator = StreamGenerator\n",
+    "count = 0\n",
+    "while True:\n",
+    "    # Get input from the user\n",
+    "    input_prompt = input(\"USER: \")\n",
+    "    if input_prompt == \"\":\n",
+    "        print(\"EXIT...\")\n",
+    "        break\n",
+    "    model_prompter.insert_prompt(input_prompt)\n",
+    "    output_stream = stream_generator(model, tokenizer, model_prompter.model_input, gen_params, device=\"cuda:0\")\n",
+    "    outputs = stream_output(output_stream)    \n",
+    "    model_prompter.update_template(outputs)\n",
+    "    count += 1"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python (awq)",
+   "language": "python",
+   "name": "awq"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

llm-awq/examples/convert_to_hf.py

@@ -0,0 +1,69 @@
+# This script demonstrates how you can convert your model into HF format
+# easily and push the quantized weights on the Hub using simple tools.
+# Make sure to have transformers > 4.34 and that you have ran
+# `huggingface-cli login` on your terminal before running this
+# script
+import os
+import argparse
+
+# This demo only support single GPU for now
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+
+from transformers import AutoConfig, AwqConfig, AutoTokenizer
+from huggingface_hub import HfApi
+
+api = HfApi()
+
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    "--model_path", type=str, help="path of the original hf model", required=True
+)
+parser.add_argument(
+    "--quantized_model_path",
+    type=str,
+    help="path of the quantized AWQ model",
+    required=True,
+)
+parser.add_argument(
+    "--quantized_model_hub_path",
+    type=str,
+    help="path of the quantized AWQ model to push on the Hub",
+    required=True,
+)
+parser.add_argument("--w_bit", type=int, default=4, help="")
+parser.add_argument("--q_group_size", default=128, type=int)
+parser.add_argument("--no_zero_point", action="store_true")
+
+args = parser.parse_args()
+
+original_model_path = args.model_path
+quantized_model_path = args.quantized_model_path
+quantized_model_hub_path = args.quantized_model_hub_path
+
+# Load the corresponding AWQConfig
+quantization_config = AwqConfig(
+    bits=args.w_bit,
+    group_size=args.q_group_size,
+    zero_point=not args.no_zero_point,
+    backend="llm-awq",
+    version="gemv",
+)
+
+# Set the attribute `quantization_config` in model's config
+config = AutoConfig.from_pretrained(original_model_path)
+config.quantization_config = quantization_config
+
+# Load tokenizer
+tok = AutoTokenizer.from_pretrained(original_model_path)
+
+# Push config and tokenizer
+config.push_to_hub(quantized_model_hub_path)
+tok.push_to_hub(quantized_model_hub_path)
+
+# Upload model weights
+api.upload_file(
+    path_or_fileobj=quantized_model_path,
+    path_in_repo="pytorch_model.bin",
+    repo_id=quantized_model_hub_path,
+    repo_type="model",
+)

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/llama_example.sh

@@ -0,0 +1,25 @@
+MODEL=llama-7b
+
+# run AWQ search (optional; we provided the pre-computed results)
+python -m awq.entry --model_path /dataset/llama-hf/$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/llama-hf/$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/llama-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
+
+# load and evaluate the real quantized model (smaller gpu memory usage)
+python -m awq.entry --model_path /dataset/llama-hf/$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/opt_example.sh

@@ -0,0 +1,25 @@
+MODEL=opt-6.7b
+
+# run AWQ search (optional; we provided the pre-computed results)
+python -m awq.entry --model_path /dataset/opt/$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/opt/$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/opt/$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/opt/$MODEL \
+    --tasks wikitext \
+    --w_bit 4 --q_group_size 128 \
+    --load_quant quant_cache/$MODEL-w4-g128-awq.pt

llm-awq/scripts/qwen_example.sh

@@ -0,0 +1,25 @@
+MODEL=qwen2.5-7b
+
+# 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/starcoder_example.sh

@@ -0,0 +1,25 @@
+MODEL=starcoder
+
+# run AWQ search (optional; we provided the pre-computed results)
+python -m awq.entry --model_path /dataset/starcoder-hf/$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/starcoder-hf/$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/starcoder-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
+
+# load and evaluate the real quantized model (smaller gpu memory usage)
+python -m awq.entry --model_path /dataset/starcoder-hf/$MODEL \
+    --tasks wikitext \
+    --w_bit 4 --q_group_size 128 \
+    --load_quant quant_cache/$MODEL-w4-g128-awq.pt

llm-awq/scripts/vicuna_example.sh

@@ -0,0 +1,25 @@
+MODEL=vicuna-7b
+
+# run AWQ search (optional; we provided the pre-computed results)
+python -m awq.entry --model_path /dataset/vicuna-hf/$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/vicuna-hf/$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/vicuna-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
+
+# load and evaluate the real quantized model (smaller gpu memory usage)
+python -m awq.entry --model_path /dataset/vicuna-hf/$MODEL \
+    --tasks wikitext \
+    --w_bit 4 --q_group_size 128 \
+    --load_quant quant_cache/$MODEL-w4-g128-awq.pt

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/benchmark.py

@@ -0,0 +1,379 @@
+# Usage:
+# Please first install awq/kernels
+# then directly run CUDA_VISIBLE_DEVICES=0 python benchmark.py
+import argparse
+import torch
+import time
+import numpy as np
+from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig, modeling_utils
+import tinychat.utils.constants
+from tinychat.utils.load_quant import load_awq_model
+from awq.quantize.quantizer import real_quantize_model_weight
+from tinychat.utils.tune import (
+    tune_all_wqlinears,
+    device_warmup,
+    tune_llava_patch_embedding,
+)
+from tinychat.modules import make_quant_norm, make_quant_attn, make_fused_mlp
+
+
+def skip(*args, **kwargs):
+    pass
+
+
+def 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/vicuna-hf/vicuna-7b",
+        help="path to the model",
+    )
+    parser.add_argument("--q_group_size", type=int, default=128)
+    parser.add_argument(
+        "--verbose",
+        default=False,
+        action="store_true",
+        help="Wheter to print more information.",
+    )
+    parser.add_argument(
+        "--max_seq_len",
+        type=int,
+        default=8192,
+        help="maximum sequence length for kv cache",
+    )
+    parser.add_argument(
+        "--max_batch_size", type=int, default=1, help="maximum batch size for kv cache"
+    )
+    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",
+    )
+    parser.add_argument(
+        "--context_length",
+        type=list,
+        nargs="+",
+        help="The length of input. And if chunk_prefilling used, this serves as the length of tokens from history rounds.",
+    )
+    parser.add_argument(
+        "--question_length",
+        type=list,
+        nargs="+",
+        help="The length of new input. Only useful and necessary when benchmarking chunk_prefilling method",
+    )
+    parser.add_argument(
+        "--precision", type=str, default="W4A16", help="compute precision"
+    )
+    args = parser.parse_args()
+    # some checks
+    assert (args.question_length is not None and args.chunk_prefilling) or (
+        not args.chunk_prefilling
+    ), "If you want to benchmark chunk prefilling, you need specify the question length and context length"
+    assert args.precision in ["W4A16", "W16A16"], "We only support W4A16/W16A16 now"
+    token_num = 256
+    # We support fixing a certain kind of length
+    if args.chunk_prefilling:
+        if len(args.context_length) == 1 and len(args.question_length) > 1:
+            args.context_length = [
+                args.context_length[0] for _ in range(len(args.question_length))
+            ]
+        elif len(args.question_length) == 1 and len(args.context_length) > 1:
+            args.question_length = [
+                args.question_length[0] for _ in range(len(args.context_length))
+            ]
+        elif len(args.question_length) != len(args.context_length):
+            raise ValueError(
+                "The number of items in the question_length and context_length is expected to be either one or equal!"
+            )
+    tinychat.utils.constants.max_batch_size = args.max_batch_size
+    tinychat.utils.constants.max_seq_len = args.max_seq_len
+    from tinychat.models import FalconForCausalLM, LlamaForCausalLM, MPTForCausalLM
+    from tinychat.models.vila_llama import VilaLlamaForCausalLM
+
+    modeling_utils._init_weights = False
+    torch.nn.init.kaiming_uniform_ = skip
+    torch.nn.init.kaiming_normal_ = skip
+    torch.nn.init.uniform_ = skip
+    torch.nn.init.normal_ = skip
+
+    device = "cuda:0"
+    model_type_dict = {
+        "llama": LlamaForCausalLM,
+        "falcon": FalconForCausalLM,
+        "mpt": MPTForCausalLM,
+    }
+
+    config = AutoConfig.from_pretrained(args.model_path, trust_remote_code=True)
+    assert args.model_type.lower() in [
+        "llama",
+        "falcon",
+        "mpt",
+        "vila",
+    ], "We only support llama & falcon & mpt & vila now"
+    if "vila" in args.model_type.lower():
+        model = VilaLlamaForCausalLM(config).half()
+        print(model)
+        if args.precision in ["W4A16"]:
+            real_quantize_model_weight(
+                model.llm,
+                w_bit=4,
+                q_config=dict(q_group_size=args.q_group_size, zero_point=True),
+                init_only=True,
+            )
+            make_quant_attn(model.llm, device, args.flash_attn)
+            make_quant_norm(model.llm)
+            make_fused_mlp(model.llm)
+        model = model.to(device)
+        device_warmup(device)
+        tune_llava_patch_embedding(model.get_vision_tower(), device=device)
+        if not args.chunk_prefilling:
+            image_num = [
+                int(int("".join(i)) * 1 / 196) for i in args.context_length
+            ]  # consider about three thirds of the history tokens are images
+            if sum(image_num) > 0:
+                image_tensor = 2 * torch.rand((max(image_num), 3, 384, 384)) - 1
+                image_tensor = image_tensor.half().to(device)
+            else:
+                image_tensor = None
+
+        print("huggingface ckpt loaded")
+
+        # warming up
+        input_ids = [1 for _ in range(2048)]
+        inputs = torch.as_tensor([input_ids], device=device)
+        out = model(
+            inputs, start_pos=0, chunk_prefilling=args.chunk_prefilling
+        )  # warmup
+
+        if not args.chunk_prefilling:
+            for i, context_length in enumerate(args.context_length):
+                context_length = int("".join(context_length))
+                time_lis = []
+                if image_num[i]:
+                    images = image_tensor[0 : image_num[i], :, :, :]
+                    input_ids = [-200 for _ in range(image_num[i])] + [
+                        1 for _ in range(context_length - 196 * image_num[i])
+                    ]
+                else:
+                    images = None
+                    input_ids = [1 for _ in range(context_length)]
+                print("-" * 80)
+                print(
+                    "Context length: {} with {} pictures".format(
+                        context_length, image_num[i]
+                    )
+                )
+                with torch.inference_mode():
+                    for i in range(10):  # Run ten times and get the average value
+                        start_pos = 0
+                        torch.cuda.synchronize()
+                        t_st = time.time()
+                        inputs = torch.as_tensor([input_ids], device=device)
+                        out = model(
+                            inputs,
+                            start_pos=start_pos,
+                            chunk_prefilling=args.chunk_prefilling,
+                            images=images,
+                        )
+                        start_pos += inputs.shape[1]
+                        torch.cuda.synchronize()
+                        t_ed = time.time()
+                        token = out[:, -1].max(1)[1].unsqueeze(1)
+                        time_lis.append(t_ed - t_st)
+                        if args.verbose:
+                            print(i, t_ed - t_st)
+                    print(f"Time To First Token: {np.mean(time_lis):.5f} s.")
+                    print("-" * 80)
+        else:
+            for i, (context_length, question_length) in enumerate(
+                zip(args.context_length, args.question_length)
+            ):
+                context_length = int("".join(context_length))
+                question_length = int("".join(question_length))
+                input_ids_old = [1 for _ in range(context_length)]
+                images = None
+                input_ids_new = [1 for _ in range(question_length)]
+                time_lis = []
+                print("-" * 80)
+                print(
+                    "History length: {} ; Question length: {}".format(
+                        context_length, question_length
+                    )
+                )
+                with torch.inference_mode():
+                    for i in range(10):  # Run ten times and get the average value
+                        # history rounds
+                        start_pos = 0
+                        if context_length > question_length:
+                            inputs = torch.as_tensor([input_ids_old], device=device)
+                            out = model(
+                                inputs,
+                                start_pos=start_pos,
+                                chunk_prefilling=args.chunk_prefilling,
+                                images=None,
+                            )
+                            start_pos += context_length
+
+                        # the present round
+                        torch.cuda.synchronize()
+                        t_st = time.time()
+                        inputs = torch.as_tensor([input_ids_new], device=device)
+                        out = model(
+                            inputs,
+                            start_pos=start_pos,
+                            chunk_prefilling=args.chunk_prefilling,
+                        )
+                        start_pos += inputs.shape[1]
+                        torch.cuda.synchronize()
+                        t_ed = time.time()
+
+                        token = out[:, -1].max(1)[1].unsqueeze(1)
+                        time_lis.append(t_ed - t_st)
+                        if args.verbose:
+                            print(i, t_ed - t_st)
+                    print(
+                        f"Time To First Token of this round: {np.mean(time_lis):.5f} s."
+                    )
+                    print("-" * 80)
+    else:
+        model = model_type_dict[args.model_type.lower()](config).half()
+        if args.precision in ["W4A16"]:
+            real_quantize_model_weight(
+                model,
+                w_bit=4,
+                q_config=dict(q_group_size=args.q_group_size, zero_point=True),
+                init_only=True,
+            )
+        model = model.to(device)
+
+        if args.precision in ["W4A16"]:
+            # tune_all_wqlinears(model)
+            make_quant_attn(model, device, args.flash_attn)
+            make_quant_norm(model)
+            make_fused_mlp(model)
+        device_warmup(device)
+
+        print("huggingface ckpt loaded")
+
+        # warming up
+        input_ids = [1 for _ in range(2048)]
+        inputs = torch.as_tensor([input_ids], device=device)
+        out = model(
+            inputs,
+            start_pos=0,
+            chunk_prefilling=args.chunk_prefilling,
+            quant=args.precision in ["W4A16"],
+        )  # warmup
+
+        if not args.chunk_prefilling:
+            for context_length in args.context_length:
+                context_length = int("".join(context_length))
+                input_ids = [1 for _ in range(context_length)]
+                time_lis = []
+                print("-" * 80)
+                print("Context length: {}".format(context_length))
+                with torch.inference_mode():
+                    for i in range(10):  # Run ten times and get the average value
+                        start_pos = 0
+                        torch.cuda.synchronize()
+                        t_st = time.time()
+                        inputs = torch.as_tensor([input_ids], device=device)
+                        out = model(
+                            inputs,
+                            start_pos=start_pos,
+                            chunk_prefilling=args.chunk_prefilling,
+                            quant=args.precision in ["W4A16"],
+                        )
+                        start_pos += inputs.shape[1]
+                        torch.cuda.synchronize()
+                        t_ed = time.time()
+                        token = torch.argmax(out, keepdim=True)[0]
+                        time_lis.append(t_ed - t_st)
+                        if args.verbose:
+                            print(i, t_ed - t_st)
+                    print(f"Time To First Token: {np.mean(time_lis):.5f} s.")
+                    # decoing throughput
+                    time_lis = []
+                    start_pos = context_length
+                    torch.cuda.synchronize()
+                    t_st = time.time()
+                    for i in range(token_num):
+                        token = model(
+                            token,
+                            start_pos=start_pos,
+                            chunk_prefilling=args.chunk_prefilling,
+                            quant=args.precision in ["W4A16"],
+                        )
+                        start_pos += 1
+                        token = torch.argmax(token, keepdim=True)[0]
+                        torch.cuda.synchronize()
+                    t_ed = time.time()
+                    time_lis.append(t_ed - t_st)
+                    print(
+                        f"Decoding throughput: {token_num/sum(time_lis):.5f} token/s."
+                    )
+                    print("-" * 80)
+        else:
+            for context_length, question_length in zip(
+                args.context_length, args.question_length
+            ):
+                context_length = int("".join(context_length))
+                question_length = int("".join(question_length))
+                input_ids_old = [1 for _ in range(context_length)]
+                input_ids_new = [1 for _ in range(question_length)]
+                time_lis = []
+                print("-" * 80)
+                print(
+                    "History length: {} ; Question length: {}".format(
+                        context_length, question_length
+                    )
+                )
+                with torch.inference_mode():
+                    for i in range(10):  # Run ten times and get the average value
+                        # history rounds
+                        start_pos = 0
+                        if context_length > question_length:
+                            inputs = torch.as_tensor([input_ids_old], device=device)
+                            out = model(
+                                inputs,
+                                start_pos=start_pos,
+                                chunk_prefilling=args.chunk_prefilling,
+                                quant=args.precision in ["W4A16"],
+                            )
+                            start_pos += inputs.shape[1]
+
+                        # the present round
+                        torch.cuda.synchronize()
+                        t_st = time.time()
+                        inputs = torch.as_tensor([input_ids_new], device=device)
+                        out = model(
+                            inputs,
+                            start_pos=start_pos,
+                            chunk_prefilling=args.chunk_prefilling,
+                            quant=args.precision in ["W4A16"],
+                        )
+                        start_pos += inputs.shape[1]
+                        torch.cuda.synchronize()
+                        t_ed = time.time()
+
+                        token = out[:, -1].max(1)[1].unsqueeze(1)
+                        time_lis.append(t_ed - t_st)
+                        if args.verbose:
+                            print(i, t_ed - t_st)
+                    print(
+                        f"Time To First Token of this round: {np.mean(time_lis):.5f} s."
+                    )
+                    print("-" * 80)
+
+
+if __name__ == "__main__":
+    main()

llm-awq/tinychat/demo.py

@@ -0,0 +1,283 @@
+import argparse
+import time
+import numpy as np
+import torch
+import torch.nn as nn
+from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig, modeling_utils
+from attributedict.collections import AttributeDict
+from tinychat.stream_generators import StreamGenerator
+import tinychat.utils.constants
+from tinychat.utils.load_quant import load_awq_model, load_awq_llama_fast
+from tinychat.utils.prompt_templates import get_prompter, get_stop_token_ids
+from tinychat.utils.tune import device_warmup, tune_all_wqlinears
+
+import os
+
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+
+# opt_params in TinyLLMEngine
+gen_params = AttributeDict(
+    [
+        ("seed", -1),  # RNG seed
+        ("n_threads", 1),  # TODO: fix this
+        ("n_predict", 512),  # new tokens to predict
+        ("n_parts", -1),  # amount of model parts (-1: determine from model dimensions)
+        ("n_ctx", 512),  # context size
+        ("n_batch", 512),  # batch size for prompt processing (must be >=32 to use BLAS)
+        ("n_keep", 0),  # number of tokens to keep from initial prompt
+        ("n_vocab", 50272),  # vocabulary size
+        # sampling parameters
+        ("logit_bias", dict()),  # logit bias for specific tokens: <int, float>
+        ("top_k", 40),  # <= 0 to use vocab size
+        ("top_p", 0.95),  # 1.0 = disabled
+        ("tfs_z", 1.00),  # 1.0 = disabled
+        ("typical_p", 1.00),  # 1.0 = disabled
+        ("temp", 0.70),  # 1.0 = disabled
+        ("repeat_penalty", 1.10),  # 1.0 = disabled
+        (
+            "repeat_last_n",
+            64,
+        ),  # last n tokens to penalize (0 = disable penalty, -1 = context size)
+        ("frequency_penalty", 0.00),  # 0.0 = disabled
+        ("presence_penalty", 0.00),  # 0.0 = disabled
+        ("mirostat", 0),  # 0 = disabled, 1 = mirostat, 2 = mirostat 2.0
+        ("mirostat_tau", 5.00),  # target entropy
+        ("mirostat_eta", 0.10),  # learning rate
+    ]
+)
+
+
+def stream_output(output_stream):
+    print(f"ASSISTANT: ", end="", flush=True)
+    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:
+            print(" ".join(output_text[pre:now]), end=" ", flush=True)
+            pre = now
+    print(" ".join(output_text[pre:]), flush=True)
+    if "timing" in outputs and outputs["timing"] is not None:
+        timing = outputs["timing"]
+        context_tokens = timing["context_tokens"]
+        context_time = timing["context_time"]
+        total_tokens = timing["total_tokens"]
+        generation_time_list = timing["generation_time_list"]
+        generation_tokens = len(generation_time_list)
+        average_speed = (context_time + np.sum(generation_time_list)) / (
+            context_tokens + generation_tokens
+        )
+        print("=" * 50)
+        print("Speed of Inference")
+        print("-" * 50)
+        print(f"TTFT : { context_time:.3f} s for {context_tokens} tokens")
+        print(
+            f"Speed of Generation : {np.average(generation_time_list)*1000:.2f} ms/token"
+        )
+        print("=" * 50)
+    return " ".join(output_text), total_tokens
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--model_type", type=str, default="LLaMa", help="type of the model"
+    )
+    parser.add_argument(
+        "--dtype", type=str, default="float16", choices=["float16", "bfloat16"]
+    )
+    parser.add_argument(
+        "--model_path",
+        type=str,
+        help="path to the model",
+    )
+    parser.add_argument(
+        "--precision", type=str, default="W4A16", help="compute precision"
+    )
+    parser.add_argument("--device", type=str, default="cuda:0")
+    parser.add_argument("--q_group_size", type=int, default=128)
+    parser.add_argument(
+        "--load_quant",
+        type=str,
+        help="path to the pre-quanted 4-bit weights",
+    )
+    parser.add_argument(
+        "--max_seq_len",
+        type=int,
+        default=2048,
+        help="maximum sequence length for kv cache",
+    )
+    parser.add_argument(
+        "--max_batch_size", type=int, default=1, help="maximum batch size for kv cache"
+    )
+    parser.add_argument(
+        "--mem_efficient_load",
+        action="store_true",
+        help="enable mem_efficient_load mod",
+    )
+    parser.add_argument(
+        "--single_round",
+        action="store_true",
+        help="whether to memorize previous conversations",
+    )
+    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",
+    )
+
+    args = parser.parse_args()
+    assert args.model_type.lower() in [
+        "llama",
+        "falcon",
+        "mpt",
+        "qwen",
+    ], "We only support llama & falcon & mpt now"
+    assert args.precision in ["W4A16", "W16A16"], "We only support W4A16/W16A16 now"
+
+    gen_params.n_predict = 1024
+    gen_params.n_vocab = 32000
+    tinychat.utils.constants.max_batch_size = args.max_batch_size
+    tinychat.utils.constants.max_seq_len = args.max_seq_len
+    tinychat.utils.constants.mem_efficient_load = args.mem_efficient_load
+    if tinychat.utils.constants.mem_efficient_load:
+        print("=" * 80)
+        print(
+            "[Info] You have activated mem_efficient_load mode.\n       Less on-chip memory will be consumed when loading the model.\n       However, the loading process will take more time."
+        )
+        print("=" * 80)
+    # TODO (Haotian): a more elegant implementation here.
+    # We need to update these global variables before models use them.
+    from tinychat.models import (
+        FalconForCausalLM,
+        LlamaForCausalLM,
+        MPTForCausalLM,
+        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
+
+    config = AutoConfig.from_pretrained(args.model_path, trust_remote_code=True)
+    if "mpt" in config.__class__.__name__.lower():
+        # config.init_device="meta"
+        tokenizer = AutoTokenizer.from_pretrained(
+            config.tokenizer_name, trust_remote_code=True
+        )
+    else:
+        tokenizer = AutoTokenizer.from_pretrained(
+            args.model_path, use_fast=False, trust_remote_code=True
+        )
+    torch_dtype = torch.float16 if args.dtype == "float16" else torch.bfloat16
+    modeling_utils._init_weights = False
+    torch.set_default_dtype(torch_dtype)
+
+    model_type_dict = {
+        "llama": LlamaForCausalLM,
+        "falcon": FalconForCausalLM,
+        "mpt": MPTForCausalLM,
+        "qwen": Qwen2ForCausalLM,
+    }
+
+    if args.precision == "W4A16":
+        if args.model_type.lower() == "llama":
+            model = model_type_dict["llama"](config).to(torch_dtype)
+            model = load_awq_llama_fast(
+                model, args.load_quant, 4, args.q_group_size, args.device
+            )
+        elif args.model_type.lower() == "qwen":
+            model = model_type_dict["qwen"](config).to(torch_dtype)
+            model = load_awq_llama_fast(
+                model, args.load_quant, 4, args.q_group_size, args.device
+            )
+        else:
+            model = model_type_dict[args.model_type.lower()](config).to(torch_dtype)
+            model = load_awq_model(
+                model, args.load_quant, 4, args.q_group_size, args.device
+            )
+    else:
+        loaded_model = AutoModelForCausalLM.from_pretrained(
+            args.model_path,
+            config=config,
+            torch_dtype=torch_dtype,
+            trust_remote_code=True,
+        )
+        model = (
+            model_type_dict[args.model_type.lower()](config)
+            .to(torch_dtype)
+            .to(args.device)
+        )
+        model.load_state_dict(loaded_model.state_dict())
+    # device warm up
+    device_warmup(args.device)
+
+    # autotune split_k_iters
+    # tune_all_wqlinears(model)
+
+    # TODO (Haotian): Verify if the StreamGenerator still works for the unmodified falcon impl.
+    stream_generator = StreamGenerator
+
+    # Optimize AWQ quantized model
+    if args.precision == "W4A16" and (
+        args.model_type.lower() == "llama" or args.model_type.lower() == "qwen"
+    ):
+        from tinychat.modules import make_quant_norm, make_quant_attn
+
+        if args.flash_attn:
+            make_quant_attn(model, args.device, args.flash_attn)
+        else:
+            make_quant_attn(model, args.device)
+        make_quant_norm(model)
+    model(
+        torch.randint(0, 1000, (1, 512), dtype=torch.int, device="cuda:0"),
+        0,
+        quant=args.precision == "W4A16",
+    )
+    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)
+    stop_token_ids = get_stop_token_ids(args.model_type, args.model_path)
+    count = 0
+    start_pos = 0
+    print("=" * 50)
+    while True:
+        # Get input from the user
+        input_prompt = input("USER: ")
+        if input_prompt == "":
+            print("EXIT...")
+            break
+        model_prompter.insert_prompt(input_prompt)
+        output_stream = stream_generator(
+            model,
+            tokenizer,
+            model_prompter.model_input,
+            start_pos,
+            gen_params,
+            device=args.device,
+            stop_token_ids=stop_token_ids,
+            chunk_prefilling=args.chunk_prefilling,
+            quant_llm=args.precision == "W4A16",
+        )
+        outputs, total_tokens = stream_output(output_stream)
+        if args.chunk_prefilling:
+            start_pos += total_tokens
+        else:
+            start_pos = 0
+        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

llm-awq/tinychat/internvl_benchmark.py

@@ -0,0 +1,167 @@
+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
+
+from torchao.quantization import quantize_, Int4WeightOnlyConfig
+
+import os
+
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+
+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 import InternVL3
+    
+    config = AutoConfig.from_pretrained(args.model_path, trust_remote_code=True)
+    config.resume_path = args.model_path
+    model = InternVL3(config).half()
+    model.language_model = model.language_model.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.language_model,
+            w_bit=4,
+            q_config=dict(q_group_size=128, zero_point=True),
+            init_only=True,
+        )
+        make_quant_attn(model.language_model, "cuda", True)
+        make_quant_norm(model.language_model)
+        make_fused_mlp(model.language_model)
+        model = model.to("cuda")
+    model = model.to(args.device)
+    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()
+
+    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/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/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/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/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_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/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_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/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/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/split_ckpt.py

@@ -0,0 +1,51 @@
+import os
+import re
+import torch
+import argparse
+
+
+def split(
+    ckpt_path: str,
+    out_folder_path: str,
+):
+    os.system(f"mkdir -p {out_folder_path}")
+    ckpt = torch.load(ckpt_path)
+    count = 0
+    for key, value in ckpt.items():
+        output_dict = {key: value}
+        output_name = out_folder_path + "/" + key + ".pt"
+        torch.save(output_dict, output_name)
+        count += 1
+    print(f"Finished splitting the original checkpoint into {count} shards.")
+
+
+def ckpt_folder_reader(ckpt_folder_path: str):
+    file_list = [f for f in os.listdir(ckpt_folder_path) if f.endswith(".pt")]
+    for ckpt in file_list:
+        print(ckpt)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--input_path",
+        type=str,
+        default=None,
+        help="Path to the original checkpoint (ends with *.pt)",
+    )
+    parser.add_argument(
+        "--output_path",
+        type=str,
+        default=None,
+        help="Folder to store the splitted checkpoint shards",
+    )
+
+    args = parser.parse_args()
+    assert (
+        args.input_path is not None
+    ), "Please specify the path to the original checkpoint."
+    if args.output_path is None:
+        suffix = r"\.pt$"
+        args.output_path = re.sub(suffix, "", args.input_path)
+
+    split(args.input_path, args.output_path)

llm-awq/tinychat/vila15_demo.py

@@ -0,0 +1,264 @@
+import argparse
+import torch
+
+from PIL import Image
+from tqdm import tqdm
+
+from transformers import AutoConfig, AutoTokenizer
+from accelerate import load_checkpoint_and_dispatch
+
+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 tinychat.utils.llava_image_processing import (
+    process_images,
+    load_images,
+    vis_images,
+)
+import tinychat.utils.constants
+
+# from tinychat.models.llava_llama import LlavaLlamaForCausalLM
+from tinychat.models.vila_llama import VilaLlamaForCausalLM
+from tinychat.stream_generators.llava_stream_gen import LlavaStreamGenerator
+from tinychat.utils.conversation_utils import gen_params, stream_output, TimeStats
+
+import os
+
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+
+
+def image_parser(args):
+    out = args.image_file.split(args.im_sep)
+    return out
+
+
+def skip(*args, **kwargs):
+    pass
+
+
+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
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        os.path.join(args.model_path, "llm"), use_fast=False
+    )
+    tinychat.utils.constants.LLAVA_DEFAULT_IMAGE_PATCH_TOKEN_IDX = (
+        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)
+    model = VilaLlamaForCausalLM(config).half()
+    tinychat.utils.constants.LLAVA_DEFAULT_IMAGE_PATCH_TOKEN_IDX = (
+        tokenizer.convert_tokens_to_ids(
+            [tinychat.utils.constants.LLAVA_DEFAULT_IMAGE_PATCH_TOKEN]
+        )[0]
+    )
+    vision_tower = model.get_vision_tower()
+    # if not vision_tower.is_loaded:
+    #     vision_tower.load_model()
+    image_processor = vision_tower.image_processor
+    # vision_tower = vision_tower.half()
+
+    if args.precision == "W16A16":
+        pbar = tqdm(range(1))
+        pbar.set_description("Loading checkpoint shards")
+        for i in pbar:
+            model.llm = load_checkpoint_and_dispatch(
+                model.llm,
+                os.path.join(args.model_path, "llm"),
+                no_split_module_classes=[
+                    "OPTDecoderLayer",
+                    "LlamaDecoderLayer",
+                    "BloomBlock",
+                    "MPTBlock",
+                    "DecoderLayer",
+                    "CLIPEncoderLayer",
+                ],
+            ).to(args.device)
+        model = model.to(args.device)
+
+    elif args.precision == "W4A16":
+        from tinychat.utils.load_quant import load_awq_model
+
+        model.llm = load_awq_model(model.llm, args.quant_path, 4, 128, args.device)
+        from tinychat.modules import (
+            make_quant_norm,
+            make_quant_attn,
+            make_fused_mlp,
+            make_fused_vision_attn,
+        )
+
+        if args.flash_attn:
+            print("Enabling flash-attention!")
+            make_quant_attn(model.llm, args.device, 1)
+        else:
+            print("Disabling flash-attention!")
+            make_quant_attn(model.llm, args.device)
+        make_quant_norm(model.llm)
+        # make_fused_mlp(model)
+        # make_fused_vision_attn(model,args.device)
+        model = model.to(args.device)
+
+    else:
+        raise NotImplementedError(f"Precision {args.precision} is not supported.")
+
+    image_files = image_parser(args)
+    image_num = len(image_files)
+    images = load_images(image_files)
+    if args.vis_image:
+        print("=" * 50)
+        print("Input Image:")
+        vis_images(image_files)
+    # Similar operation in model_worker.py
+    image_tensor = process_images(images, image_processor, model.config)
+    if type(image_tensor) is list:
+        image_tensor = [
+            image.to(args.device, dtype=torch.float16) for image in image_tensor
+        ]
+    else:
+        image_tensor = image_tensor.to(args.device, dtype=torch.float16)
+
+    device_warmup(args.device)
+    tune_llava_patch_embedding(vision_tower, device=args.device)
+
+    stream_generator = LlavaStreamGenerator
+
+    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: "
+
+    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 image here
+            image_token = get_image_token(model, args.model_path)
+            image_token_holder = (
+                tinychat.utils.constants.LLAVA_DEFAULT_IM_TOKEN_PLACE_HOLDER
+            )
+            im_token_count = input_prompt.count(image_token_holder)
+            if im_token_count == 0:
+                model_prompter.insert_prompt(image_token * image_num + input_prompt)
+            else:
+                assert im_token_count == image_num
+                input_prompt = input_prompt.replace(image_token_holder, image_token)
+                model_prompter.insert_prompt(input_prompt)
+        else:
+            model_prompter.insert_prompt(input_prompt)
+            if args.chunk_prefilling:
+                image_tensor = None  # Can insert more images in future
+        output_stream = stream_generator(
+            model,
+            tokenizer,
+            model_prompter.model_input,
+            start_pos,
+            gen_params,
+            device=args.device,
+            stop_token_ids=stop_token_ids,
+            image_tensor=image_tensor,
+            chunk_prefilling=args.chunk_prefilling,
+        )
+        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(
+        "--precision", type=str, default="W4A16", help="compute precision"
+    )
+    parser.add_argument(
+        "--image-file",
+        type=str,
+        default="https://llava.hliu.cc/file=/nobackup/haotian/code/LLaVA/llava/serve/examples/extreme_ironing.jpg",
+    )
+    parser.add_argument(
+        "--im-sep",
+        type=str,
+        default=",",
+    )
+    parser.add_argument("--device", type=str, default="cuda")
+    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",
+    )
+    args = parser.parse_args()
+    main(args)