Switch backend to use llm-compressor

Author: mgoin

auto_fp8/__init__.py

@@ -1,5 +1,4 @@
-from .config import BaseQuantizeConfig
-from .modeling import AutoFP8ForCausalLM
+from .modeling import AutoFP8ForCausalLM, BaseQuantizeConfig
 
 __all__ = [
     "AutoFP8ForCausalLM",

auto_fp8/config.py

@@ -1,173 +1,80 @@
-import re
-from typing import List, Optional, Tuple
-
-import torch
-from transformers import AutoModelForCausalLM
-
-from auto_fp8.config import BaseQuantizeConfig
-from auto_fp8.quantize import (
-    quantize_activations,
-    quantize_weights,
-    save_quantized_model,
-)
-
-
-class AutoFP8ForCausalLM:
+import os
+from typing import List, Optional
+
+from transformers import AutoConfig, AutoTokenizer
+from datasets import Dataset
+from llmcompressor.transformers import SparseAutoModelForCausalLM
+from llmcompressor.transformers import oneshot
+from llmcompressor.modifiers.quantization import QuantizationModifier
+
+class BaseQuantizeConfig:
+    """Configuration for model quantization.
+
+    Args:
+        quant_method: Type/precision of quantization method to use.
+            At the moment, this is just "fp8" which specifically means
+            the fp8_e4m3 format in pytorch.
+        activation_scheme: Choice of either "dynamic" or "static" quantization
+            of activtions. If "static", then calibration samples are required
+            during quantization to produce accurate per-tensor scales for
+            activations of Linear modules.
+        ignore_patterns: List of patterns used to ignore layers. If a string
+            starts with "re:", then everything afterwards is used as python
+            regex style matching i.e. re.search(), for each Linear layer.
+            By default, "lm_head" is included to ignore the embedding
+            Linear layer usually at the end of decoder LLMs
+    """
     def __init__(
         self,
-        model: AutoModelForCausalLM,
-        quantize_config: BaseQuantizeConfig,
+        quant_method: str = "fp8",
+        activation_scheme: str = "static",
+        ignore_patterns: List[str] = ["lm_head"],
     ):
+        self.quant_method = quant_method
+        self.activation_scheme = activation_scheme
+        self.ignore_patterns = ignore_patterns
+
+
+class AutoFP8ForCausalLM:
+    def __init__(self, model: SparseAutoModelForCausalLM, quantize_config: BaseQuantizeConfig):
         self.model = model
         self.model_type = self.model.config.model_type
         self.config = self.model.config
-
-        # Gather the Linear module names that we want to ignore
-        quantize_config.ignored_layers = get_layers_to_ignore(
-            self.model, quantize_config.ignore_patterns
-        )
-
-        if quantize_config.kv_cache_quant_targets:
-            kv_cache_quant_layers = get_kv_cache_quant_layers(
-                self.model, quantize_config.kv_cache_quant_targets
-            )
-            if len(kv_cache_quant_layers) == 0:
-                raise ValueError(
-                    f"Could not find any kv cache layers using kv_cache_quant_targets={quantize_config.kv_cache_quant_targets}, please fix your argument."
-                )
-            quantize_config.kv_cache_quant_layers = kv_cache_quant_layers
-
         self.quantize_config = quantize_config
 
     @classmethod
-    def from_pretrained(
-        cls,
-        pretrained_model_name_or_path: str,
-        quantize_config: BaseQuantizeConfig,
-        **model_init_kwargs,
-    ):
-        """Load the un-quantized pretrained model"""
-
-        def skip(*args, **kwargs):
-            pass
-
-        torch.nn.init.kaiming_uniform_ = skip
-        torch.nn.init.uniform_ = skip
-        torch.nn.init.normal_ = skip
-
-        # Parameters related to loading from Hugging Face Hub
-        cache_dir = model_init_kwargs.pop("cache_dir", None)
-        force_download = model_init_kwargs.pop("force_download", False)
-        resume_download = model_init_kwargs.pop("resume_download", False)
-        proxies = model_init_kwargs.pop("proxies", None)
-        local_files_only = model_init_kwargs.pop("local_files_only", False)
-        use_auth_token = model_init_kwargs.pop("use_auth_token", None)
-        revision = model_init_kwargs.pop("revision", None)
-        subfolder = model_init_kwargs.pop("subfolder", "")
-        commit_hash = model_init_kwargs.pop("_commit_hash", None)
-
-        cached_file_kwargs = {
-            "cache_dir": cache_dir,
-            "force_download": force_download,
-            "proxies": proxies,
-            "resume_download": resume_download,
-            "local_files_only": local_files_only,
-            "use_auth_token": use_auth_token,
-            "revision": revision,
-            "subfolder": subfolder,
-            "_commit_hash": commit_hash,
-        }
-
-        torch.cuda.empty_cache()
-
-        # Important defaults
-        if "torch_dtype" not in model_init_kwargs:
-            model_init_kwargs["torch_dtype"] = "auto"
-
-        if "device_map" not in model_init_kwargs:
-            model_init_kwargs["device_map"] = "auto"
-
-        merged_kwargs = {**model_init_kwargs, **cached_file_kwargs}
-        print("Loading model with the following kwargs:", merged_kwargs)
-        model = AutoModelForCausalLM.from_pretrained(
-            pretrained_model_name_or_path, **merged_kwargs
+    def from_pretrained(cls, pretrained_model_name_or_path: str, quantize_config: BaseQuantizeConfig, **kwargs):
+        config = AutoConfig.from_pretrained(pretrained_model_name_or_path)
+        model = SparseAutoModelForCausalLM.from_pretrained(
+            pretrained_model_name_or_path,
+            config=config,
+            device_map="auto",
+            torch_dtype="auto",
+            **kwargs
         )
-
-        model_config = model.config.to_dict()
-        seq_len_keys = ["max_position_embeddings", "seq_length", "n_positions"]
-        if any(k in model_config for k in seq_len_keys):
-            for key in seq_len_keys:
-                if key in model_config:
-                    model.seqlen = model_config[key]
-                    break
-        else:
-            print("Can't get model's sequence length, setting to 2048.")
-            model.seqlen = 2048
-        model.eval()
-
         return cls(model, quantize_config)
 
-    def quantize(self, calibration_tokens: Optional[torch.Tensor] = None):
-
-        # Always quantize the weights as they do not require calibration data
-        quantize_weights(self.model, self.quantize_config)
-
-        if self.quantize_config.activation_scheme == "static":
-            assert (
-                calibration_tokens is not None
-            ), "Calibration tokens required for activation quantization"
-
-
-            def _prepare_calibration_data(calibration_tokens):
-                if hasattr(calibration_tokens, "input_ids"):
-                    return calibration_tokens.input_ids
-                return calibration_tokens
+    def quantize(self, dataset: Optional[Dataset] = None):
+        assert self.quantize_config.activation_scheme == "static"
+        assert dataset is not None, "Calibration tokens required for static activation quantization"
 
-            quantize_activations(
-                self.model,
-                self.quantize_config,
-                _prepare_calibration_data(calibration_tokens),
-            )
-
-    def save_quantized(self, save_dir):
-        save_quantized_model(
-            self.model,
-            quant_config=self.quantize_config,
-            save_dir=save_dir,
+        recipe = QuantizationModifier(
+            targets="Linear",
+            scheme="FP8",
+            ignore=self.quantize_config.ignore_patterns
         )
 
+        oneshot(
+            model=self.model,
+            dataset=dataset,
+            recipe=recipe,
+        )
 
-def get_layers_to_ignore(model, ignore_patterns) -> List[str]:
-    ignored_layers = set()
-
-    for name, linear in model.named_modules():
-        if not isinstance(linear, torch.nn.Linear):
-            continue
-
-        for ignore_pattern in ignore_patterns:
-            regex_prefix = "re:"
-            if ignore_pattern.startswith(regex_prefix):
-                # check if name matches regex and add to set if true
-                regex_pattern = ignore_pattern[len(regex_prefix) :]
-                if re.search(regex_pattern, name):
-                    ignored_layers.add(name)
-            else:
-                # else, exact match
-                if ignore_pattern == name:
-                    ignored_layers.add(name)
-
-    return list(ignored_layers)
-
-
-def get_kv_cache_quant_layers(model, kv_cache_quant_targets: Tuple[str]) -> List[str]:
-    kv_cache_quant_layers = []
-
-    for name, linear in model.named_modules():
-        if not isinstance(linear, torch.nn.Linear):
-            continue
-
-        for output_quant_target in kv_cache_quant_targets:
-            if name.endswith(output_quant_target):
-                kv_cache_quant_layers.append(name)
+    def save_quantized(self, save_directory: str):
+        self.save_pretrained(save_directory, save_compressed=True)
 
-    return kv_cache_quant_layers
+    def save_pretrained(self, save_directory: str, save_compressed: bool = True):   
+        self.model.save_pretrained(save_directory, save_compressed=save_compressed)
+        tokenizer = AutoTokenizer.from_pretrained(self.model.config._name_or_path)
+        tokenizer.save_pretrained(save_directory)
+        print(f"Saved final checkpoint to {os.path.abspath(save_directory)}")

auto_fp8/modeling.py

@@ -3,20 +3,27 @@
 
 from auto_fp8 import AutoFP8ForCausalLM, BaseQuantizeConfig
 
-pretrained_model_dir = "meta-llama/Meta-Llama-3-8B-Instruct"
-quantized_model_dir = "Meta-Llama-3-8B-Instruct-FP8"
+pretrained_model_dir = "facebook/opt-125m"
+quantized_model_dir = "opt-125m-FP8"
 
 tokenizer = AutoTokenizer.from_pretrained(pretrained_model_dir, use_fast=True)
 tokenizer.pad_token = tokenizer.eos_token
 
+MAX_SEQUENCE_LENGTH = 2048
 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")
+def preprocess(example):
+    example = tokenizer.apply_chat_template(example["messages"], tokenize=False)
+    return tokenizer(
+        example,
+        padding=False,
+        max_length=MAX_SEQUENCE_LENGTH,
+        truncation=True,
+        add_special_tokens=False,
+    )
+ds = ds.map(preprocess, remove_columns=ds.column_names)
 
 quantize_config = BaseQuantizeConfig(quant_method="fp8", activation_scheme="static")
 
-model = AutoFP8ForCausalLM.from_pretrained(
-    pretrained_model_dir, quantize_config=quantize_config
-)
-model.quantize(examples)
+model = AutoFP8ForCausalLM.from_pretrained(pretrained_model_dir, quantize_config)
+model.quantize(ds)
 model.save_quantized(quantized_model_dir)