refactor linear

Signed-off-by: realliujiaxu <realliujiaxu@163.com>

Author: realliujiaxu

vllm_ascend/ops/linear.py

@@ -19,7 +19,6 @@
 
 import torch
 import torch.distributed as dist
-import torch.nn as nn
 import torch_npu
 from torch.distributed import ProcessGroup
 from torch.nn.parameter import Parameter
@@ -63,14 +62,11 @@ def __init__(
         *,
         return_bias: bool = True,
     ):
-        self.comm_group = None
-        if prefix.find("gate_up_proj") != -1 and mlp_tp_enable():
-            self.comm_group = get_mlp_tp_group()
-        else:
+        # if self has attr `tp_size`, this means it has been customized by subclass
+        if not hasattr(self, "tp_size"):
             self.comm_group = get_tp_group()
-
-        self.tp_size = self.comm_group.world_size
-        self.tp_rank = self.comm_group.rank_in_group
+            self.tp_size = self.comm_group.world_size
+            self.tp_rank = self.comm_group.rank_in_group
 
         self.input_size_per_partition = input_size
         self.output_size_per_partition = divide(output_size, self.tp_size)
@@ -81,6 +77,8 @@ def __init__(
                 divide(output_size, self.tp_size)
                 for output_size in self.output_sizes
             ]
+        # skip ColumnParallelLinear.__init__, as it will create weight_loader with default tp group
+        # we will create weight_loader by customized comm group
         AscendLinearBase.__init__(self,
                                   input_size,
                                   output_size,
@@ -164,6 +162,7 @@ def __init__(
         self.output_size_per_partition = output_size
         self.output_partition_sizes = [output_size]
 
+        # skip RowParallelLinear.__init__, as it will create weight_loader with default tp group
         AscendLinearBase.__init__(self,
                                   input_size,
                                   output_size,
@@ -526,7 +525,7 @@ def __init__(
         self.output_sizes = [
             self.num_heads * self.head_size * tp_size,  # q_proj
             self.num_kv_heads * self.head_size * tp_size,  # k_proj
-            self.num_kv_heads * self.head_size * tp_size,  # v_proj 
+            self.num_kv_heads * self.head_size * tp_size,  # v_proj
         ]
         AscendColumnParallelLinear.__init__(self,
                                             input_size=input_size,
@@ -593,22 +592,15 @@ def __init__(
         return_bias: bool = True,
         disable_tp: bool = False,
     ):
-        nn.Module.__init__(self)
-
-        # Keep input parameters
-        self.input_size = input_size
-        self.output_size = output_size
-        self.skip_bias_add = skip_bias_add
-        if params_dtype is None:
-            params_dtype = torch.get_default_dtype()
-        self.params_dtype = params_dtype
-        self.quant_config = quant_config
-        self.prefix = prefix
-        if quant_config is None:
-            self.quant_method: Optional[
-                QuantizeMethodBase] = UnquantizedLinearMethod()
+        if hasattr(self, "tp_rank") and hasattr(self, "tp_size"):
+            tp_rank = self.tp_rank
+            tp_size = self.tp_size
+            super().__init__(input_size, output_size, skip_bias_add,
+                             params_dtype, quant_config, prefix, return_bias,
+                             disable_tp)
+            self.tp_rank = tp_rank
+            self.tp_size = tp_size
         else:
-            self.quant_method = quant_config.get_quant_method(self,
-                                                              prefix=prefix)
-        self.return_bias = return_bias
-        self.disable_tp = disable_tp
+            super().__init__(input_size, output_size, skip_bias_add,
+                             params_dtype, quant_config, prefix, return_bias,
+                             disable_tp)