Last fixes

Author: mgoin

auto_fp8/modeling.py

@@ -75,7 +75,7 @@ def skip(*args, **kwargs):
             model_init_kwargs["device_map"] = "auto"
 
         merged_kwargs = {**model_init_kwargs, **cached_file_kwargs}
-        print(merged_kwargs)
+        print("Loading model with the following kwargs:", merged_kwargs)
         model = AutoModelForCausalLM.from_pretrained(
             pretrained_model_name_or_path, **merged_kwargs
         )
@@ -102,10 +102,10 @@ def _prepare_calibration_data(calibration_tokens):
                 return calibration_tokens.input_ids
             return calibration_tokens
 
-        if self.quantize_config.activation_scheme == "dynamic":
-            quantize_weights(self.model)
-        else:
-            quantize_weights(self.model)
+        # Always quantize the weights as they do not require calibration data
+        quantize_weights(self.model)
+
+        if self.quantize_config.activation_scheme == "static":
             quantize_activations(
                 self.model, _prepare_calibration_data(calibration_tokens)
             )

example.py

@@ -1,14 +1,16 @@
 from transformers import AutoTokenizer
 from auto_fp8 import AutoFP8ForCausalLM, BaseQuantizeConfig
 
-pretrained_model_dir = "facebook/opt-125m"
-quantized_model_dir = "opt-125m-fp8"
+pretrained_model_dir = "meta-llama/Meta-Llama-3-8B-Instruct"
+quantized_model_dir = "Meta-Llama-3-8B-Instruct-FP8"
 
 tokenizer = AutoTokenizer.from_pretrained(pretrained_model_dir, use_fast=True)
-examples = ["auto-fp8 is an easy-to-use model quantization library"]
+examples = ["auto_fp8 is an easy-to-use model quantization library"]
 examples = tokenizer(examples, return_tensors="pt").to("cuda")
 
-quantize_config = BaseQuantizeConfig(quant_method="fp8", activation_scheme="static")
+quantize_config = BaseQuantizeConfig(
+    quant_method="fp8", activation_scheme="dynamic"
+)  # or "static"
 
 model = AutoFP8ForCausalLM.from_pretrained(
     pretrained_model_dir, quantize_config=quantize_config