Update README.md

README.md

@@ -113,58 +113,16 @@ For more information please refer to the original model card [`meta-llama/Meta-L
 
 Llama 3.1 405B Instruct has been quantized using [AutoAWQ](https://github.com/casperhansen/AutoAWQ) from FP16 down to INT4 using the GEMM kernels performing zero-point quantization with a group size of 128.
 
-In order to quantize Llama 3.1 405B Instruct, we had to first install `torch` and `autoawq` as follows:
-
-```bash
-pip install "torch>=2.2.0,<2.3.0" autoawq --upgrade
-```
-
-Otherwise the quantization may fail, since the AutoAWQ kernels are built with PyTorch 2.2.1, meaning that those will break with PyTorch 2.3.0.
-
-Then we install the latest version of `transformers` as follows:
-
-```bash
-pip install "transformers>=4.43.0" --upgrade
-```
-
-And then we can run the following script, adapted from [`AutoAWQ/examples/quantize.py`](https://github.com/casper-hansen/AutoAWQ/blob/main/examples/quantize.py) as follows:
-
-```python
-from awq import AutoAWQForCausalLM
-from transformers import AutoTokenizer
-
-model_path = "meta-llama/Meta-Llama-3.1-405B-Instruct"
-quant_path = "hugging-quants/Meta-Llama-3.1-405B-Instruct-AWQ-INT4"
-quant_config = {
-    "zero_point": True,
-    "q_group_size": 128,
-    "w_bit": 4,
-    "version": "GEMM",
-}
-
-# Load model
-model = AutoAWQForCausalLM.from_pretrained(
-    model_path, **{"low_cpu_mem_usage": True, "use_cache": False}
-)
-tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
-
-# Quantize
-model.quantize(tokenizer, quant_config=quant_config)
-
-# Save quantized model
-model.save_quantized(quant_path)
-tokenizer.save_pretrained(quant_path)
-
-print(f'Model is quantized and saved at "{quant_path}"')
-```
-
 ## Quantized Model Usage
 
-In order to use the current quantized model, we offer support for different alternatives:
+> [!NOTE]
+> In order to run the inference with Llama 3.1 405B Instruct AWQ in INT4, around 203 GiB of VRAM are needed only for loading the model checkpoint, without including the KV cache or the CUDA graphs, meaning that there should be a bit over that VRAM available.
+
+In order to use the current quantized model, support is offered for different solutions:
 
 ### 🤗 transformers
 
-To run the inference on top of Llama 3.1 405B Instruct AWQ in INT4 precision, we can instantiate the AWQ model as any other causal language modeling model via `AutoModelForCausalLM` and run the inference normally.
+To run the inference on top of Llama 3.1 405B Instruct AWQ in INT4 precision, the AWQ model can be instantiated as any other causal language modeling model via `AutoModelForCausalLM` and run the inference normally.
 
 ```python
 import torch
@@ -235,4 +193,54 @@ The AutoAWQ script has been adapted from [AutoAWQ/examples/generate.py](https://
 
 ### 🤗 Text Generation Inference (TGI)
 
-Coming soon!
+Coming soon!
+
+## Quantization Reproduction
+
+> [!NOTE]
+> In order to quantize Llama 3.1 405B Instruct using AutoAWQ, you will need to use an instance with at least enough CPU RAM to fit the whole model i.e. ~800GiB, and an NVIDIA GPU with 80GiB of VRAM to quantize it.
+
+In order to quantize Llama 3.1 405B Instruct, first install `torch` and `autoawq` as follows:
+
+```bash
+pip install "torch>=2.2.0,<2.3.0" autoawq --upgrade
+```
+
+Otherwise the quantization may fail, since the AutoAWQ kernels are built with PyTorch 2.2.1, meaning that those will break with PyTorch 2.3.0.
+
+Then install the latest version of `transformers` as follows:
+
+```bash
+pip install "transformers>=4.43.0" --upgrade
+```
+
+And then, run the following script, adapted from [`AutoAWQ/examples/quantize.py`](https://github.com/casper-hansen/AutoAWQ/blob/main/examples/quantize.py) as follows:
+
+```python
+from awq import AutoAWQForCausalLM
+from transformers import AutoTokenizer
+
+model_path = "meta-llama/Meta-Llama-3.1-405B-Instruct"
+quant_path = "hugging-quants/Meta-Llama-3.1-405B-Instruct-AWQ-INT4"
+quant_config = {
+  "zero_point": True,
+  "q_group_size": 128,
+  "w_bit": 4,
+  "version": "GEMM",
+}
+
+# Load model
+model = AutoAWQForCausalLM.from_pretrained(
+  model_path, **{"low_cpu_mem_usage": True, "use_cache": False}
+)
+tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+
+# Quantize
+model.quantize(tokenizer, quant_config=quant_config)
+
+# Save quantized model
+model.save_quantized(quant_path)
+tokenizer.save_pretrained(quant_path)
+
+print(f'Model is quantized and saved at "{quant_path}"')
+```