fix the logic to verify whether AutoAWQ has been successfully install…

…ed (#2844)
InternLM · Dec 3, 2024 · 0dedd73 · 0dedd73
1 parent 8fbfed6
commit 0dedd73
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/lmdeploy/pytorch/backends/cuda/awq_modules.py b/lmdeploy/pytorch/backends/cuda/awq_modules.py
@@ -53,8 +53,6 @@ class AwqLinearW4A16Impl(LinearW4A16Impl):
 
     def __init__(self, in_features: int, out_features: int, w_bit: int,
                  group_size: int):
-        from awq.modules.linear.gemm import AWQ_INSTALLED
-        assert AWQ_INSTALLED
         self.in_features = in_features
         self.out_features = out_features
         self.w_bit = w_bit

diff --git a/lmdeploy/pytorch/backends/cuda/op_backend.py b/lmdeploy/pytorch/backends/cuda/op_backend.py
@@ -48,7 +48,11 @@ def get_layer_impl_builder(cls, layer_type: OpType):
             from .activation import TritonSiluAndMulBuilder
             return TritonSiluAndMulBuilder
         elif layer_type == OpType.LinearW4A16:
-            from awq.modules.linear.gemm import AWQ_INSTALLED
+            try:
+                from awq.modules.linear.gemm import awq_ext  # noqa: F401
+                AWQ_INSTALLED = True
+            except Exception:
+                AWQ_INSTALLED = False
             if AWQ_INSTALLED:
                 from .awq_modules import AwqLinearW4A16Builder
                 return AwqLinearW4A16Builder