quantization

Author: lucylq

README.md

@@ -103,9 +103,9 @@ Designed for interactive graphical conversations using the familiar web browser
 
 Quantization is the process of converting a model into a more memory-efficient representation.  Quantization is particularly important for accelerators -- to take advantage of the available memory bandwidth, and fit in the often limited high-speed memory in accelerators – and mobile devices – to fit in the typically very limited memory of mobile devices.
 
-Depending on the model and the target device, different quantization recipes may be applied.  Torchchat contains two example configurations to optimize performance for GPU-based systems `config/data/cuda.json` , and mobile systems `config/data/mobile.json`.  The GPU configuration is targeted towards optimizing for memory bandwidth which is a scarce resource in powerful GPUs (and to a less degree, memory footprint to fit large models into a device's memory).  The mobile configuration is targeted towards optimizing for memory fotoprint because in many devices, a single application is limited to as little as GB or less of memory.
+Depending on the model and the target device, different quantization recipes may be applied.  Torchchat contains two example configurations to optimize performance for GPU-based systems `config/data/server.json` , and mobile systems `config/data/mobile.json`.  The GPU configuration is targeted towards optimizing for memory bandwidth which is a scarce resource in powerful GPUs (and to a less degree, memory footprint to fit large models into a device's memory).  The mobile configuration is targeted towards optimizing for memory fotoprint because in many devices, a single application is limited to as little as GB or less of memory.
 
-You can use the quantization recipes in conjunction with any of the `chat`, `generate` and `browser` commands to test their impact and accelerate model execution. You will apply these recipes to the export comamnds below, to optimize the exported models.  To adapt these recipes or wrote your own, please refer to the [quantization overview](docs/quantization.md).
+You can use the quantization recipes in conjunction with any of the `chat`, `generate` and `browser` commands to test their impact and accelerate model execution. You will apply these recipes to the export commands below, to optimize the exported models.  To adapt these recipes or wrote your own, please refer to the [quantization overview](docs/quantization.md).
 
 ---
 *TO BE REPLACED BY SUITABLE ORDING PROVIDED BY LEGAL*
@@ -215,31 +215,32 @@ python3 torchchat.py chat codellama
 AOT compiles models into machine code before execution, enhancing performance and predictability. It's particularly beneficial for frequently used models or those requiring quick start times. However, it may lead to larger binary sizes and lacks the runtime flexibility of eager mode.
 
 **Examples**
-The following example uses the Stories15M model.
+The following example uses the Llama3 8B model.
 ```
 # Compile
-python3 torchchat.py export stories15M --output-dso-path stories15M.so
+python3 torchchat.py export llama3 --output-dso-path llama3.so
 
 # Execute
-python3 torchchat.py generate --dso-path stories15M.so --prompt "Hello my name is"
+python3 torchchat.py generate llama3 --quantize config/data/server.json--dso-path llama3.so --prompt "Hello my name is"
 ```
 
-NOTE: The exported model will be large. We suggest you quantize the model, explained further down, before deploying the model on device.
+NOTE: We use `--quantize config/data/server.json` to quantize the llama3 model to reduce model size and improve performance for on-device use cases.
 
 ### ExecuTorch
 
 ExecuTorch enables you to optimize your model for execution on a mobile or embedded device, but can also be used on desktop for testing.
 Before running ExecuTorch commands, you must first set-up ExecuTorch in torchchat, see [Set-up Executorch](docs/executorch_setup.md).
 
 **Examples**
-The following example uses the Stories15M model.
+The following example uses the Llama3 8B model.
 ```
 # Compile
-python3 torchchat.py export stories15M --output-pte-path stories15M.pte
+python3 torchchat.py export llama3 --quantize config/data/mobile.json --output-pte-path llama3.pte
 
 # Execute
-python3 torchchat.py generate --device cpu --pte-path stories15M.pte --prompt "Hello my name is"
+python3 torchchat.py generate llama3 --device cpu --pte-path llama3.pte --prompt "Hello my name is"
 ```
+NOTE: We use `--quantize config/data/mobile.json` to quantize the llama3 model to reduce model size and improve performance for on-device use cases.
 
 See below under [Mobile Execution](#run-mobile) if you want to deploy and execute a model in your iOS or Android app.
 

config/data/cuda.json renamed to config/data/server.json

@@ -4,15 +4,58 @@
 ## Introduction
 Quantization focuses on reducing the precision of model parameters and computations from floating-point to lower-bit integers, such as 8-bit integers. This approach aims to minimize memory requirements, accelerate inference speeds, and decrease power consumption, making models more feasible for deployment on edge devices with limited computational resources. While quantization can potentially degrade the model's performance, the methods supported by torchchat are designed to mitigate this effect, maintaining a balance between efficiency and accuracy.
 
-## Supported quantization techniques
+## Weight Quantization
+| compression | FP Precision | weight quantization (bitwidth)| weight quantization (group size) | dynamic activation quantization | Eager | AOTI | ExecuTorch |
+|--|--|--|--|--|--|--|--|
+| linear (asymmetric) | fp32, fp16, bf16 | [8, 4]* | [32, 64, 128, 256]** | | ✅ | ✅ | ✅ |
+| linear (symmetric) | | | | a8w4dq | ✅ | ✅ | ✅ |
+| linear (asymmetric) with GPTQ | | | | | ✅ | ✅ | ❌ |
+| linear (asymmetric) with HQQ | | | | | ✅ | ✅ | ❌ |
+
+## Embedding Quantization
+Due to the larger vocabulary size of llama3, we also recommend quantizing the embeddings to further reduce the model size for on-device usecases.
+
+| compression | FP Precision | weight quantization (bitwidth)| weight quantization (group size) | dynamic activation quantization | Eager | AOTI | ExecuTorch |
+|--|--|--|--|--|--|--|--|
+| embedding table (symmetric) | fp32, fp16, bf16 | [8, 4]* | [32, 64, 128, 256]** | | ✅ | ✅ | ✅ |
+
+*These are the only valid bitwidth options.
+
+**There are many valid group size options, including 512, 1024, etc.
+
+## Examples
+We can mix and match weight quantization with embedding quantization. Quantization options are passed in json format either as a config file (see [server.json](../config/data/server.json) and [mobile.json](../config/data/mobile.json)) or a json string.
+
+* Config file
+  ```
+  --quantize quant_config.json
+  ```
+* Single quantization scheme
+  ```
+  --quantize '{"linear:a8w4dq": {"groupsize" : 256}}'
+  ```
+* Multiple quantization schemes
+  ```
+  --quantize '{"embedding": {"bitwidth": 4, "groupsize":32}, "linear:a8w4dq": {"groupsize" : 256}}'
+  ```
+Quantization recipes can be applied in conjunction with any of the `chat`, `generate`, `browser` and `export` commands. Below are examples showcasing eager mode with `generate` and AOTI and ExecuTorch with `export`.
+### Eager mode
+```
+python3 generate.py [--compile] --checkpoint-path ${MODEL_PATH} --prompt "Hello, my name is" --quantize '{"embedding" : {"bitwidth": 8, "groupsize": 0}}' --device cpu
+```
+### AOTI
+```
+python3 torchchat.py export llama3 --quantize '{"embedding": {"bitwidth": 4, "groupsize":32}, "linear:a8w4dq": {"groupsize" : 256}}' --output-dso-path llama3.dso
+
+python3 generate.py --dso-path llama3.dso  --prompt "Hello my name is"
+```
+### ExecuTorch
+```
+python3 torchchat.py export llama3 --quantize '{"embedding": {"bitwidth": 4, "groupsize":32}, "linear:a8w4dq": {"groupsize" : 256}}' --output-pte-path llama3.pte
+
+python3 generate.py --pte-path llama3.pte  --prompt "Hello my name is"
+```
 
-| compression | FP precision |  weight quantization | dynamic activation quantization |
-|--|--|--|--|
-embedding table (symmetric) | fp32, fp16, bf16 | 8b (group/channel), 4b (group/channel) | n/a |
-linear operator (symmetric) | fp32, fp16, bf16 | 8b (group/channel) | n/a |
-linear operator (asymmetric) | n/a | 4b (group), a6w4dq | a8w4dq (group) |
-linear operator (asymmetric) with GPTQ | n/a | 4b (group) | n/a |
-linear operator (asymmetric) with HQQ | n/a |  work in progress | n/a |
 
 ## Model precision (dtype precision setting)
 You can generate models (for both export and generate, with eager, torch.compile, AOTI, ET, for all backends - mobile at present will primarily support fp32, with all options) specify the precision of the model with
@@ -24,7 +67,7 @@ python3 export.py --dtype [bf16 | fp16 | fp32] ...
 ```
 
 Unlike gpt-fast which uses bfloat16 as default, torchchat uses float32 as the default. As a consequence you will have to set to --dtype bf16 or --dtype fp16 on server / desktop for best performance.
-Support for FP16 and BF16 is limited in many embedded processors.  Additional executorch support for 16-bit floating point types may be added in the future based on hardware support.
+Support for FP16 and BF16 is limited in many embedded processors.  Additional ExecuTorch support for 16-bit floating point types may be added in the future based on hardware support.
 
 ## Making your models fit and execute fast!
 
@@ -214,12 +257,12 @@ python3 generate.py [--compile] --checkpoint-path ${MODEL_PATH} --prompt "Hello,
 ```
 
 ```
-python3 export.py --checkpoint-path ${MODEL_PATH} -d fp32 --quantize '{"linear:int4": {"groupsize" : 32} }' [ --output-pte-path ${MODEL_OUT}/${MODEL_NAME}_int4-gw32.pte | --output-dso-path ${MODEL_OUT}/${MODEL_NAME}_int4-gw32.dso]
+python3 export.py --checkpoint-path ${MODEL_PATH} -d fp32 --quantize '{"linear:int4": {"groupsize" : 32} }' --output-dso-path ${MODEL_OUT}/${MODEL_NAME}_int4-gw32.dso
 ```
 Now you can run your model with the same command as before:
 
 ```
-python3 generate.py [ --pte-path ${MODEL_OUT}/${MODEL_NAME}_int4-gw32.pte | --dso-path ${MODEL_OUT}/${MODEL_NAME}_int4-gw32.dso]  --prompt "Hello my name is"
+python3 generate.py --dso-path ${MODEL_OUT}/${MODEL_NAME}_int4-gw32.dso  --prompt "Hello my name is"
 ```
 
 ## 4-Bit Integer Linear Quantization  (a8w4dq)
@@ -247,12 +290,12 @@ python3 generate.py [--compile] --checkpoint-path ${MODEL_PATH} --prompt "Hello,
 ```
 
 ```
-python3 export.py --checkpoint-path ${MODEL_PATH} -d fp32 --quantize '{"linear:gptq": {"groupsize" : 32} }' [ --output-pte-path ${MODEL_OUT}/${MODEL_NAME}_gptq.pte | ...dso... ]
+python3 export.py --checkpoint-path ${MODEL_PATH} -d fp32 --quantize '{"linear:gptq": {"groupsize" : 32} }' --output-dso-path ${MODEL_OUT}/${MODEL_NAME}_gptq.dso
 ```
 Now you can run your model with the same command as before:
 
 ```
-python3 generate.py [ --pte-path ${MODEL_OUT}/${MODEL_NAME}_gptq.pte | ...dso...]  --prompt "Hello my name is"
+python3 generate.py --dso-path ${MODEL_OUT}/${MODEL_NAME}_gptq.dso  --prompt "Hello my name is"
 ```
 
 ## 4-bit Integer Linear Quantization with HQQ (hqq)
@@ -267,12 +310,12 @@ python3 generate.py [--compile] --checkpoint-path ${MODEL_PATH} --prompt "Hello,
 ```
 
 ```
-python3 export.py --checkpoint-path ${MODEL_PATH} -d fp32 --quantize '{"linear:hqq": {"groupsize" : 32} }' [ --output-pte-path ${MODEL_OUT}/${MODEL_NAME}_hqq.pte | ...dso... ]
+python3 export.py --checkpoint-path ${MODEL_PATH} -d fp32 --quantize '{"linear:hqq": {"groupsize" : 32} }' --output-dso-path ${MODEL_OUT}/${MODEL_NAME}_hqq.dso
 ```
 Now you can run your model with the same command as before:
 
 ```
-python3 generate.py [ --pte-path ${MODEL_OUT}/${MODEL_NAME}_hqq.pte | ...dso...]  --prompt "Hello my name is"
+python3 generate.py --dso-path ${MODEL_OUT}/${MODEL_NAME}_hqq.dso --prompt "Hello my name is"
 
 
 ## Adding additional quantization schemes

docs/quantization.md

@@ -94,7 +94,6 @@ def download_and_convert(
         # overwriting if necessary.
         if os.path.isdir(model_dir):
             shutil.rmtree(model_dir)
-        os.makedirs(model_dir, exist_ok=True)
         shutil.move(temp_dir, model_dir)
 
     finally: