refactor

Author: guocuimi

src/layers/linear/multi_parallel_linear.cpp

@@ -3,6 +3,7 @@
 #include <glog/logging.h>
 #include <torch/torch.h>
 
+#include "layers/linear/linear.h"
 #include "model_parallel/parallel_args.h"
 #include "parallel_linear.h"
 #include "quantization/quant_args.h"
@@ -19,38 +20,35 @@ MultiColumnParallelLinearImpl::MultiColumnParallelLinearImpl(
     const ParallelArgs& parallel_args,
     const torch::TensorOptions& options) {
   // check if the linear layers can be fused
-  fused_ = quant_args.can_be_fused();
-  if (fused_) {
+  std::shared_ptr<MultiParallelLinearImpl> linear;
+  if (quant_args.can_be_fused()) {
     // fused linear layer
-    fused_linear_ = register_module("fused_linear",
-                                    FusedColumnParallelLinear(in_features,
-                                                              out_features_vec,
-                                                              prefixes,
-                                                              bias,
-                                                              gather_output,
-                                                              parallel_args,
-                                                              options),
-                                    /*selector=*/nullptr);
+    linear = register_module("fused_linear",
+                             FusedColumnParallelLinear(in_features,
+                                                       out_features_vec,
+                                                       prefixes,
+                                                       bias,
+                                                       gather_output,
+                                                       parallel_args,
+                                                       options),
+                             /*selector=*/nullptr);
   } else {
     // non-fused linear layers
-    grouped_linear_ =
-        register_module("grouped_linear",
-                        GroupedColumnParallelLinear(in_features,
-                                                    out_features_vec,
-                                                    prefixes,
-                                                    bias,
-                                                    gather_output,
-                                                    parallel_args,
-                                                    options),
-                        /*selector=*/nullptr);
+    linear = register_module("grouped_linear",
+                             GroupedColumnParallelLinear(in_features,
+                                                         out_features_vec,
+                                                         prefixes,
+                                                         bias,
+                                                         gather_output,
+                                                         parallel_args,
+                                                         options),
+                             /*selector=*/nullptr);
   }
+  linear_ = linear;
 }
 
 std::vector<torch::Tensor> MultiColumnParallelLinearImpl::forward(
     torch::Tensor input) {
-  if (fused_) {
-    return fused_linear_(input);
-  }
-  return grouped_linear_(input);
+  return linear_(input);
 }
 }  // namespace llm

src/layers/linear/multi_parallel_linear.h

@@ -3,11 +3,10 @@
 #include <glog/logging.h>
 #include <torch/torch.h>
 
-// #include "linear.h"
+#include "linear.h"
 #include "model_parallel/parallel_args.h"
 #include "module/module.h"
 #include "module/module_holder.h"
-#include "parallel_linear.h"
 #include "quantization/quant_args.h"
 
 namespace llm {
@@ -25,18 +24,8 @@ class MultiColumnParallelLinearImpl : public Module {
 
   std::vector<torch::Tensor> forward(torch::Tensor input);
 
-  // whether the linear layer is fused
-  bool fused() const { return fused_; }
-
  private:
-  // non-fused linear layers
-  GroupedColumnParallelLinear grouped_linear_{nullptr};
-
-  // fused linear layer
-  FusedColumnParallelLinear fused_linear_{nullptr};
-
-  // whether the linear layer is fused
-  bool fused_ = false;
+  MultiParallelLinear linear_{nullptr};
 };
 LLM_MODULE(MultiColumnParallelLinear);