Add Qwen test

Author: mgoin

auto_fp8/modeling.py

@@ -113,12 +113,15 @@ def skip(*args, **kwargs):
 
     def quantize(self, calibration_tokens: Optional[torch.Tensor] = None):
 <<<<<<< HEAD
+<<<<<<< HEAD
 =======
         def _prepare_calibration_data(calibration_tokens):
             if hasattr(calibration_tokens, "input_ids"):
                 return calibration_tokens.input_ids
             return calibration_tokens
 >>>>>>> 3ee9283 (Support calibrating kv cache scales)
+=======
+>>>>>>> 2739d61 (Add Qwen test)
 
         # Always quantize the weights as they do not require calibration data
         quantize_weights(self.model, self.quantize_config)
@@ -128,15 +131,21 @@ def _prepare_calibration_data(calibration_tokens):
                 calibration_tokens is not None
             ), "Calibration tokens required for activation quantization"
 <<<<<<< HEAD
+<<<<<<< HEAD
+=======
+>>>>>>> 2739d61 (Add Qwen test)
 
 
             def _prepare_calibration_data(calibration_tokens):
                 if hasattr(calibration_tokens, "input_ids"):
                     return calibration_tokens.input_ids
                 return calibration_tokens
 
+<<<<<<< HEAD
 =======
 >>>>>>> 3ee9283 (Support calibrating kv cache scales)
+=======
+>>>>>>> 2739d61 (Add Qwen test)
             quantize_activations(
                 self.model,
                 self.quantize_config,

tests/test_auto_fp8.py

@@ -1,16 +1,21 @@
 import os
 import shutil
 
+<<<<<<< HEAD
 <<<<<<< HEAD
 import pytest
 =======
 >>>>>>> 3ee9283 (Support calibrating kv cache scales)
+=======
+import pytest
+>>>>>>> 2739d61 (Add Qwen test)
 import safetensors.torch
 from transformers import AutoTokenizer
 
 from auto_fp8 import AutoFP8ForCausalLM, BaseQuantizeConfig
 
 MODELS = [
+<<<<<<< HEAD
     ("facebook/opt-125m", 160),
     ("Qwen/Qwen2-0.5B-Instruct", 620),
 ]
@@ -24,6 +29,15 @@ def test_dynamic_quantization():
     model_id = "facebook/opt-125m"
     quantized_model_dir = "opt-125m-fp8-dynamic"
 >>>>>>> 3ee9283 (Support calibrating kv cache scales)
+=======
+    "facebook/opt-125m",
+    "Qwen/Qwen2-0.5B-Instruct",
+]
+
+@pytest.mark.parametrize("model_id", MODELS)
+def test_dynamic_quantization(model_id):
+    quantized_model_dir = model_id.split("/")[-1] + "-fp8-dynamic"
+>>>>>>> 2739d61 (Add Qwen test)
 
     quantize_config = BaseQuantizeConfig(
         quant_method="fp8", activation_scheme="dynamic"
@@ -54,10 +68,16 @@ def test_static_quantization(model_id, target_size):
     assert model_size < target_size
 
 
+<<<<<<< HEAD
 def test_static_quantization():
     model_id = "facebook/opt-125m"
     quantized_model_dir = "opt-125m-fp8-static"
 >>>>>>> 3ee9283 (Support calibrating kv cache scales)
+=======
+@pytest.mark.parametrize("model_id", MODELS)
+def test_static_quantization(model_id):
+    quantized_model_dir = model_id.split("/")[-1] + "-fp8-static"
+>>>>>>> 2739d61 (Add Qwen test)
 
     tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)
     examples = ["auto-fp8 is an easy-to-use model quantization library"]
@@ -117,10 +137,9 @@ def test_kv_cache_static_quantization(model_id, target_size):
     target_size = target_size * (1024 * 1024)
     assert model_size < target_size
 
-
-def test_kv_cache_static_quantization():
-    model_id = "facebook/opt-125m"
-    quantized_model_dir = "opt-125m-fp8-static-kv"
+@pytest.mark.parametrize("model_id", MODELS)
+def test_kv_cache_static_quantization(model_id):
+    quantized_model_dir = model_id.split("/")[-1] + "-fp8-static-kv"
 
     tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)
     examples = ["auto-fp8 is an easy-to-use model quantization library"]