refactor: clean up legacy load_state_dict for linear layers

Author: guocuimi

src/layers/linear.h

@@ -26,12 +26,6 @@ class ParallelLinearImpl : public Module {
 
   virtual void verify_loaded_weights(const std::string& prefix = "") const = 0;
 
-  // load state dict with a transform function
-  virtual void load_state_dict(const StateDict& /*state_dict*/,
-                               TensorTransform /*transform_func*/) {
-    LOG(FATAL) << "not implemented";
-  }
-
   // special load_state_dict for fused cases
   virtual void load_state_dict(const StateDict& /*state_dict*/,
                                const std::vector<std::string>& /*prefixes*/) {

src/layers/linear_impl.cpp

@@ -55,23 +55,15 @@ torch::Tensor ColumnParallelLinearImpl::forward(torch::Tensor input) {
 
 // load the weight from the checkpoint
 void ColumnParallelLinearImpl::load_state_dict(const StateDict& state_dict) {
-  // call load_state_dict with identity transform
-  load_state_dict(state_dict,
-                  [](const torch::Tensor& tensor) { return tensor; });
-}
-
-void ColumnParallelLinearImpl::load_state_dict(const StateDict& state_dict,
-                                               TensorTransform transform_func) {
-  CHECK(transform_func != nullptr) << "transform_func must be provided";
   const auto rank = parallel_args_.rank();
   const auto world_size = parallel_args_.world_size();
 
   // load sharded weights on dim 0
-  LOAD_SHARDED_WEIGHT_WITH_TRANSFORM(weight, 0);
+  LOAD_SHARDED_WEIGHT(weight, 0);
 
   if (bias_.defined()) {
     // load sharded bias on dim 0
-    LOAD_SHARDED_WEIGHT_WITH_TRANSFORM(bias, 0);
+    LOAD_SHARDED_WEIGHT(bias, 0);
   }
 }
 

src/layers/linear_impl.h

@@ -26,10 +26,6 @@ class ColumnParallelLinearImpl : public ParallelLinearImpl {
   // load the weight from the checkpoint
   void load_state_dict(const StateDict& state_dict) override;
 
-  // load state dict with a transform function
-  void load_state_dict(const StateDict& state_dict,
-                       TensorTransform transform_func) override;
-
   // special load_state_dict for fused cases
   void load_state_dict(const StateDict& state_dict,
                        const std::vector<std::string>& prefixes) override;

src/layers/qkv_linear.h

@@ -27,8 +27,11 @@ class QKVColumnParallelLinearImpl : public Module {
                               const ParallelArgs& parallel_args,
                               const torch::TensorOptions& options);
 
-  std::vector<torch::Tensor> forward(torch::Tensor input) {
-    return parallel_linear_->forward(input);
+  // returns (query, key, value)
+  std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> forward(
+      torch::Tensor input) {
+    const auto qkv = parallel_linear_->forward(input);
+    return {qkv[0], qkv[1], qkv[2]};
   }
 
  private:

src/layers/qkv_linear_test.cpp

@@ -67,19 +67,19 @@ TEST_P(QKVColumnParallelLinearTest, LoadFusedWeight) {
 
     // generate random input and compare with the output
     auto input = torch::randn({n_tokens, hidden_size}, options);
-    auto qkv = linear.forward(input);
+    const auto [q, k, v] = linear.forward(input);
 
     const int64_t kv_shard_id =
         n_kv_heads >= n_shards ? shard_id : n_kv_heads * shard_id / n_shards;
 
     auto query = input.matmul(query_chunks[shard_id].t());
-    EXPECT_TRUE(torch::allclose(qkv[0], query, /*rtol=*/1e-5, /*atol=*/1e-5));
+    EXPECT_TRUE(torch::allclose(q, query, /*rtol=*/1e-5, /*atol=*/1e-5));
 
     auto key = input.matmul(key_chunks[kv_shard_id].t());
-    EXPECT_TRUE(torch::allclose(qkv[1], key, /*rtol=*/1e-5, /*atol=*/1e-5));
+    EXPECT_TRUE(torch::allclose(k, key, /*rtol=*/1e-5, /*atol=*/1e-5));
 
     auto value = input.matmul(value_chunks[kv_shard_id].t());
-    EXPECT_TRUE(torch::allclose(qkv[2], value, /*rtol=*/1e-5, /*atol=*/1e-5));
+    EXPECT_TRUE(torch::allclose(v, value, /*rtol=*/1e-5, /*atol=*/1e-5));
   }
 }
 

src/models/alibaba/qwen2.h

@@ -133,10 +133,9 @@ class QWen2AttentionImpl : public Module {
                         const InputParameters& input_params) {
     // (num_tokens, dim) x (dim, n_local_heads * head_dim)
     // => (num_tokens, n_local_heads * head_dim)
-    const auto qkv = qkv_proj_(x);
+    const auto [q, k, v] = qkv_proj_(x);
     // calculate attention, output: (num_tokens, n_local_heads * head_dim)
-    const auto output =
-        atten_(qkv[0], qkv[1], qkv[2], positions, kv_cache, input_params);
+    const auto output = atten_(q, k, v, positions, kv_cache, input_params);
     return o_proj_(output);
   }
 

src/models/google/gemma.h

@@ -128,11 +128,10 @@ class GemmaAttentionImpl : public Module {
                         const InputParameters& input_params) {
     // (num_tokens, dim) x (dim, n_local_heads * head_dim)
     // => (num_tokens, n_local_heads * head_dim)
-    const auto qkv = qkv_proj_(x);
+    const auto [q, k, v] = qkv_proj_(x);
     // calculate attention,
     // output: (num_tokens, n_local_heads*head_dim)
-    const auto output =
-        atten_(qkv[0], qkv[1], qkv[2], positions, kv_cache, input_params);
+    const auto output = atten_(q, k, v, positions, kv_cache, input_params);
     return o_proj_(output);
   }
 

src/models/google/gemma2.h

@@ -132,11 +132,10 @@ class Gemma2AttentionImpl : public Module {
                         const InputParameters& input_params) {
     // (num_tokens, dim) x (dim, n_local_heads * head_dim)
     // => (num_tokens, n_local_heads * head_dim)
-    const auto qkv = qkv_proj_(x);
+    const auto [q, k, v] = qkv_proj_(x);
     // calculate attention,
     // output: (num_tokens, n_local_heads*head_dim)
-    const auto output =
-        atten_(qkv[0], qkv[1], qkv[2], positions, kv_cache, input_params);
+    const auto output = atten_(q, k, v, positions, kv_cache, input_params);
     return o_proj_(output);
   }
 

src/models/meta/llama.h

@@ -127,10 +127,9 @@ class LlamaAttentionImpl : public Module {
                         const InputParameters& input_params) {
     // (num_tokens, dim) x (dim, n_local_heads * head_dim)
     // => (num_tokens, n_local_heads * head_dim)
-    const auto qkv = qkv_proj_(x);
+    const auto [q, k, v] = qkv_proj_(x);
     // calculate attention, output: (num_tokens, n_local_heads * head_dim)
-    const auto output =
-        atten_(qkv[0], qkv[1], qkv[2], positions, kv_cache, input_params);
+    const auto output = atten_(q, k, v, positions, kv_cache, input_params);
     return o_proj_(output);
   }