From 555334dec6d708da1f9601cd0ee0084585beb768 Mon Sep 17 00:00:00 2001
From: mgoin <michael@neuralmagic.com>
Date: Mon, 10 Jun 2024 21:36:54 +0000
Subject: [PATCH] Fix inline code and add checkpoint format

---
 docs/source/quantization/fp8.rst | 83 ++++++++++++++++++++++----------
 1 file changed, 58 insertions(+), 25 deletions(-)

diff --git a/docs/source/quantization/fp8.rst b/docs/source/quantization/fp8.rst
index a7aaae8a2805..0c88d8d71509 100644
--- a/docs/source/quantization/fp8.rst
+++ b/docs/source/quantization/fp8.rst
@@ -9,21 +9,22 @@ Please visit the HF collection of `quantized FP8 checkpoints of popular LLMs rea
 
 The FP8 types typically supported in hardware have two distinct representations, each useful in different scenarios:
 
-- **E4M3**: Consists of 1 sign bit, 4 exponent bits, and 3 bits of mantissa. It can store values up to +/-448 and `nan`.
-- **E5M2**: Consists of 1 sign bit, 5 exponent bits, and 2 bits of mantissa. It can store values up to +/-57344, +/- `inf`, and `nan`. The tradeoff for the increased dynamic range is lower precision of the stored values.
+- **E4M3**: Consists of 1 sign bit, 4 exponent bits, and 3 bits of mantissa. It can store values up to +/-448 and ``nan``.
+- **E5M2**: Consists of 1 sign bit, 5 exponent bits, and 2 bits of mantissa. It can store values up to +/-57344, +/- ``inf``, and ``nan``. The tradeoff for the increased dynamic range is lower precision of the stored values.
 
 Quick Start with Online Dynamic Quantization
 -------------------------------------
 
-Dynamic quantization of an original precision BF16/FP16 model to FP8 can be achieved with vLLM without any calibration data required. You can enable the feature by specifying `--quantization="fp8"` in the command line or setting `quantization="fp8"` in the LLM constructor.
+Dynamic quantization of an original precision BF16/FP16 model to FP8 can be achieved with vLLM without any calibration data required. You can enable the feature by specifying ``--quantization="fp8"`` in the command line or setting ``quantization="fp8"`` in the LLM constructor.
 
-In this mode, all Linear modules (except for the final `lm_head`) have their weights quantized down to FP8_E4M3 precision with a per-tensor scale. Activations have their minimum and maximum values calculated during each forward pass to provide a dynamic per-tensor scale for high accuracy. As a result, latency improvements are limited in this mode.
+In this mode, all Linear modules (except for the final ``lm_head``) have their weights quantized down to FP8_E4M3 precision with a per-tensor scale. Activations have their minimum and maximum values calculated during each forward pass to provide a dynamic per-tensor scale for high accuracy. As a result, latency improvements are limited in this mode.
 
 .. code-block:: python
 
     from vllm import LLM
     model = LLM("facebook/opt-125m", quantization="fp8")
     # INFO 06-10 17:55:42 model_runner.py:157] Loading model weights took 0.1550 GB
+    result = model.generate("Hello, my name is")
 
 .. warning::
 
@@ -39,12 +40,12 @@ For offline quantization to FP8, please install the `AutoFP8 library <https://gi
     git clone https://github.com/neuralmagic/AutoFP8.git
     pip install -e AutoFP8
 
-This package introduces the `AutoFP8ForCausalLM` and `BaseQuantizeConfig` objects for managing how your model will be compressed.
+This package introduces the ``AutoFP8ForCausalLM`` and ``BaseQuantizeConfig`` objects for managing how your model will be compressed.
 
 Offline Quantization with Dynamic Activation Scaling Factors
 ------------------------------------------------------------
 
-You can use AutoFP8 to produce checkpoints with their weights quantized to FP8 ahead of time and let vLLM handle calculating dynamic scales for the activations at runtime for maximum accuracy. You can enable this with the `activation_scheme="dynamic"` argument.
+You can use AutoFP8 to produce checkpoints with their weights quantized to FP8 ahead of time and let vLLM handle calculating dynamic scales for the activations at runtime for maximum accuracy. You can enable this with the ``activation_scheme="dynamic"`` argument.
 
 .. warning::
 
@@ -57,15 +58,18 @@ You can use AutoFP8 to produce checkpoints with their weights quantized to FP8 a
     pretrained_model_dir = "meta-llama/Meta-Llama-3-8B-Instruct"
     quantized_model_dir = "Meta-Llama-3-8B-Instruct-FP8-Dynamic"
 
+    # Define quantization config with static activation scales
+    quantize_config = BaseQuantizeConfig(quant_method="fp8", activation_scheme="dynamic")
     # For dynamic activation scales, there is no need for calbration examples
     examples = []
-    quantize_config = BaseQuantizeConfig(quant_method="fp8", activation_scheme="dynamic")
 
+    # Load the model, quantize, and save checkpoint
     model = AutoFP8ForCausalLM.from_pretrained(pretrained_model_dir, quantize_config)
     model.quantize(examples)
     model.save_quantized(quantized_model_dir)
 
-In the output of the above script, you should be able to see the quantized Linear modules replaced in the model definition. Note that the `lm_head` Linear module at the end is currently skipped by default.
+In the output of the above script, you should be able to see the quantized Linear modules (FP8DynamicLinear) replaced in the model definition. 
+Note that the ``lm_head`` Linear module at the end is currently skipped by default.
 
 .. code-block:: text
 
