ScaleLLM/src/layers/linear/qkv_parallel_linear_test.cpp at 1e1e948750681752ace1fc608aedda7430263840 · vectorch-ai/ScaleLLM · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
#include "qkv_parallel_linear.h"

#include <gtest/gtest.h>
#include <torch/torch.h>

#include <tuple>

#include "model_loader/state_dict.h"

namespace llm {

class QKVColumnParallelLinearTest
    : public ::testing::TestWithParam<std::tuple<int64_t /*n_tokens*/,
                                                 int64_t /*n_heads*/,
                                                 int64_t /*n_kv_heads*/,
                                                 int64_t /*n_shards*/,
                                                 int64_t /*head_dim*/,
                                                 int64_t /*hidden_size*/>> {};

TEST_P(QKVColumnParallelLinearTest, LoadFusedWeight) {
  const auto& [n_tokens, n_heads, n_kv_heads, n_shards, head_dim, hidden_size] =
      GetParam();

  ASSERT_LE(n_kv_heads, n_heads);
  ASSERT_LE(n_shards, n_heads);

  torch::Device device(torch::kCPU);
  torch::ScalarType dtype(torch::kFloat);
  const auto options = torch::dtype(dtype).device(device);
  std::unordered_map<std::string, torch::Tensor> state_dict_data;
  // Allocate transposed weight matrix
  state_dict_data["query.weight"] =
      torch::randn({n_heads * head_dim, hidden_size}, options);
  state_dict_data["key.weight"] =
      torch::randn({n_kv_heads * head_dim, hidden_size}, options);
  state_dict_data["value.weight"] =
      torch::randn({n_kv_heads * head_dim, hidden_size}, options);

  // weight is not sharded
  StateDict state_dict(state_dict_data);

  int64_t n_kv_shards = std::min(n_kv_heads, n_shards);
  auto query_chunks = state_dict_data["query.weight"].chunk(
      /*chunks=*/n_shards, /*dim=*/0);
  auto key_chunks = state_dict_data["key.weight"].chunk(
      /*chunks=*/n_kv_shards, /*dim=*/0);
  auto value_chunks = state_dict_data["value.weight"].chunk(
      /*chunks=*/n_kv_shards, /*dim=*/0);

  // test load weight
  for (int32_t shard_id = 0; shard_id < n_shards; ++shard_id) {
    QuantArgs quant_args;
    ParallelArgs parallel_args(shard_id, n_shards, nullptr);
    QKVColumnParallelLinearImpl linear(
        hidden_size,
        n_heads,
        n_kv_heads,
        head_dim,
        std::vector<std::string>{"query.", "key.", "value."},
        /*bias=*/false,
        /*gather_output=*/false,
        quant_args,
        parallel_args,
        options);
    EXPECT_EQ(linear.load(state_dict), 3);
    EXPECT_TRUE(linear.verify());

    // generate random input and compare with the output
    auto input = torch::randn({n_tokens, hidden_size}, options);
    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(q, query, /*rtol=*/1e-5, /*atol=*/1e-5));

    auto key = input.matmul(key_chunks[kv_shard_id].t());
    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(v, value, /*rtol=*/1e-5, /*atol=*/1e-5));
  }
}

INSTANTIATE_TEST_SUITE_P(
    QKVLinearTestSuite,
    QKVColumnParallelLinearTest,
    ::testing::Combine(::testing::Values(10, 32),      // n_tokens
                       ::testing::Values(8, 16, 32),   // n_heads
                       ::testing::Values(1, 2, 4, 8),  // n_kv_heads
                       ::testing::Values(1, 2, 4, 8),  // n_shards
                       ::testing::Values(1, 32, 64),   // head_dim
                       ::testing::Values(1, 32, 64))   // hidden_size
);

}  // namespace llm