fix format (#4815)

op_builder/gptq.py

@@ -1,19 +1,17 @@
-import os
-import torch
 import re
 
+import torch
+
 from .builder import Builder
-from .utils import append_nvcc_threads, get_cuda_cc_flag
+from .utils import append_nvcc_threads
 
-class GPTQBuilder(Builder):
 
+class GPTQBuilder(Builder):
     NAME = "cu_gptq"
     PREBUILT_IMPORT_PATH = "colossalai._C.cu_gptq"
 
     def __init__(self):
-        super().__init__(name=GPTQBuilder.NAME,
+        super().__init__(name=GPTQBuilder.NAME, prebuilt_import_path=GPTQBuilder.PREBUILT_IMPORT_PATH)
-        prebuilt_import_path=GPTQBuilder.PREBUILT_IMPORT_PATH)
-
 
     def include_dirs(self):
         ret = [self.csrc_abs_path("gptq"), self.get_cuda_home_include()]
@@ -21,32 +19,38 @@ class GPTQBuilder(Builder):
 
     def sources_files(self):
         ret = [
-            self.csrc_abs_path(fname) for fname in [
+            self.csrc_abs_path(fname)
-                'gptq/linear_gptq.cpp',
+            for fname in [
-                'gptq/column_remap.cu',
+                "gptq/linear_gptq.cpp",
-                'gptq/cuda_buffers.cu',
+                "gptq/column_remap.cu",
-                'gptq/q4_matmul.cu',
+                "gptq/cuda_buffers.cu",
-                'gptq/q4_matrix.cu'
+                "gptq/q4_matmul.cu",
+                "gptq/q4_matrix.cu",
             ]
         ]
         return ret
 
     def cxx_flags(self):
-        return ['-O3'] + self.version_dependent_macros
+        return ["-O3"] + self.version_dependent_macros
 
     def nvcc_flags(self):
-        extra_cuda_flags = ['-v',
+        extra_cuda_flags = [
-            '-std=c++14', '-U__CUDA_NO_HALF_OPERATORS__', '-U__CUDA_NO_HALF_CONVERSIONS__',
+            "-v",
-            '-U__CUDA_NO_HALF2_OPERATORS__', '-DTHRUST_IGNORE_CUB_VERSION_CHECK', "-lcublas", "-std=c++17"
+            "-std=c++14",
+            "-U__CUDA_NO_HALF_OPERATORS__",
+            "-U__CUDA_NO_HALF_CONVERSIONS__",
+            "-U__CUDA_NO_HALF2_OPERATORS__",
+            "-DTHRUST_IGNORE_CUB_VERSION_CHECK",
+            "-lcublas",
+            "-std=c++17",
         ]
 
-
         for arch in torch.cuda.get_arch_list():
-            res = re.search(r'sm_(\d+)', arch)
+            res = re.search(r"sm_(\d+)", arch)
             if res:
                 arch_cap = res[1]
                 if int(arch_cap) >= 80:
-                    extra_cuda_flags.extend(['-gencode', f'arch=compute_{arch_cap},code={arch}'])
+                    extra_cuda_flags.extend(["-gencode", f"arch=compute_{arch_cap},code={arch}"])
 
-        ret = ['-O3', '--use_fast_math'] + self.version_dependent_macros + extra_cuda_flags
+        ret = ["-O3", "--use_fast_math"] + self.version_dependent_macros + extra_cuda_flags
         return append_nvcc_threads(ret)