@@ -97,7 +101,8 @@ In the output of the above script, you should be able to see the quantized Linea
     )
     Saving the model to Meta-Llama-3-8B-Instruct-FP8-Dynamic
 
-Your model checkpoint with quantized weights should be available at `quantized_model_dir`. We can see that the weights are smaller than the original BF16 precision.
+Your model checkpoint with quantized weights should be available at ``Meta-Llama-3-8B-Instruct-FP8/``.
+We can see that the weights are smaller than the original BF16 precision.
 
 .. code-block:: bash
 
@@ -118,11 +123,13 @@ Finally, you can load the quantized model checkpoint directly in vLLM.
 
     from vllm import LLM
     model = LLM(model="Meta-Llama-3-8B-Instruct-FP8-Dynamic/")
+    # INFO 06-10 21:15:41 model_runner.py:159] Loading model weights took 8.4596 GB
+    result = model.generate("Hello, my name is")
 
 Offline Quantization with Static Activation Scaling Factors
 -----------------------------------------------------------
 
-For the best inference performance, you can use AutoFP8 with calibration data to produce per-tensor static scales for both the weights and activations by enabling the `activation_scheme="static"` argument.
+For the best inference performance, you can use AutoFP8 with calibration data to produce per-tensor static scales for both the weights and activations by enabling the ``activation_scheme="static"`` argument.
 
 .. code-block:: python
 
@@ -136,34 +143,60 @@ For the best inference performance, you can use AutoFP8 with calibration data to
     tokenizer = AutoTokenizer.from_pretrained(pretrained_model_dir, use_fast=True)
     tokenizer.pad_token = tokenizer.eos_token
 
+    # Load and tokenize 512 dataset samples for calibration of activation scales
     ds = load_dataset("mgoin/ultrachat_2k", split="train_sft").select(range(512))
     examples = [tokenizer.apply_chat_template(batch["messages"], tokenize=False) for batch in ds]
     examples = tokenizer(examples, padding=True, truncation=True, return_tensors="pt").to("cuda")
 
+    # Define quantization config with static activation scales
     quantize_config = BaseQuantizeConfig(quant_method="fp8", activation_scheme="static")
 
+    # Load the model, quantize, and save checkpoint
     model = AutoFP8ForCausalLM.from_pretrained(pretrained_model_dir, quantize_config)
     model.quantize(examples)
     model.save_quantized(quantized_model_dir)
 
-Your model checkpoint with quantized weights should be available at `Meta-Llama-3-8B-Instruct-FP8/`. We can see that the weights are smaller than the original BF16 precision.
-
-.. code-block:: bash
-
-    ls -lh Meta-Llama-3-8B-Instruct-FP8/
-    total 8.5G
-    -rw-rw-r-- 1 user user  869 Jun  7 14:43 config.json
-    -rw-rw-r-- 1 user user  194 Jun  7 14:43 generation_config.json
-    -rw-rw-r-- 1 user user 4.7G Jun  7 14:43 model-00001-of-00002.safetensors
-    -rw-rw-r-- 1 user user 3.9G Jun  7 14:43 model-00002-of-00002.safetensors
-    -rw-rw-r-- 1 user user  43K Jun  7 14:43 model.safetensors.index.json
-    -rw-rw-r-- 1 user user  296 Jun  7 14:43 special_tokens_map.json
-    -rw-rw-r-- 1 user user  50K Jun  7 14:43 tokenizer_config.json
-    -rw-rw-r-- 1 user user 8.7M Jun  7 14:43 tokenizer.json
-
+Your model checkpoint with quantized weights and activations should be available at ``Meta-Llama-3-8B-Instruct-FP8/``.
 Finally, you can load the quantized model checkpoint directly in vLLM.
 
 .. code-block:: python
 
     from vllm import LLM
     model = LLM(model="Meta-Llama-3-8B-Instruct-FP8/")
+    # INFO 06-10 21:15:41 model_runner.py:159] Loading model weights took 8.4596 GB
+    result = model.generate("Hello, my name is")
+
+FP8 checkpoint structure explanation
+-----------------------------------------------------------
+
+Here we detail the structure for the FP8 checkpoints.
+
+The following is necessary to be present in the model's ``config.json``:
+
+.. code-block:: yaml
+    "quantization_config": {
+        "quant_method": "fp8",
+        "activation_scheme": "static" or "dynamic"
+      },
+
+
+Each quantized layer in the state_dict will have these tensors:
+
+* If the config has `"activation_scheme": "static"`:
+
+.. code-block:: text
+    model.layers.0.mlp.down_proj.weight              < F8_E4M3
+    model.layers.0.mlp.down_proj.input_scale         < F32
+    model.layers.0.mlp.down_proj.weight_scale        < F32
+
+* If the config has `"activation_scheme": "dynamic"`:
+
+.. code-block:: text
+    model.layers.0.mlp.down_proj.weight              < F8_E4M3
+    model.layers.0.mlp.down_proj.weight_scale        < F32
+
+
+Additionally, there can be `FP8 kv-cache scaling factors <https://github.com/vllm-project/vllm/pull/4893>`_ contained within quantized checkpoints specified through the ``.kv_scale`` parameter present on the Attention Module, such as:
+
+.. code-block:: text
+    model.layers.0.self_attn.kv_scale	             < F32
\ No newline at end of file