|
| 1 | +#include "cpu_model.h" |
| 2 | + |
| 3 | +#include <algorithm> |
| 4 | +#include <cmath> |
| 5 | +#include <format> |
| 6 | +#include <memory> |
| 7 | +#include <set> |
| 8 | + |
| 9 | +#include "dequant.h" |
| 10 | + |
| 11 | +namespace onebit { |
| 12 | + |
| 13 | +namespace { |
| 14 | + |
| 15 | +// KV cache length. Kept well under the training context so the reference |
| 16 | +// stays small; raise it when a test needs longer sequences. |
| 17 | +constexpr uint32_t kMaxCtx = 4096; |
| 18 | + |
| 19 | +void rms_norm(const float* x, const float* w, float* out, uint32_t n, float eps) { |
| 20 | + double ss = 0.0; |
| 21 | + for (uint32_t i = 0; i < n; ++i) ss += double(x[i]) * x[i]; |
| 22 | + const float scale = 1.0f / std::sqrt(float(ss / n) + eps); |
| 23 | + for (uint32_t i = 0; i < n; ++i) out[i] = w[i] * (x[i] * scale); |
| 24 | +} |
| 25 | + |
| 26 | +void rope(float* x, uint32_t n_heads, uint32_t head_dim, uint32_t pos, float theta, RopeStyle style) { |
| 27 | + const uint32_t half = head_dim / 2; |
| 28 | + for (uint32_t i = 0; i < half; ++i) { |
| 29 | + const float inv_freq = 1.0f / std::pow(theta, float(2 * i) / float(head_dim)); |
| 30 | + const float angle = float(pos) * inv_freq; |
| 31 | + const float c = std::cos(angle), s = std::sin(angle); |
| 32 | + for (uint32_t h = 0; h < n_heads; ++h) { |
| 33 | + float* v = x + h * head_dim; |
| 34 | + const uint32_t a = style == RopeStyle::Neox ? i : 2 * i; |
| 35 | + const uint32_t b = style == RopeStyle::Neox ? i + half : 2 * i + 1; |
| 36 | + const float x0 = v[a], x1 = v[b]; |
| 37 | + v[a] = x0 * c - x1 * s; |
| 38 | + v[b] = x0 * s + x1 * c; |
| 39 | + } |
| 40 | + } |
| 41 | +} |
| 42 | + |
| 43 | +float silu(float x) { return x / (1.0f + std::exp(-x)); } |
| 44 | + |
| 45 | +class Loader { |
| 46 | +public: |
| 47 | + explicit Loader(const GgufFile& f) : f_(f) {} |
| 48 | + |
| 49 | + // Dequantizes tensor `name`, checking its shape is exactly `shape`. |
| 50 | + std::expected<std::vector<float>, std::string> take(const std::string& name, std::vector<uint64_t> shape) { |
| 51 | + const GgufTensor* t = f_.tensor(name); |
| 52 | + if (!t) return std::unexpected(std::format("missing tensor {}", name)); |
| 53 | + if (t->ne != shape) return std::unexpected(std::format("tensor {} has an unexpected shape", name)); |
| 54 | + std::vector<float> out(t->n_elements()); |
| 55 | + if (auto r = dequantize(t->type, t->data, out.data(), out.size()); !r) |
| 56 | + return std::unexpected(std::format("tensor {}: {}", name, r.error())); |
| 57 | + used_.insert(name); |
| 58 | + return out; |
| 59 | + } |
| 60 | + |
| 61 | + bool has(const std::string& name) const { return f_.tensor(name) != nullptr; } |
| 62 | + |
| 63 | + std::expected<void, std::string> check_all_used() const { |
| 64 | + for (const GgufTensor& t : f_.tensors()) |
| 65 | + if (!used_.contains(t.name)) return std::unexpected(std::format("tensor {} is not used by the reference", t.name)); |
| 66 | + return {}; |
| 67 | + } |
| 68 | + |
| 69 | +private: |
| 70 | + const GgufFile& f_; |
| 71 | + std::set<std::string> used_; |
| 72 | +}; |
| 73 | + |
| 74 | +} // namespace |
| 75 | + |
| 76 | +std::expected<CpuModel, std::string> CpuModel::load(const std::string& gguf_path, size_t n_threads) { |
| 77 | + auto file = GgufFile::open(gguf_path); |
| 78 | + if (!file) return std::unexpected(file.error()); |
| 79 | + auto cfg = read_model_config(*file); |
| 80 | + if (!cfg) return std::unexpected(cfg.error()); |
| 81 | + |
| 82 | + CpuModel m; |
| 83 | + m.cfg_ = *cfg; |
| 84 | + const ModelConfig& c = m.cfg_; |
| 85 | + const uint64_t E = c.n_embd, D = c.head_dim, Q = uint64_t{c.n_head} * D, KV = uint64_t{c.n_head_kv} * D; |
| 86 | + const uint64_t F = c.n_ff, V = c.n_vocab; |
| 87 | + |
| 88 | + Loader ld(*file); |
| 89 | +#define TAKE(dst, name, ...) \ |
| 90 | + do { \ |
| 91 | + auto r_ = ld.take(name, {__VA_ARGS__}); \ |
| 92 | + if (!r_) return std::unexpected(std::move(r_.error())); \ |
| 93 | + dst = std::move(*r_); \ |
| 94 | + } while (0) |
| 95 | + |
| 96 | + TAKE(m.tok_embd_, "token_embd.weight", E, V); |
| 97 | + TAKE(m.output_norm_, "output_norm.weight", E); |
| 98 | + if (ld.has("output.weight")) TAKE(m.output_, "output.weight", E, V); |
| 99 | + |
| 100 | + m.layers_.resize(c.n_layer); |
| 101 | + for (uint32_t i = 0; i < c.n_layer; ++i) { |
| 102 | + Layer& L = m.layers_[i]; |
| 103 | + const std::string p = std::format("blk.{}.", i); |
| 104 | + TAKE(L.attn_norm, p + "attn_norm.weight", E); |
| 105 | + TAKE(L.wq, p + "attn_q.weight", E, Q); |
| 106 | + TAKE(L.wk, p + "attn_k.weight", E, KV); |
| 107 | + TAKE(L.wv, p + "attn_v.weight", E, KV); |
| 108 | + TAKE(L.wo, p + "attn_output.weight", Q, E); |
| 109 | + TAKE(L.q_norm, p + "attn_q_norm.weight", D); |
| 110 | + TAKE(L.k_norm, p + "attn_k_norm.weight", D); |
| 111 | + TAKE(L.ffn_norm, p + "ffn_norm.weight", E); |
| 112 | + TAKE(L.w_gate, p + "ffn_gate.weight", E, F); |
| 113 | + TAKE(L.w_up, p + "ffn_up.weight", E, F); |
| 114 | + TAKE(L.w_down, p + "ffn_down.weight", F, E); |
| 115 | + } |
| 116 | +#undef TAKE |
| 117 | + if (auto r = ld.check_all_used(); !r) return std::unexpected(r.error()); |
| 118 | + |
| 119 | + m.n_ctx_ = std::min(c.n_ctx_train, kMaxCtx); |
| 120 | + const size_t cache = size_t{c.n_layer} * m.n_ctx_ * KV; |
| 121 | + m.k_cache_.assign(cache, 0.0f); |
| 122 | + m.v_cache_.assign(cache, 0.0f); |
| 123 | + m.x_.resize(E); |
| 124 | + m.xn_.resize(E); |
| 125 | + m.q_.resize(Q); |
| 126 | + m.k_.resize(KV); |
| 127 | + m.v_.resize(KV); |
| 128 | + m.attn_.resize(Q); |
| 129 | + m.scores_.resize(m.n_ctx_); |
| 130 | + m.gate_.resize(F); |
| 131 | + m.up_.resize(F); |
| 132 | + m.logits_.resize(V); |
| 133 | + m.pool_ = std::make_unique<ThreadPool>(n_threads); |
| 134 | + return m; |
| 135 | +} |
| 136 | + |
| 137 | +void CpuModel::reset() { n_past_ = 0; } |
| 138 | + |
| 139 | +void CpuModel::matvec(const std::vector<float>& w, const float* x, float* y, uint32_t rows, uint32_t cols) { |
| 140 | + pool_->parallel_for(rows, [&](size_t begin, size_t end) { |
| 141 | + for (size_t r = begin; r < end; ++r) { |
| 142 | + const float* row = w.data() + r * cols; |
| 143 | + float acc = 0.0f; |
| 144 | + for (uint32_t i = 0; i < cols; ++i) acc += row[i] * x[i]; |
| 145 | + y[r] = acc; |
| 146 | + } |
| 147 | + }); |
| 148 | +} |
| 149 | + |
| 150 | +std::expected<const std::vector<float>*, std::string> CpuModel::forward(int32_t token) { |
| 151 | + const ModelConfig& c = cfg_; |
| 152 | + if (token < 0 || uint32_t(token) >= c.n_vocab) return std::unexpected(std::format("token {} out of range", token)); |
| 153 | + if (n_past_ >= n_ctx_) return std::unexpected(std::format("context full ({} tokens)", n_ctx_)); |
| 154 | + |
| 155 | + const uint32_t E = c.n_embd, D = c.head_dim, H = c.n_head, HKV = c.n_head_kv; |
| 156 | + const uint32_t KV = HKV * D, pos = n_past_, group = H / HKV; |
| 157 | + const float att_scale = 1.0f / std::sqrt(float(D)); |
| 158 | + |
| 159 | + std::copy_n(tok_embd_.data() + size_t(token) * E, E, x_.data()); |
| 160 | + |
| 161 | + for (uint32_t l = 0; l < c.n_layer; ++l) { |
| 162 | + const Layer& L = layers_[l]; |
| 163 | + |
| 164 | + rms_norm(x_.data(), L.attn_norm.data(), xn_.data(), E, c.rms_eps); |
| 165 | + matvec(L.wq, xn_.data(), q_.data(), H * D, E); |
| 166 | + matvec(L.wk, xn_.data(), k_.data(), KV, E); |
| 167 | + matvec(L.wv, xn_.data(), v_.data(), KV, E); |
| 168 | + for (uint32_t h = 0; h < H; ++h) rms_norm(q_.data() + h * D, L.q_norm.data(), q_.data() + h * D, D, c.rms_eps); |
| 169 | + for (uint32_t h = 0; h < HKV; ++h) rms_norm(k_.data() + h * D, L.k_norm.data(), k_.data() + h * D, D, c.rms_eps); |
| 170 | + rope(q_.data(), H, D, pos, c.rope_theta, c.rope_style); |
| 171 | + rope(k_.data(), HKV, D, pos, c.rope_theta, c.rope_style); |
| 172 | + |
| 173 | + float* kc = k_cache_.data() + (size_t(l) * n_ctx_) * KV; |
| 174 | + float* vc = v_cache_.data() + (size_t(l) * n_ctx_) * KV; |
| 175 | + std::copy_n(k_.data(), KV, kc + size_t(pos) * KV); |
| 176 | + std::copy_n(v_.data(), KV, vc + size_t(pos) * KV); |
| 177 | + |
| 178 | + for (uint32_t h = 0; h < H; ++h) { |
| 179 | + const float* q = q_.data() + h * D; |
| 180 | + const uint32_t kvh = h / group; |
| 181 | + float mx = -INFINITY; |
| 182 | + for (uint32_t t = 0; t <= pos; ++t) { |
| 183 | + const float* k = kc + size_t(t) * KV + kvh * D; |
| 184 | + float s = 0.0f; |
| 185 | + for (uint32_t i = 0; i < D; ++i) s += q[i] * k[i]; |
| 186 | + scores_[t] = s * att_scale; |
| 187 | + mx = std::max(mx, scores_[t]); |
| 188 | + } |
| 189 | + double sum = 0.0; |
| 190 | + for (uint32_t t = 0; t <= pos; ++t) { |
| 191 | + scores_[t] = std::exp(scores_[t] - mx); |
| 192 | + sum += scores_[t]; |
| 193 | + } |
| 194 | + float* out = attn_.data() + h * D; |
| 195 | + std::fill_n(out, D, 0.0f); |
| 196 | + for (uint32_t t = 0; t <= pos; ++t) { |
| 197 | + const float p = float(scores_[t] / sum); |
| 198 | + const float* v = vc + size_t(t) * KV + kvh * D; |
| 199 | + for (uint32_t i = 0; i < D; ++i) out[i] += p * v[i]; |
| 200 | + } |
| 201 | + } |
| 202 | + matvec(L.wo, attn_.data(), xn_.data(), E, H * D); |
| 203 | + for (uint32_t i = 0; i < E; ++i) x_[i] += xn_[i]; |
| 204 | + |
| 205 | + rms_norm(x_.data(), L.ffn_norm.data(), xn_.data(), E, c.rms_eps); |
| 206 | + matvec(L.w_gate, xn_.data(), gate_.data(), c.n_ff, E); |
| 207 | + matvec(L.w_up, xn_.data(), up_.data(), c.n_ff, E); |
| 208 | + for (uint32_t i = 0; i < c.n_ff; ++i) gate_[i] = silu(gate_[i]) * up_[i]; |
| 209 | + matvec(L.w_down, gate_.data(), xn_.data(), E, c.n_ff); |
| 210 | + for (uint32_t i = 0; i < E; ++i) x_[i] += xn_[i]; |
| 211 | + } |
| 212 | + |
| 213 | + rms_norm(x_.data(), output_norm_.data(), xn_.data(), E, c.rms_eps); |
| 214 | + matvec(output_.empty() ? tok_embd_ : output_, xn_.data(), logits_.data(), c.n_vocab, E); |
| 215 | + ++n_past_; |
| 216 | + return &logits_; |
| 217 | +} |
| 218 | + |
| 219 | +} // namespace onebit |
0 commit comments