SPB Git

spb/forge Public MIT

Forge — LLM training from scratch in pure C++20 + Metal on Apple Silicon.

llmtrainingcppmetalapple-silicon
C++ 61.2% C 23% Python 7.6% TeX 7.2% CMake 1.1%

Add QAT, MoE, and architecture-variant knobs — all config-selected

- quant "int8"|"ternary": per-row fake-quant each forward (BitNet-style
  absmean for ternary), straight-through estimator backward, f32 masters;
  wired through the Linear quantization seam
- n_experts/moe_top_k/n_shared_experts: softmax router, renormalized top-k
  gates (topk_renorm + row_scale ops, CPU+Metal), differentiable
  load-balance loss, DeepSeek-style always-active shared experts;
  v1 computes experts densely (correctness first)
- qk_norm (Qwen3/Gemma3), final_softcap (Gemma2), scale_embeddings (Gemma)
- new kernels: quant.metal, moe.metal, softcap in elementwise.metal
- CPU references + parity tests for every new op and full-model variants
  (QAT int8/ternary, MoE 4+1shared, qk-norm+softcap+embed-scale)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
simon-pierre boucher committed 5 days ago (Aug 5, 2026) parent ecaa53d

Showing 17 changed files with +849 and −10

added configs/gpt-50m-moe.json +41 −0
@@ -0,0 +1,41 @@
1 +{
2 + "_comment": "MoE variant: 4 experts (d_ff 1728 each) + top-2 softmax router per block, ~152M total params but only ~2 experts' worth active per token. v1 computes all experts densely (correctness first), so steps cost ~4x the dense MLP; the aux load-balance loss (E * sum mean_gate^2, weight 0.01) keeps the router from collapsing. Same data/steps as gpt-50m.",
3 + "model": {
4 + "name": "gpt-50m-moe",
5 + "n_layers": 10,
6 + "d_model": 640,
7 + "n_heads": 10,
8 + "n_kv_heads": 10,
9 + "d_ff": 1728,
10 + "vocab_size": 4096,
11 + "context_length": 1024,
12 + "tied_embeddings": true,
13 + "use_rope": true,
14 + "rope_theta": 10000.0,
15 + "norm": "rmsnorm",
16 + "norm_eps": 1e-06,
17 + "activation": "swiglu",
18 + "dropout": 0.0,
19 + "n_experts": 4,
20 + "moe_top_k": 2,
21 + "moe_aux_weight": 0.01
22 + },
23 + "train": {
24 + "lr": 0.0005,
25 + "min_lr_ratio": 0.1,
26 + "warmup_steps": 117,
27 + "max_steps": 1170,
28 + "beta1": 0.9,
29 + "beta2": 0.95,
30 + "eps": 1e-08,
31 + "weight_decay": 0.1,
32 + "grad_clip": 1.0,
33 + "batch_size": 8,
34 + "grad_accum_steps": 8,
35 + "precision": "f32",
36 + "checkpoint_every": 200,
37 + "eval_every": 100,
38 + "eval_batches": 20,
39 + "seed": 1337
40 + }
41 +}
added configs/gpt-50m-ternary.json +39 −0
@@ -0,0 +1,39 @@
1 +{
2 + "_comment": "gpt-50m with BitNet-style ternary QAT: every linear weight is fake-quantized to {-s, 0, +s} (per-row absmean scale) each forward, gradients flow via straight-through estimator into f32 master weights. Embeddings/lm_head stay f32. Same data/steps as gpt-50m for A/B comparison.",
3 + "model": {
4 + "name": "gpt-50m-ternary",
5 + "n_layers": 10,
6 + "d_model": 640,
7 + "n_heads": 10,
8 + "n_kv_heads": 10,
9 + "d_ff": 1728,
10 + "vocab_size": 4096,
11 + "context_length": 1024,
12 + "tied_embeddings": true,
13 + "use_rope": true,
14 + "rope_theta": 10000.0,
15 + "norm": "rmsnorm",
16 + "norm_eps": 1e-06,
17 + "activation": "swiglu",
18 + "dropout": 0.0,
19 + "quant": "ternary"
20 + },
21 + "train": {
22 + "lr": 0.0005,
23 + "min_lr_ratio": 0.1,
24 + "warmup_steps": 117,
25 + "max_steps": 1170,
26 + "beta1": 0.9,
27 + "beta2": 0.95,
28 + "eps": 1e-08,
29 + "weight_decay": 0.1,
30 + "grad_clip": 1.0,
31 + "batch_size": 8,
32 + "grad_accum_steps": 8,
33 + "precision": "f32",
34 + "checkpoint_every": 200,
35 + "eval_every": 100,
36 + "eval_batches": 20,
37 + "seed": 1337
38 + }
39 +}
modified src/kernels/elementwise.metal +18 −0
@@ -52,6 +52,24 @@ kernel void gelu_f32(device const float* x [[buffer(0)]],
52 52 out[gid] = 0.5f * v * (1.0f + precise::tanh(k * (v + 0.044715f * v * v * v)));
53 53 }
54 54
55 +// Gemma-style logit softcap: out = cap * tanh(x / cap).
56 +kernel void softcap_f32(device const float* x [[buffer(0)]],
57 + device float* out [[buffer(1)]],
58 + constant float& cap [[buffer(2)]],
59 + uint gid [[thread_position_in_grid]]) {
60 + out[gid] = cap * precise::tanh(x[gid] / cap);
61 +}
62 +
63 +// dx += dout * (1 - (y/cap)^2), using the forward output y.
64 +kernel void softcap_bwd_f32(device const float* y [[buffer(0)]],
65 + device const float* dout [[buffer(1)]],
66 + device float* dx [[buffer(2)]],
67 + constant float& cap [[buffer(3)]],
68 + uint gid [[thread_position_in_grid]]) {
69 + const float t = y[gid] / cap;
70 + dx[gid] = fma(dout[gid], 1.0f - t * t, dx[gid]);
71 +}
72 +
55 73 // ---- backward / accumulation kernels (all ACCUMULATE into their outputs) ----
56 74
57 75 kernel void accum_f32(device float* dst [[buffer(0)]],
added src/kernels/moe.metal +120 −0
@@ -0,0 +1,120 @@
1 +// Author: Simon-Pierre Boucher — contact@spboucher.ai
2 +//
3 +// Mixture-of-experts kernels. All operate on router-sized rows (E experts,
4 +// E <= 64 in practice), so every kernel is one thread per row with a plain
5 +// loop — no threadgroup reductions needed. row_scale kernels are flat
6 +// elementwise over [N, C] activations.
7 +#include <metal_stdlib>
8 +using namespace metal;
9 +
10 +// Generic row softmax backward: dx += p ∘ (dout − dot(dout, p)), row = last
11 +// dim. Thread-per-row; meant for small C (the router's E columns).
12 +kernel void softmax_bwd_f32(device const float* P [[buffer(0)]],
13 + device const float* DOUT [[buffer(1)]],
14 + device float* DX [[buffer(2)]],
15 + constant uint& C [[buffer(3)]],
16 + uint row [[thread_position_in_grid]]) {
17 + device const float* prob = P + ulong(row) * C;
18 + device const float* dout = DOUT + ulong(row) * C;
19 + device float* dx = DX + ulong(row) * C;
20 + float dot = 0.0f;
21 + for (uint j = 0; j < C; ++j) dot += dout[j] * prob[j];
22 + for (uint j = 0; j < C; ++j) dx[j] += prob[j] * (dout[j] - dot);
23 +}
24 +
25 +// Top-k gating: keep the k largest of each row of P [N, E], renormalize the
26 +// kept entries to sum to 1, zero the rest. Ties broken by lower index, which
27 +// keeps forward and backward selections identical.
28 +// p = (E, k)
29 +kernel void topk_renorm_f32(device const float* P [[buffer(0)]],
30 + device float* OUT [[buffer(1)]],
31 + constant uint2& p [[buffer(2)]],
32 + uint row [[thread_position_in_grid]]) {
33 + const uint E = p.x, K = p.y;
34 + device const float* in = P + ulong(row) * E;
35 + device float* out = OUT + ulong(row) * E;
36 +
37 + bool kept[64];
38 + for (uint j = 0; j < E; ++j) kept[j] = false;
39 + float S = 0.0f;
40 + for (uint sel = 0; sel < K; ++sel) {
41 + float best = -FLT_MAX;
42 + uint arg = 0;
43 + for (uint j = 0; j < E; ++j)
44 + if (!kept[j] && in[j] > best) { best = in[j]; arg = j; }
45 + kept[arg] = true;
46 + S += best;
47 + }
48 + const float inv = 1.0f / max(S, 1e-12f);
49 + for (uint j = 0; j < E; ++j) out[j] = kept[j] ? in[j] * inv : 0.0f;
50 +}
51 +
52 +// Backward of topk_renorm: for kept entries, g_i = p_i / S with S the kept
53 +// sum, so dp_i += (dg_i − Σ_j dg_j g_j) / S; dropped entries get zero. The
54 +// kept set is recomputed from P with the same tie-breaking as forward.
55 +// p = (E, k)
56 +kernel void topk_renorm_bwd_f32(device const float* P [[buffer(0)]],
57 + device const float* DOUT [[buffer(1)]],
58 + device float* DP [[buffer(2)]],
59 + constant uint2& p [[buffer(3)]],
60 + uint row [[thread_position_in_grid]]) {
61 + const uint E = p.x, K = p.y;
62 + device const float* in = P + ulong(row) * E;
63 + device const float* dout = DOUT + ulong(row) * E;
64 + device float* dp = DP + ulong(row) * E;
65 +
66 + bool kept[64];
67 + for (uint j = 0; j < E; ++j) kept[j] = false;
68 + float S = 0.0f;
69 + for (uint sel = 0; sel < K; ++sel) {
70 + float best = -FLT_MAX;
71 + uint arg = 0;
72 + for (uint j = 0; j < E; ++j)
73 + if (!kept[j] && in[j] > best) { best = in[j]; arg = j; }
74 + kept[arg] = true;
75 + S += best;
76 + }
77 + const float inv = 1.0f / max(S, 1e-12f);
78 + float dot = 0.0f;
79 + for (uint j = 0; j < E; ++j)
80 + if (kept[j]) dot += dout[j] * in[j] * inv;
81 + for (uint j = 0; j < E; ++j)
82 + if (kept[j]) dp[j] += (dout[j] - dot) * inv;
83 +}
84 +
85 +// 12-byte layout matching the host-side struct (uint3 would pad to 16).
86 +struct RowScaleParams { uint C, E, e; };
87 +
88 +// out[i, c] = x[i, c] * G[i, e] — scale each row of X by one gate column.
89 +// Flat over N*C.
90 +kernel void row_scale_f32(device const float* X [[buffer(0)]],
91 + device const float* G [[buffer(1)]],
92 + device float* OUT [[buffer(2)]],
93 + constant RowScaleParams& p [[buffer(3)]],
94 + uint gid [[thread_position_in_grid]]) {
95 + const uint i = gid / p.C;
96 + OUT[gid] = X[gid] * G[ulong(i) * p.E + p.e];
97 +}
98 +
99 +// dst[i, c] += x[i, c] * G[i, e] — accumulating variant (dx of row_scale).
100 +kernel void row_scale_acc_f32(device const float* X [[buffer(0)]],
101 + device const float* G [[buffer(1)]],
102 + device float* DST [[buffer(2)]],
103 + constant RowScaleParams& p [[buffer(3)]],
104 + uint gid [[thread_position_in_grid]]) {
105 + const uint i = gid / p.C;
106 + DST[gid] = fma(X[gid], G[ulong(i) * p.E + p.e], DST[gid]);
107 +}
108 +
109 +// dG[i, e] += dot(dout[i, :], x[i, :]) — gate gradient, one thread per row.
110 +kernel void row_scale_gate_bwd_f32(device const float* DOUT [[buffer(0)]],
111 + device const float* X [[buffer(1)]],
112 + device float* DG [[buffer(2)]],
113 + constant RowScaleParams& p [[buffer(3)]],
114 + uint row [[thread_position_in_grid]]) {
115 + device const float* dout = DOUT + ulong(row) * p.C;
116 + device const float* x = X + ulong(row) * p.C;
117 + float acc = 0.0f;
118 + for (uint j = 0; j < p.C; ++j) acc += dout[j] * x[j];
119 + DG[ulong(row) * p.E + p.e] += acc;
120 +}
added src/kernels/quant.metal +37 −0
@@ -0,0 +1,37 @@
1 +// Author: Simon-Pierre Boucher — contact@spboucher.ai
2 +//
3 +// Quantization-aware-training kernel: fake-quantize a weight matrix row by
4 +// row (per-output-channel scales). One thread per row, two strided passes
5 +// over the row (reduce scale, then quantize). Rows are weight-matrix sized
6 +// (C = in_features, a few thousand max) and this runs once per layer per
7 +// forward, so a simple thread-per-row layout is plenty.
8 +//
9 +// Backward is the straight-through estimator (identity), so there is no
10 +// backward kernel — autograd routes dout straight into the master weight's
11 +// grad.
12 +#include <metal_stdlib>
13 +using namespace metal;
14 +
15 +// p = (C, mode); mode 0 = int8 (absmax/127), mode 1 = ternary BitNet-style
16 +// (scale = mean|w|, values in {-s, 0, +s}).
17 +kernel void fake_quant_f32(device const float* W [[buffer(0)]],
18 + device float* OUT [[buffer(1)]],
19 + constant uint2& p [[buffer(2)]],
20 + uint row [[thread_position_in_grid]]) {
21 + const uint C = p.x;
22 + device const float* w = W + ulong(row) * C;
23 + device float* out = OUT + ulong(row) * C;
24 +
25 + if (p.y == 0) { // int8: symmetric absmax
26 + float amax = 0.0f;
27 + for (uint j = 0; j < C; ++j) amax = max(amax, fabs(w[j]));
28 + const float s = max(amax / 127.0f, 1e-12f);
29 + for (uint j = 0; j < C; ++j) out[j] = rint(w[j] / s) * s;
30 + } else { // ternary: BitNet b1.58 absmean
31 + float asum = 0.0f;
32 + for (uint j = 0; j < C; ++j) asum += fabs(w[j]);
33 + const float s = max(asum / float(C), 1e-12f);
34 + for (uint j = 0; j < C; ++j)
35 + out[j] = clamp(rint(w[j] / s), -1.0f, 1.0f) * s;
36 + }
37 +}
modified src/nn/attention.h +26 −1
@@ -28,7 +28,9 @@ public:
28 28 n_kv_heads_(cfg.n_kv_heads),
29 29 head_dim_(cfg.head_dim()),
30 30 use_rope_(cfg.use_rope),
31 − rope_theta_(cfg.rope_theta) {
31 + rope_theta_(cfg.rope_theta),
32 + qk_norm_(cfg.qk_norm),
33 + norm_eps_(cfg.norm_eps) {
32 34 const float std = 0.02f;
33 35 const int64_t C = cfg.d_model;
34 36 const int64_t Ckv = cfg.n_kv_heads * cfg.head_dim();
@@ -36,10 +38,21 @@ public:
36 38 wk_ = std::make_unique<Linear>(C, Ckv, false, std, rng);
37 39 wv_ = std::make_unique<Linear>(C, Ckv, false, std, rng);
38 40 wo_ = std::make_unique<Linear>(C, C, false, proj_std, rng); // residual projection
41 + const ops::QuantMode qm = quant_mode_from(cfg.quant);
42 + wq_->set_quant(qm);
43 + wk_->set_quant(qm);
44 + wv_->set_quant(qm);
45 + wo_->set_quant(qm);
39 46 absorb("wq", *wq_);
40 47 absorb("wk", *wk_);
41 48 absorb("wv", *wv_);
42 49 absorb("wo", *wo_);
50 + if (qk_norm_) {
51 + q_norm_w_ = register_param("q_norm.weight",
52 + Tensor::full({cfg.head_dim()}, 1.0f));
53 + k_norm_w_ = register_param("k_norm.weight",
54 + Tensor::full({cfg.head_dim()}, 1.0f));
55 + }
43 56 }
44 57
45 58 Var forward(const Var& x) const override {
@@ -51,6 +64,15 @@ public:
51 64 Var k = wk_->forward(x2d).reshaped({B, T, Ckv});
52 65 Var v = wv_->forward(x2d).reshaped({B, T, Ckv});
53 66
67 + if (qk_norm_) {
68 + // Per-head RMSNorm before RoPE (Qwen3/Gemma3): normalize over
69 + // head_dim by viewing heads as rows.
70 + q = ops::rmsnorm(q.reshaped({B * T * n_heads_, head_dim_}), q_norm_w_,
71 + norm_eps_).reshaped({B, T, C});
72 + k = ops::rmsnorm(k.reshaped({B * T * n_kv_heads_, head_dim_}), k_norm_w_,
73 + norm_eps_).reshaped({B, T, Ckv});
74 + }
75 +
54 76 if (use_rope_) {
55 77 q = ops::rope(q, n_heads_, rope_theta_, 0);
56 78 k = ops::rope(k, n_kv_heads_, rope_theta_, 0);
@@ -65,7 +87,10 @@ private:
65 87 int64_t n_heads_, n_kv_heads_, head_dim_;
66 88 bool use_rope_;
67 89 float rope_theta_;
90 + bool qk_norm_;
91 + float norm_eps_;
68 92 std::unique_ptr<Linear> wq_, wk_, wv_, wo_;
93 + Var q_norm_w_, k_norm_w_; // defined only when qk_norm
69 94 };
70 95
71 96 } // namespace forge::nn
modified src/nn/linear.cpp +1 −0
@@ -7,6 +7,7 @@ namespace forge::nn {
7 7
8 8 Var Linear::forward(const Var& x) const {
9 9 Var w = mask_.defined() ? ops::mul(weight_, mask_) : weight_;
10 + if (quant_ != ops::QuantMode::None) w = ops::fake_quant(w, quant_);
10 11 Var y = ops::matmul(x, w, /*transpose_a=*/false, /*transpose_b=*/true);
11 12 if (bias_.defined()) y = ops::add_bias(y, bias_);
12 13 return y;
modified src/nn/linear.h +14 −0
@@ -2,9 +2,18 @@
2 2 #pragma once
3 3
4 4 #include "nn/module.h"
5 +#include "ops/ops.h"
6 +
7 +#include <string>
5 8
6 9 namespace forge::nn {
7 10
11 +inline ops::QuantMode quant_mode_from(const std::string& s) {
12 + if (s == "int8") return ops::QuantMode::Int8;
13 + if (s == "ternary") return ops::QuantMode::Ternary;
14 + return ops::QuantMode::None;
15 +}
16 +
8 17 // y = x ⋅ Wᵀ (+ b). Weight stored [out, in] (PyTorch convention), x is
9 18 // [N, in] 2-D — callers flatten [B,T,C] with Var::reshaped.
10 19 //
@@ -32,12 +41,17 @@ public:
32 41 mask_ = mask.defined() ? Var(std::move(mask), /*requires_grad=*/false) : Var();
33 42 }
34 43
44 + // QAT: fake-quantize the weight each forward (STE backward). The master
45 + // weight and optimizer state stay f32.
46 + void set_quant(ops::QuantMode mode) { quant_ = mode; }
47 +
35 48 const Var& weight() const { return weight_; }
36 49
37 50 protected:
38 51 Var weight_;
39 52 Var bias_; // undefined if bias disabled
40 53 Var mask_; // undefined if no mask
54 + ops::QuantMode quant_ = ops::QuantMode::None;
41 55 };
42 56
43 57 } // namespace forge::nn
modified src/nn/mlp.h +82 −3
@@ -6,11 +6,24 @@
6 6 #include "ops/ops.h"
7 7
8 8 #include <memory>
9 +#include <string>
10 +#include <vector>
9 11
10 12 namespace forge::nn {
11 13
14 +// Interface so the MoE variant slots into TransformerBlock without branching
15 +// at every call site. aux() is the block's auxiliary loss ([1] Var, undefined
16 +// when the variant has none).
17 +class MLPBase : public Module {
18 +public:
19 + virtual ~MLPBase() = default;
20 + // x: [N, C] 2-D
21 + virtual Var forward(const Var& x) const = 0;
22 + virtual Var aux() const { return Var(); }
23 +};
24 +
12 25 // SwiGLU: w2( silu(x w1) ⊙ (x w3) ) | GELU: proj( gelu(fc(x)) )
13 −class MLP : public Module {
26 +class MLP : public MLPBase {
14 27 public:
15 28 MLP(const ModelConfig& cfg, float proj_std, std::mt19937_64& rng)
16 29 : swiglu_(cfg.activation == "swiglu") {
@@ -21,13 +34,16 @@ public:
21 34 w3_ = std::make_unique<Linear>(C, F, false, std, rng);
22 35 }
23 36 w2_ = std::make_unique<Linear>(F, C, false, proj_std, rng); // residual projection
37 + const ops::QuantMode qm = quant_mode_from(cfg.quant);
38 + w1_->set_quant(qm);
39 + if (w3_) w3_->set_quant(qm);
40 + w2_->set_quant(qm);
24 41 absorb("w1", *w1_);
25 42 if (w3_) absorb("w3", *w3_);
26 43 absorb("w2", *w2_);
27 44 }
28 45
29 − // x: [N, C] 2-D
30 − Var forward(const Var& x) const {
46 + Var forward(const Var& x) const override {
31 47 if (swiglu_) {
32 48 Var gate = ops::silu(w1_->forward(x));
33 49 Var up = w3_->forward(x);
@@ -41,4 +57,67 @@ private:
41 57 std::unique_ptr<Linear> w1_, w2_, w3_;
42 58 };
43 59
60 +// Mixture-of-experts MLP: softmax router over n_experts expert MLPs (d_ff
61 +// each), top-k gates renormalized to sum 1. v1 computes EVERY expert densely
62 +// and weights by the (mostly zero) gates — correctness first; token
63 +// gather/scatter sparsity is a later optimization. The auxiliary
64 +// load-balance loss is the differentiable proxy E · Σ_e (mean_i gate
65 +// probs)² — minimized (at 1.0) by a uniform router — scaled by
66 +// moe_aux_weight and added to the training loss by Transformer::loss.
67 +class MoEMLP : public MLPBase {
68 +public:
69 + MoEMLP(const ModelConfig& cfg, float proj_std, std::mt19937_64& rng)
70 + : n_experts_(cfg.n_experts), top_k_(cfg.moe_top_k),
71 + aux_weight_(cfg.moe_aux_weight) {
72 + router_ = std::make_unique<Linear>(cfg.d_model, n_experts_, false, 0.02f, rng);
73 + absorb("router", *router_);
74 + experts_.reserve(size_t(n_experts_));
75 + for (int64_t e = 0; e < n_experts_; ++e) {
76 + experts_.push_back(std::make_unique<MLP>(cfg, proj_std, rng));
77 + absorb("experts." + std::to_string(e), *experts_.back());
78 + }
79 + // Always-active experts, added ungated (DeepSeek/Kimi style): they
80 + // absorb common knowledge so routed experts can specialize.
81 + for (int64_t s = 0; s < cfg.n_shared_experts; ++s) {
82 + shared_.push_back(std::make_unique<MLP>(cfg, proj_std, rng));
83 + absorb("shared." + std::to_string(s), *shared_.back());
84 + }
85 + }
86 +
87 + Var forward(const Var& x) const override {
88 + const int64_t N = x.value().size(0);
89 + Var probs = ops::softmax(router_->forward(x)); // [N, E]
90 + Var gates = ops::topk_renorm(probs, top_k_); // [N, E], rows sum to 1
91 +
92 + Var y = ops::row_scale(experts_[0]->forward(x), gates, 0);
93 + for (int64_t e = 1; e < n_experts_; ++e)
94 + y = ops::add(y, ops::row_scale(experts_[size_t(e)]->forward(x), gates, e));
95 + for (const auto& s : shared_) y = ops::add(y, s->forward(x));
96 +
97 + if (aux_weight_ > 0.0f) {
98 + // Column means of probs via a constant 1/N row — every step is an
99 + // existing autograd op, so the router gets balance gradients.
100 + Var mean_row(Tensor::full({1, N}, 1.0f / float(N)), /*requires_grad=*/false);
101 + Var col_mean = ops::matmul(mean_row, probs); // [1, E]
102 + Var sq = ops::mul(col_mean, col_mean); // [1, E]
103 + Var ones_e(Tensor::full({1, n_experts_}, 1.0f), /*requires_grad=*/false);
104 + Var sum = ops::matmul(sq, ones_e, false, true); // [1, 1]
105 + aux_ = ops::scale(sum.reshaped({1}), aux_weight_ * float(n_experts_));
106 + } else {
107 + aux_ = Var();
108 + }
109 + return y;
110 + }
111 +
112 + Var aux() const override { return aux_; }
113 +
114 +private:
115 + int64_t n_experts_, top_k_;
116 + float aux_weight_;
117 + std::unique_ptr<Linear> router_;
118 + std::vector<std::unique_ptr<MLP>> experts_;
119 + std::vector<std::unique_ptr<MLP>> shared_;
120 + mutable Var aux_; // set by the last forward; consumed by Transformer::loss
121 +};
122 +
44 123 } // namespace forge::nn
modified src/nn/transformer.h +25 −6
@@ -36,8 +36,9 @@ public:
36 36 TransformerBlock(const ModelConfig& cfg, float proj_std, std::mt19937_64& rng)
37 37 : norm1_(std::make_unique<Norm>(cfg)),
38 38 attn_(std::make_unique<CausalSelfAttention>(cfg, proj_std, rng)),
39 − norm2_(std::make_unique<Norm>(cfg)),
40 − mlp_(std::make_unique<MLP>(cfg, proj_std, rng)) {
39 + norm2_(std::make_unique<Norm>(cfg)) {
40 + if (cfg.n_experts > 0) mlp_ = std::make_unique<MoEMLP>(cfg, proj_std, rng);
41 + else mlp_ = std::make_unique<MLP>(cfg, proj_std, rng);
41 42 absorb("norm1", *norm1_);
42 43 absorb("attn", *attn_);
43 44 absorb("norm2", *norm2_);
@@ -51,11 +52,14 @@ public:
51 52 return ops::add(h, m.reshaped({B, T, C}));
52 53 }
53 54
55 + // MoE load-balance loss of the last forward (undefined for dense MLP).
56 + Var moe_aux() const { return mlp_->aux(); }
57 +
54 58 private:
55 59 std::unique_ptr<Norm> norm1_;
56 60 std::unique_ptr<AttentionBase> attn_;
57 61 std::unique_ptr<Norm> norm2_;
58 − std::unique_ptr<MLP> mlp_;
62 + std::unique_ptr<MLPBase> mlp_;
59 63 };
60 64
61 65 // Decoder-only transformer, entirely shaped by ModelConfig.
@@ -93,6 +97,8 @@ public:
93 97 Var forward(const Tensor& ids) const {
94 98 const int64_t B = ids.shape()[0], T = ids.shape()[1];
95 99 Var x = tok_emb_->forward(ids);
100 + if (cfg_.scale_embeddings)
101 + x = ops::scale(x, std::sqrt(float(cfg_.d_model)));
96 102 if (pos_emb_) {
97 103 Tensor pos = Tensor::empty({1, T}, DType::I32);
98 104 for (int64_t t = 0; t < T; ++t) pos.data<int32_t>()[t] = int32_t(t);
@@ -106,14 +112,27 @@ public:
106 112 for (const auto& blk : blocks_) x = blk->forward(x);
107 113 x = final_norm_->forward(x);
108 114 Var x2d = x.reshaped({B * T, cfg_.d_model});
109 − return ops::matmul(x2d, lm_head_weight_, false, true); // [B*T, V]
115 + Var logits = ops::matmul(x2d, lm_head_weight_, false, true); // [B*T, V]
116 + if (cfg_.final_softcap > 0.0f)
117 + logits = ops::softcap(logits, cfg_.final_softcap);
118 + return logits;
110 119 }
111 120
112 − // targets: [B, T] with ignore_index=-1 → scalar mean CE loss
121 + // targets: [B, T] with ignore_index=-1 → scalar mean CE loss. With MoE,
122 + // each block's load-balance term is added (also during eval — its scale
123 + // is ~moe_aux_weight, negligible next to CE but keeps train/val
124 + // comparable).
113 125 Var loss(const Tensor& ids, const Tensor& targets) const {
114 126 Var logits = forward(ids);
115 127 Tensor tflat = targets;
116 − return ops::cross_entropy(logits, tflat.view({targets.numel()}));
128 + Var l = ops::cross_entropy(logits, tflat.view({targets.numel()}));
129 + if (cfg_.n_experts > 0 && cfg_.moe_aux_weight > 0.0f) {
130 + for (const auto& blk : blocks_) {
131 + Var a = blk->moe_aux();
132 + if (a.defined()) l = ops::add(l, a);
133 + }
134 + }
135 + return l;
117 136 }
118 137
119 138 const ModelConfig& config() const { return cfg_; }
modified src/ops/cpu/cpu_ops.cpp +136 −0
@@ -3,10 +3,12 @@
3 3
4 4 #include <dispatch/dispatch.h>
5 5
6 +#include <algorithm>
6 7 #include <cassert>
7 8 #include <cmath>
8 9 #include <cstdio>
9 10 #include <cstdlib>
11 +#include <limits>
10 12
11 13 namespace forge::cpu {
12 14
@@ -179,6 +181,22 @@ void gelu_backward(const Tensor& x, const Tensor& dout, Tensor& dx) {
179 181 }
180 182 }
181 183
184 +void softcap(const Tensor& x, float cap, Tensor& out) {
185 + const float* px = x.data<float>();
186 + float* po = out.data<float>();
187 + for (int64_t i = 0; i < x.numel(); ++i) po[i] = cap * std::tanh(px[i] / cap);
188 +}
189 +
190 +void softcap_backward(const Tensor& y, const Tensor& dout, float cap, Tensor& dx) {
191 + const float* py = y.data<float>();
192 + const float* pd = dout.data<float>();
193 + float* pdx = dx.data<float>();
194 + for (int64_t i = 0; i < y.numel(); ++i) {
195 + const float t = py[i] / cap;
196 + pdx[i] += pd[i] * (1.0f - t * t);
197 + }
198 +}
199 +
182 200 // ---- norms --------------------------------------------------------------------
183 201
184 202 void rmsnorm(const Tensor& x, const Tensor& w, float eps, Tensor& out) {
@@ -319,6 +337,124 @@ void softmax_backward(const Tensor& p, const Tensor& dout, Tensor& dx) {
319 337 }
320 338 }
321 339
340 +// ---- QAT / MoE ------------------------------------------------------------------
341 +
342 +void fake_quant(const Tensor& w, int mode, Tensor& out) {
343 + const int64_t R = w.size(0), C = w.size(1);
344 + const float* pw = w.data<float>();
345 + float* po = out.data<float>();
346 + for (int64_t i = 0; i < R; ++i) {
347 + const float* row = pw + i * C;
348 + float* orow = po + i * C;
349 + if (mode == 0) { // int8: symmetric absmax
350 + float amax = 0.0f;
351 + for (int64_t j = 0; j < C; ++j) amax = std::max(amax, std::fabs(row[j]));
352 + const float s = std::max(amax / 127.0f, 1e-12f);
353 + for (int64_t j = 0; j < C; ++j) orow[j] = std::rint(row[j] / s) * s;
354 + } else { // ternary: BitNet b1.58 absmean
355 + float asum = 0.0f;
356 + for (int64_t j = 0; j < C; ++j) asum += std::fabs(row[j]);
357 + const float s = std::max(asum / float(C), 1e-12f);
358 + for (int64_t j = 0; j < C; ++j)
359 + orow[j] = std::min(1.0f, std::max(-1.0f, std::rint(row[j] / s))) * s;
360 + }
361 + }
362 +}
363 +
364 +namespace {
365 +// Kept set of the k largest entries, ties to the lower index (matches Metal).
366 +void topk_select(const float* row, int64_t E, int64_t k, bool* kept, float* sum) {
367 + for (int64_t j = 0; j < E; ++j) kept[j] = false;
368 + float S = 0.0f;
369 + for (int64_t sel = 0; sel < k; ++sel) {
370 + float best = -std::numeric_limits<float>::max();
371 + int64_t arg = 0;
372 + for (int64_t j = 0; j < E; ++j)
373 + if (!kept[j] && row[j] > best) { best = row[j]; arg = j; }
374 + kept[arg] = true;
375 + S += best;
376 + }
377 + *sum = S;
378 +}
379 +} // namespace
380 +
381 +void topk_renorm(const Tensor& p, int64_t k, Tensor& out) {
382 + const int64_t E = p.shape().back();
383 + const int64_t N = p.numel() / E;
384 + const float* pp = p.data<float>();
385 + float* po = out.data<float>();
386 + bool kbuf[64];
387 + for (int64_t i = 0; i < N; ++i) {
388 + const float* row = pp + i * E;
389 + float S = 0.0f;
390 + topk_select(row, E, k, kbuf, &S);
391 + const float inv = 1.0f / std::max(S, 1e-12f);
392 + for (int64_t j = 0; j < E; ++j) po[i * E + j] = kbuf[j] ? row[j] * inv : 0.0f;
393 + }
394 +}
395 +
396 +void topk_renorm_backward(const Tensor& p, const Tensor& dout, int64_t k, Tensor& dp) {
397 + const int64_t E = p.shape().back();
398 + const int64_t N = p.numel() / E;
399 + const float* pp = p.data<float>();
400 + const float* pd = dout.data<float>();
401 + float* pdp = dp.data<float>();
402 + bool kbuf[64];
403 + for (int64_t i = 0; i < N; ++i) {
404 + const float* row = pp + i * E;
405 + const float* drow = pd + i * E;
406 + float S = 0.0f;
407 + topk_select(row, E, k, kbuf, &S);
408 + const float inv = 1.0f / std::max(S, 1e-12f);
409 + float dot = 0.0f;
410 + for (int64_t j = 0; j < E; ++j)
411 + if (kbuf[j]) dot += drow[j] * row[j] * inv;
412 + for (int64_t j = 0; j < E; ++j)
413 + if (kbuf[j]) pdp[i * E + j] += (drow[j] - dot) * inv;
414 + }
415 +}
416 +
417 +void row_scale(const Tensor& x, const Tensor& gates, int64_t e, Tensor& out) {
418 + const int64_t C = x.shape().back();
419 + const int64_t N = x.numel() / C;
420 + const int64_t E = gates.shape().back();
421 + const float* px = x.data<float>();
422 + const float* pg = gates.data<float>();
423 + float* po = out.data<float>();
424 + for (int64_t i = 0; i < N; ++i) {
425 + const float s = pg[i * E + e];
426 + for (int64_t j = 0; j < C; ++j) po[i * C + j] = px[i * C + j] * s;
427 + }
428 +}
429 +
430 +void row_scale_accumulate(const Tensor& x, const Tensor& gates, int64_t e, Tensor& dst) {
431 + const int64_t C = x.shape().back();
432 + const int64_t N = x.numel() / C;
433 + const int64_t E = gates.shape().back();
434 + const float* px = x.data<float>();
435 + const float* pg = gates.data<float>();
436 + float* pd = dst.data<float>();
437 + for (int64_t i = 0; i < N; ++i) {
438 + const float s = pg[i * E + e];
439 + for (int64_t j = 0; j < C; ++j) pd[i * C + j] += px[i * C + j] * s;
440 + }
441 +}
442 +
443 +void row_scale_gate_backward(const Tensor& dout, const Tensor& x, int64_t e,
444 + Tensor& dgates) {
445 + const int64_t C = x.shape().back();
446 + const int64_t N = x.numel() / C;
447 + const int64_t E = dgates.shape().back();
448 + const float* pd = dout.data<float>();
449 + const float* px = x.data<float>();
450 + float* pg = dgates.data<float>();
451 + for (int64_t i = 0; i < N; ++i) {
452 + float acc = 0.0f;
453 + for (int64_t j = 0; j < C; ++j) acc += pd[i * C + j] * px[i * C + j];
454 + pg[i * E + e] += acc;
455 + }
456 +}
457 +
322 458 // ---- embedding ------------------------------------------------------------------
323 459
324 460 void embedding(const Tensor& weight, const Tensor& ids, Tensor& out) {
modified src/ops/cpu/cpu_ops.h +18 −0
@@ -36,6 +36,9 @@ void silu(const Tensor& x, Tensor& out);
36 36 void silu_backward(const Tensor& x, const Tensor& dout, Tensor& dx);
37 37 void gelu(const Tensor& x, Tensor& out); // tanh approximation (GPT-2)
38 38 void gelu_backward(const Tensor& x, const Tensor& dout, Tensor& dx);
39 +void softcap(const Tensor& x, float cap, Tensor& out); // cap*tanh(x/cap)
40 +// dx += dout * (1 - (y/cap)^2) — takes the forward OUTPUT y
41 +void softcap_backward(const Tensor& y, const Tensor& dout, float cap, Tensor& dx);
39 42
40 43 // ---- norms (row-wise over the last dim) -----------------------------------
41 44 // x: [N, C], w: [C] (b: [C] for layernorm), out: [N, C]
@@ -78,6 +81,21 @@ void attention_backward(const Tensor& q, const Tensor& k, const Tensor& v,
78 81 int64_t n_heads, int64_t n_kv_heads, float scale,
79 82 Tensor& dq, Tensor& dk, Tensor& dv);
80 83
84 +// ---- QAT / MoE --------------------------------------------------------------
85 +// Per-row fake quantization of a 2-D weight; mode 0 = int8 (absmax/127),
86 +// mode 1 = ternary (absmean, BitNet-style). Backward is STE — identity.
87 +void fake_quant(const Tensor& w, int mode, Tensor& out);
88 +// Keep top-k per row of p [N,E], renormalize kept entries to sum 1, zero the
89 +// rest (ties broken by lower index, same as the Metal kernel).
90 +void topk_renorm(const Tensor& p, int64_t k, Tensor& out);
91 +void topk_renorm_backward(const Tensor& p, const Tensor& dout, int64_t k, Tensor& dp);
92 +// out[i,:] = x[i,:] * gates[i,e]; _accumulate does dst += (dx path);
93 +// gate_backward does dgates[i,e] += dot(dout[i,:], x[i,:]).
94 +void row_scale(const Tensor& x, const Tensor& gates, int64_t e, Tensor& out);
95 +void row_scale_accumulate(const Tensor& x, const Tensor& gates, int64_t e, Tensor& dst);
96 +void row_scale_gate_backward(const Tensor& dout, const Tensor& x, int64_t e,
97 + Tensor& dgates);
98 +
81 99 // ---- fused softmax cross-entropy -------------------------------------------
82 100 // logits: [N, V]; targets: [N] i32 (ignore_index = -1). Returns mean loss
83 101 // over valid rows. dlogits (if non-null) receives (softmax − onehot)/n_valid
modified src/ops/metal/metal_ops.cpp +60 −0
@@ -218,6 +218,14 @@ void gelu(const Tensor& x, Tensor& out) {
218 218 encode_flat("gelu_f32", {&x, &out}, nullptr, 0, x.numel());
219 219 }
220 220
221 +void softcap(const Tensor& x, float cap, Tensor& out) {
222 + encode_flat("softcap_f32", {&x, &out}, &cap, sizeof(cap), x.numel());
223 +}
224 +
225 +void softcap_backward(const Tensor& y, const Tensor& dout, float cap, Tensor& dx) {
226 + encode_flat("softcap_bwd_f32", {&y, &dout, &dx}, &cap, sizeof(cap), y.numel());
227 +}
228 +
221 229 // ---- row reductions --------------------------------------------------------------
222 230
223 231 void softmax(const Tensor& x, Tensor& out) {
@@ -269,6 +277,58 @@ void layernorm(const Tensor& x, const Tensor& w, const Tensor& b, float eps, Ten
269 277 }
270 278
271 279
280 +// ---- QAT / MoE -------------------------------------------------------------------
281 +
282 +void fake_quant(const Tensor& w, int mode, Tensor& out) {
283 + check(w.ndim() == 2 && w.numel() == out.numel(), "fake_quant: bad shapes");
284 + const uint32_t p[2] = {uint32_t(w.size(1)), uint32_t(mode)};
285 + encode_flat("fake_quant_f32", {&w, &out}, p, sizeof(p), w.size(0));
286 +}
287 +
288 +void softmax_backward(const Tensor& p, const Tensor& dout, Tensor& dx) {
289 + const int64_t C = p.shape().back();
290 + const uint32_t c32 = uint32_t(C);
291 + encode_flat("softmax_bwd_f32", {&p, &dout, &dx}, &c32, sizeof(c32), p.numel() / C);
292 +}
293 +
294 +void topk_renorm(const Tensor& p, int64_t k, Tensor& out) {
295 + const int64_t E = p.shape().back();
296 + check(E <= 64, "topk_renorm: E must be <= 64 (kernel's kept[] bound)");
297 + const uint32_t pr[2] = {uint32_t(E), uint32_t(k)};
298 + encode_flat("topk_renorm_f32", {&p, &out}, pr, sizeof(pr), p.numel() / E);
299 +}
300 +
301 +void topk_renorm_backward(const Tensor& p, const Tensor& dout, int64_t k, Tensor& dp) {
302 + const int64_t E = p.shape().back();
303 + check(E <= 64, "topk_renorm_backward: E must be <= 64");
304 + const uint32_t pr[2] = {uint32_t(E), uint32_t(k)};
305 + encode_flat("topk_renorm_bwd_f32", {&p, &dout, &dp}, pr, sizeof(pr), p.numel() / E);
306 +}
307 +
308 +namespace {
309 +struct RowScaleParams { uint32_t C, E, e; };
310 +} // namespace
311 +
312 +void row_scale(const Tensor& x, const Tensor& gates, int64_t e, Tensor& out) {
313 + const RowScaleParams p{uint32_t(x.shape().back()), uint32_t(gates.shape().back()),
314 + uint32_t(e)};
315 + encode_flat("row_scale_f32", {&x, &gates, &out}, &p, sizeof(p), x.numel());
316 +}
317 +
318 +void row_scale_accumulate(const Tensor& x, const Tensor& gates, int64_t e, Tensor& dst) {
319 + const RowScaleParams p{uint32_t(x.shape().back()), uint32_t(gates.shape().back()),
320 + uint32_t(e)};
321 + encode_flat("row_scale_acc_f32", {&x, &gates, &dst}, &p, sizeof(p), x.numel());
322 +}
323 +
324 +void row_scale_gate_backward(const Tensor& dout, const Tensor& x, int64_t e,
325 + Tensor& dgates) {
326 + const RowScaleParams p{uint32_t(x.shape().back()), uint32_t(dgates.shape().back()),
327 + uint32_t(e)};
328 + encode_flat("row_scale_gate_bwd_f32", {&dout, &x, &dgates}, &p, sizeof(p),
329 + x.numel() / x.shape().back());
330 +}
331 +
272 332 // ---- backward / training ops ---------------------------------------------------
273 333
274 334 void accumulate(Tensor& dst, const Tensor& src) {
modified src/ops/metal/metal_ops.h +18 −0
@@ -88,6 +88,8 @@ void scale(const Tensor& a, float s, Tensor& out);
88 88 void add_bias(const Tensor& x, const Tensor& bias, Tensor& out);
89 89 void silu(const Tensor& x, Tensor& out);
90 90 void gelu(const Tensor& x, Tensor& out);
91 +void softcap(const Tensor& x, float cap, Tensor& out); // cap*tanh(x/cap)
92 +void softcap_backward(const Tensor& y, const Tensor& dout, float cap, Tensor& dx);
91 93
92 94 void softmax(const Tensor& x, Tensor& out); // rows = last dim
93 95 void rmsnorm(const Tensor& x, const Tensor& w, float eps, Tensor& out);
@@ -140,6 +142,22 @@ void flash_attention_backward(const Tensor& q, const Tensor& k, const Tensor& v,
140 142 void cross_entropy(const Tensor& logits, const Tensor& targets, int64_t n_valid,
141 143 Tensor& losses, Tensor& loss_out, Tensor* dlogits);
142 144
145 +// ---- QAT / MoE --------------------------------------------------------------
146 +// Per-row fake quantization of a 2-D weight; mode 0 = int8 (absmax/127),
147 +// mode 1 = ternary (absmean, BitNet-style). Backward is STE — no kernel.
148 +void fake_quant(const Tensor& w, int mode, Tensor& out);
149 +// dx += p ∘ (dout − dot(dout, p)) per row; thread-per-row, small last dims.
150 +void softmax_backward(const Tensor& p, const Tensor& dout, Tensor& dx);
151 +// Keep top-k per row of p [N,E], renormalize kept to sum 1, zero the rest.
152 +void topk_renorm(const Tensor& p, int64_t k, Tensor& out);
153 +void topk_renorm_backward(const Tensor& p, const Tensor& dout, int64_t k, Tensor& dp);
154 +// out[i,:] = x[i,:] * gates[i,e]; the _accumulate variant does dst += (dx path).
155 +void row_scale(const Tensor& x, const Tensor& gates, int64_t e, Tensor& out);
156 +void row_scale_accumulate(const Tensor& x, const Tensor& gates, int64_t e, Tensor& dst);
157 +// dgates[i,e] += dot(dout[i,:], x[i,:])
158 +void row_scale_gate_backward(const Tensor& dout, const Tensor& x, int64_t e,
159 + Tensor& dgates);
160 +
143 161 // Fused optimizer update for one tensor; all state f32.
144 162 void adamw_step(Tensor& w, const Tensor& g, Tensor& m, Tensor& v,
145 163 float lr, float beta1, float beta2, int64_t t, float eps, float wd,
modified src/ops/ops.cpp +96 −0
@@ -208,6 +208,102 @@ Var gelu(const Var& x) {
208 208 return result;
209 209 }
210 210
211 +Var softcap(const Var& x, float cap) {
212 + Tensor out = Tensor::empty(x.value().shape());
213 + if (gpu()) metal::softcap(x.value(), cap, out);
214 + else cpu::softcap(x.value(), cap, out);
215 +
216 + const bool needs = grad_needed({&x});
217 + Var result(std::move(out), needs);
218 + if (needs) {
219 + Tape::get().record([x, cap, result]() {
220 + if (!x.requires_grad()) return;
221 + if (gpu()) metal::softcap_backward(result.value(), result.grad(), cap, x.grad());
222 + else cpu::softcap_backward(result.value(), result.grad(), cap, x.grad());
223 + });
224 + }
225 + return result;
226 +}
227 +
228 +Var softmax(const Var& x) {
229 + Tensor out = Tensor::empty(x.value().shape());
230 + if (gpu()) metal::softmax(x.value(), out);
231 + else cpu::softmax(x.value(), out);
232 +
233 + const bool needs = grad_needed({&x});
234 + Var result(std::move(out), needs);
235 + if (needs) {
236 + Tape::get().record([x, result]() {
237 + if (!x.requires_grad()) return;
238 + if (gpu()) metal::softmax_backward(result.value(), result.grad(), x.grad());
239 + else cpu::softmax_backward(result.value(), result.grad(), x.grad());
240 + });
241 + }
242 + return result;
243 +}
244 +
245 +Var fake_quant(const Var& w, QuantMode mode) {
246 + if (mode == QuantMode::None) return w;
247 + const int m = mode == QuantMode::Int8 ? 0 : 1;
248 + Tensor out = Tensor::empty(w.value().shape());
249 + if (gpu()) metal::fake_quant(w.value(), m, out);
250 + else cpu::fake_quant(w.value(), m, out);
251 +
252 + const bool needs = grad_needed({&w});
253 + Var result(std::move(out), needs);
254 + if (needs) {
255 + // Straight-through estimator: d(quant(w))/dw ≈ I.
256 + Tape::get().record([w, result]() {
257 + if (w.requires_grad()) accumulate(w.grad(), result.grad());
258 + });
259 + }
260 + return result;
261 +}
262 +
263 +Var topk_renorm(const Var& probs, int64_t k) {
264 + Tensor out = Tensor::empty(probs.value().shape());
265 + if (gpu()) metal::topk_renorm(probs.value(), k, out);
266 + else cpu::topk_renorm(probs.value(), k, out);
267 +
268 + const bool needs = grad_needed({&probs});
269 + Var result(std::move(out), needs);
270 + if (needs) {
271 + Tape::get().record([probs, k, result]() {
272 + if (!probs.requires_grad()) return;
273 + if (gpu())
274 + metal::topk_renorm_backward(probs.value(), result.grad(), k, probs.grad());
275 + else
276 + cpu::topk_renorm_backward(probs.value(), result.grad(), k, probs.grad());
277 + });
278 + }
279 + return result;
280 +}
281 +
282 +Var row_scale(const Var& x, const Var& gates, int64_t e) {
283 + Tensor out = Tensor::empty(x.value().shape());
284 + if (gpu()) metal::row_scale(x.value(), gates.value(), e, out);
285 + else cpu::row_scale(x.value(), gates.value(), e, out);
286 +
287 + const bool needs = grad_needed({&x, &gates});
288 + Var result(std::move(out), needs);
289 + if (needs) {
290 + Tape::get().record([x, gates, e, result]() {
291 + const Tensor& dout = result.grad();
292 + if (x.requires_grad()) {
293 + if (gpu()) metal::row_scale_accumulate(dout, gates.value(), e, x.grad());
294 + else cpu::row_scale_accumulate(dout, gates.value(), e, x.grad());
295 + }
296 + if (gates.requires_grad()) {
297 + if (gpu())
298 + metal::row_scale_gate_backward(dout, x.value(), e, gates.grad());
299 + else
300 + cpu::row_scale_gate_backward(dout, x.value(), e, gates.grad());
301 + }
302 + });
303 + }
304 + return result;
305 +}
306 +
211 307 Var rmsnorm(const Var& x, const Var& w, float eps) {
212 308 Tensor out = Tensor::empty(x.value().shape());
213 309 if (gpu()) metal::rmsnorm(x.value(), w.value(), eps, out);
modified src/ops/ops.h +19 −0
@@ -32,6 +32,25 @@ Var scale(const Var& a, float s);
32 32 Var silu(const Var& x);
33 33 Var gelu(const Var& x);
34 34
35 +// Gemma-style soft capping: cap * tanh(x / cap). Bounds logits smoothly.
36 +Var softcap(const Var& x, float cap);
37 +
38 +// Row softmax over the last dim (autograd; the attention op has its own
39 +// fused softmax — this one is for the MoE router's small rows).
40 +Var softmax(const Var& x);
41 +
42 +// QAT: per-row fake quantization of a 2-D weight. Forward quantizes
43 +// (int8 absmax / ternary absmean); backward is the straight-through
44 +// estimator — gradients flow unchanged into the f32 master weight.
45 +enum class QuantMode { None, Int8, Ternary };
46 +Var fake_quant(const Var& w, QuantMode mode);
47 +
48 +// MoE gating: keep top-k per row, renormalize kept entries to sum 1, zero
49 +// the rest. Gradients flow through kept entries only.
50 +Var topk_renorm(const Var& probs, int64_t k);
51 +// y[i,:] = x[i,:] * gates[i,e] — weight expert e's output by its gate column.
52 +Var row_scale(const Var& x, const Var& gates, int64_t e);
53 +
35 54 Var rmsnorm(const Var& x, const Var& w, float eps);
36 55 Var layernorm(const Var& x, const Var& w, const Var& b, float eps);
37 56
modified tests/test_ops.cpp +99 −0
@@ -264,6 +264,78 @@ void test_metal_rowops() {
264 264 }
265 265 }
266 266
267 +void test_metal_moe_quant_ops() {
268 + std::printf("metal QAT/MoE op parity\n");
269 + std::mt19937 rng(4242);
270 +
271 + // fake_quant: both modes, odd and wide rows
272 + for (int mode : {0, 1}) {
273 + forge::Tensor w = forge::Tensor::empty({33, 257});
274 + fill_random(w, rng, -0.2f, 0.2f);
275 + forge::Tensor ref = forge::Tensor::empty(w.shape());
276 + forge::Tensor gpu = forge::Tensor::empty(w.shape());
277 + forge::cpu::fake_quant(w, mode, ref);
278 + forge::metal::fake_quant(w, mode, gpu);
279 + forge::metal::sync();
280 + expect_close(gpu, ref, mode == 0 ? "fake_quant int8" : "fake_quant ternary");
281 + }
282 +
283 + // topk_renorm forward + backward, row_scale family
284 + const int64_t N = 130, E = 8, C = 37, K = 2;
285 + forge::Tensor probs = forge::Tensor::empty({N, E});
286 + fill_random(probs, rng, 0.01f, 1.0f);
287 + forge::Tensor ref = forge::Tensor::empty({N, E});
288 + forge::Tensor gpu = forge::Tensor::empty({N, E});
289 + forge::cpu::topk_renorm(probs, K, ref);
290 + forge::metal::topk_renorm(probs, K, gpu);
291 + forge::metal::sync();
292 + expect_close(gpu, ref, "topk_renorm fwd");
293 +
294 + forge::Tensor dout = forge::Tensor::empty({N, E});
295 + fill_random(dout, rng);
296 + forge::Tensor dref = forge::Tensor::zeros({N, E});
297 + forge::Tensor dgpu = forge::Tensor::zeros({N, E});
298 + forge::cpu::topk_renorm_backward(probs, dout, K, dref);
299 + forge::metal::topk_renorm_backward(probs, dout, K, dgpu);
300 + forge::metal::sync();
301 + expect_close(dgpu, dref, "topk_renorm bwd");
302 +
303 + forge::Tensor x = forge::Tensor::empty({N, C});
304 + fill_random(x, rng);
305 + forge::Tensor yref = forge::Tensor::empty({N, C});
306 + forge::Tensor ygpu = forge::Tensor::empty({N, C});
307 + forge::cpu::row_scale(x, ref, 3, yref);
308 + forge::metal::row_scale(x, ref, 3, ygpu);
309 + forge::metal::sync();
310 + expect_close(ygpu, yref, "row_scale fwd");
311 +
312 + forge::Tensor aref = forge::Tensor::zeros({N, C});
313 + forge::Tensor agpu = forge::Tensor::zeros({N, C});
314 + forge::cpu::row_scale_accumulate(x, ref, 3, aref);
315 + forge::metal::row_scale_accumulate(x, ref, 3, agpu);
316 + forge::metal::sync();
317 + expect_close(agpu, aref, "row_scale accumulate");
318 +
319 + forge::Tensor dxo = forge::Tensor::empty({N, C});
320 + fill_random(dxo, rng);
321 + forge::Tensor gref = forge::Tensor::zeros({N, E});
322 + forge::Tensor ggpu = forge::Tensor::zeros({N, E});
323 + forge::cpu::row_scale_gate_backward(dxo, x, 3, gref);
324 + forge::metal::row_scale_gate_backward(dxo, x, 3, ggpu);
325 + forge::metal::sync();
326 + expect_close(ggpu, gref, "row_scale gate bwd");
327 +
328 + // generic softmax backward (router-sized rows)
329 + forge::Tensor sm = forge::Tensor::empty({N, E});
330 + forge::cpu::softmax(probs, sm);
331 + forge::Tensor sref = forge::Tensor::zeros({N, E});
332 + forge::Tensor sgpu = forge::Tensor::zeros({N, E});
333 + forge::cpu::softmax_backward(sm, dout, sref);
334 + forge::metal::softmax_backward(sm, dout, sgpu);
335 + forge::metal::sync();
336 + expect_close(sgpu, sref, "softmax bwd");
337 +}
338 +
267 339 void test_batched_encoding() {
268 340 std::printf("batched encoding (many dispatches, one sync)\n");
269 341 std::mt19937 rng(55);
@@ -502,7 +574,34 @@ int main() {
502 574 cfg.norm = "layernorm"; cfg.activation = "gelu"; cfg.use_rope = false;
503 575 cfg.n_kv_heads = 2; cfg.tied_embeddings = false;
504 576 test_backend_parity_model("layernorm/gelu/pos-emb/mha", cfg);
577 +
578 + // QAT: full model with fake-quantized linears (STE backward)
579 + forge::ModelConfig qcfg;
580 + qcfg.n_layers = 2; qcfg.d_model = 16; qcfg.n_heads = 2; qcfg.n_kv_heads = 1;
581 + qcfg.d_ff = 24; qcfg.vocab_size = 11; qcfg.context_length = 8;
582 + qcfg.tied_embeddings = true;
583 + qcfg.quant = "ternary";
584 + test_backend_parity_model("QAT ternary linears", qcfg);
585 + qcfg.quant = "int8";
586 + test_backend_parity_model("QAT int8 linears", qcfg);
587 +
588 + // MoE: router + top-2 of 4 experts + 1 shared expert + aux loss
589 + forge::ModelConfig mcfg;
590 + mcfg.n_layers = 2; mcfg.d_model = 16; mcfg.n_heads = 2; mcfg.n_kv_heads = 1;
591 + mcfg.d_ff = 24; mcfg.vocab_size = 11; mcfg.context_length = 8;
592 + mcfg.tied_embeddings = true;
593 + mcfg.n_experts = 4; mcfg.moe_top_k = 2; mcfg.n_shared_experts = 1;
594 + test_backend_parity_model("MoE 4+1shared top-2 + aux", mcfg);
595 +
596 + // Architecture-variant knobs: QK-norm + logit softcap + embed scaling
597 + forge::ModelConfig vcfg;
598 + vcfg.n_layers = 2; vcfg.d_model = 16; vcfg.n_heads = 2; vcfg.n_kv_heads = 1;
599 + vcfg.d_ff = 24; vcfg.vocab_size = 11; vcfg.context_length = 8;
600 + vcfg.tied_embeddings = true;
601 + vcfg.qk_norm = true; vcfg.final_softcap = 30.0f; vcfg.scale_embeddings = true;
602 + test_backend_parity_model("qk-norm + softcap + embed-scale", vcfg);
505 603 }
604 + test_metal_moe_quant_ops();
506 605 test_flash_attention();
507 606 test_adamw_kernel();
508 607
509 608