.

Author: mgoin

auto_fp8/__init__.py

@@ -1,2 +1,7 @@
 from .modeling import AutoFP8ForCausalLM
 from .config import BaseQuantizeConfig
+
+__all__ = [
+    "AutoFP8ForCausalLM",
+    "BaseQuantizeConfig",
+]

auto_fp8/modeling.py

@@ -1,5 +1,5 @@
 import torch
-from transformers import AutoConfig, AutoModelForCausalLM, PreTrainedModel
+from transformers import AutoModelForCausalLM, PreTrainedModel
 from auto_fp8.quantize import (
     quantize_weights,
     quantize_activations,
@@ -14,8 +14,6 @@ def __init__(
         model: PreTrainedModel,
         quantize_config: BaseQuantizeConfig,
     ):
-        # super().__init__()
-
         self.model = model
         self.model_type = self.model.config.model_type
         self.quantize_config = quantize_config
@@ -30,11 +28,6 @@ def from_pretrained(
     ):
         """Load the un-quantized pretrained model"""
 
-        # if not torch.cuda.is_available():
-        #     raise EnvironmentError(
-        #         "Load pretrained model to do quantization requires CUDA available."
-        #     )
-
         def skip(*args, **kwargs):
             pass
 

auto_fp8/quantize.py

@@ -2,10 +2,9 @@
 import re
 from typing import Tuple
 import torch
-import torch.functional as F
 import transformers
 import tqdm
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoTokenizer
 
 
 # HACK: Override the dtype_byte_size function in transformers to support float8 types
@@ -59,8 +58,10 @@ def per_tensor_quantize(tensor: torch.Tensor) -> Tuple[torch.Tensor, float]:
 
 
 def fp8_gemm(A, A_scale, B, B_scale, bias, out_dtype):
-    cuda_compute_capability = torch.cuda.get_device_capability()
-    if cuda_compute_capability >= (9, 0):
+    native_fp8_support = (
+        torch.cuda.is_available() and torch.cuda.get_device_capability() >= (9, 0)
+    )
+    if native_fp8_support:
         output, _ = torch._scaled_mm(
             A,
             B.t(),

examples/quantize.py

@@ -4,7 +4,6 @@
 from typing import Tuple
 
 import torch
-import torch.functional as F
 import transformers
 import tqdm
 from datasets import load_dataset