Add check for quantization FP8

Author: gty111

gllm/layers/linear.py

@@ -7,6 +7,7 @@
 from gllm.dist_utils import (get_tp_size, get_tp_rank, divide, 
                              split_tensor_along_last_dim, tensor_model_parallel_all_reduce)
 from gllm.layers.quantization.fp8 import fp8LinearMethod
+from gllm.utils import get_device_capability
 
 class LinearBase(torch.nn.Module):
     """Base linear layer.
@@ -53,6 +54,8 @@ def create_weights(self,
                             requires_grad=False)
             self.register_parameter('weight', weight)
         elif self.quant_config['quant_method'] == 'fp8':
+            if get_device_capability() < 89:
+                raise Exception(f'FP8 quantizaiton method is not supported on device capability less than 89 (current is {get_device_capability()})')
             self.activation_scheme = self.quant_config['activation_scheme']
             self.block_quant = 'weight_block_size' in self.quant_config
             if self.block_quant:

gllm/utils.py

@@ -8,6 +8,7 @@
 import tempfile
 import logging
 import tqdm
+import torch
 
 from logger import logger
 from functools import partial
@@ -167,4 +168,9 @@ def get_dtype_bytes(dtype):
         info = torch.finfo(dtype)
     else:
         info = torch.iinfo(dtype)
-    return info.bits // 8  # bits => bytes
+    return info.bits // 8  # bits => bytes
+
+def get_device_capability():
+    device = torch.cuda.current_device()
+    capability_arr = torch.cuda.get_device_capability(device)
+    return capability_arr[0]*10 + capability_arr[1]