llm.cpp/llmcpp/optim_test.cpp at master · zhangpiu/llm.cpp · 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
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
#include "optim.hpp"
#include "gtest/gtest.h"

TEST(Optimizer, SGD) {
  /*
torch.set_printoptions(precision=6)
torch.manual_seed(42)
m = nn.Linear(3, 2)
optimizer = torch.optim.SGD(m.parameters(), lr=0.01)
x = torch.randn(4, 3)
for _ in range(10):
  y = m(x)
  loss = torch.sum(y)
  optimizer.zero_grad()
  loss.backward()
  optimizer.step()
  */

  nn::ManualSeed(42);
  int B = 4, in_features = 3, out_features = 2;
  nn::Linear m(in_features, out_features, true);

  // forward
  std::vector<float> x(B * in_features), y(B * out_features);
  nn::NormalFill(absl::MakeSpan(x));
  auto xm = MakeConstMatrix(x.data(), B, in_features);
  auto ym = MakeMatrix(y.data(), B, out_features);

  // optimizer
  std::vector<nn::Parameter*> parameters;
  m.Parameters(&parameters);
  optim::SGD sgd(parameters, 0.01);

  // backward
  std::vector<float> y_grad(y.size(), 1.0f);
  std::vector<float> x_grad(x.size(), 0.f);
  auto y_gradm = MakeConstMatrix(y_grad.data(), B, out_features);
  auto x_gradm = MakeMatrix(x_grad.data(), B, in_features);

  int step = 10;
  for (int i = 0; i < step; ++i) {
    m.Forward(xm, ym);
    sgd.ZeroGrad();
    m.Backward(xm, y_gradm, x_gradm);
    sgd.Step();
  }

  auto weight = m.weight_->span<float>();
  auto bias = m.bias_->span<float>();
  std::vector<float> expected_weight = {0.732981, 0.469633,  -0.589639,
                                        0.821935, -0.136072, -0.337878};
  std::vector<float> expected_bias = {-0.681086, -0.060932};
  for (size_t i = 0; i < expected_weight.size(); ++i) {
    EXPECT_NEAR(expected_weight[i], weight[i], 1e-5);
  }
  for (size_t i = 0; i < expected_bias.size(); ++i) {
    EXPECT_NEAR(expected_bias[i], bias[i], 1e-5);
  }
}

TEST(Optimizer, AdamW) {
  /*
torch.set_printoptions(precision=6)
torch.manual_seed(42)
m = nn.Linear(3, 2)
optimizer = torch.optim.AdamW(m.parameters(), lr=0.01, betas=(0.9, 0.999),
eps=1e-8, weight_decay=0.001)
x = torch.randn(4, 3)
for _ in range(10):
  y = m(x)
  loss = torch.sum(y)
  optimizer.zero_grad()
  loss.backward()
  optimizer.step()
  */

  nn::ManualSeed(42);
  int B = 4, in_features = 3, out_features = 2;
  nn::Linear m(in_features, out_features, true);

  // forward
  std::vector<float> x(B * in_features), y(B * out_features);
  nn::NormalFill(absl::MakeSpan(x));
  auto xm = MakeConstMatrix(x.data(), B, in_features);
  auto ym = MakeMatrix(y.data(), B, out_features);

  // optimizer
  std::vector<nn::Parameter*> parameters;
  m.Parameters(&parameters);
  optim::AdamW adam_w(parameters, 0.01f, 0.9f, 0.999f, 1e-8f, 0.001f);

  // backward
  std::vector<float> y_grad(y.size(), 1.0f);
  std::vector<float> x_grad(x.size(), 0.f);
  auto y_gradm = MakeConstMatrix(y_grad.data(), B, out_features);
  auto x_gradm = MakeMatrix(x_grad.data(), B, in_features);

  int step = 10;
  for (int i = 0; i < step; ++i) {
    m.Forward(xm, ym);
    adam_w.ZeroGrad();
    m.Backward(xm, y_gradm, x_gradm);
    adam_w.Step(i + 1);
  }

  auto weight = m.weight_->span<float>();
  auto bias = m.bias_->span<float>();
  std::vector<float> expected_weight = {0.541358, 0.379162,  -0.235239,
                                        0.630303, -0.226482, 0.016497};
  std::vector<float> expected_bias = {-0.381053, 0.239038};
  for (size_t i = 0; i < expected_weight.size(); ++i) {
    EXPECT_NEAR(expected_weight[i], weight[i], 1e-5);
  }
  for (size_t i = 0; i < expected_bias.size(); ++i) {
    EXPECT_NEAR(expected_bias[i], bias[i], 1e-5);
  }
}