spb/forge Public MIT

1// Author: Simon-Pierre Boucher — contact@spboucher.ai2#pragma once34#include "nn/config.h"5#include "nn/linear.h"6#include "ops/cpu/cpu_ops.h"7#include "ops/metal/metal_ops.h"8#include "ops/ops.h"910#include <memory>11#include <string>12#include <vector>1314namespace forge::nn {1516// Interface so the MoE variant slots into TransformerBlock without branching17// at every call site. aux() is the block's auxiliary loss ([1] Var, undefined18// when the variant has none).19class MLPBase : public Module {20public:21    virtual ~MLPBase() = default;22    // x: [N, C] 2-D23    virtual Var forward(const Var& x) const = 0;24    virtual Var aux() const { return Var(); }25    // noaux load-balance bias step (V3); no-op for dense/softmax variants.26    virtual void bias_update(float /*gamma*/) {}27};2829// SwiGLU: w2( silu(x w1) ⊙ (x w3) )  |  GELU/ReLU²: w2( act(x w1) )30class MLP : public MLPBase {31public:32    MLP(const ModelConfig& cfg, float proj_std, std::mt19937_64& rng)33        : swiglu_(cfg.activation == "swiglu"), relu2_(cfg.activation == "relu2") {34        const float std = 0.02f;35        const int64_t C = cfg.d_model, F = cfg.d_ff;36        w1_ = std::make_unique<Linear>(C, F, false, std, rng);37        if (swiglu_) {38            w3_ = std::make_unique<Linear>(C, F, false, std, rng);39        }40        w2_ = std::make_unique<Linear>(F, C, false, proj_std, rng); // residual projection41        const ops::QuantMode qm = quant_mode_from(cfg.quant);42        w1_->set_quant(qm);43        if (w3_) w3_->set_quant(qm);44        w2_->set_quant(qm);45        absorb("w1", *w1_);46        if (w3_) absorb("w3", *w3_);47        absorb("w2", *w2_);48    }4950    Var forward(const Var& x) const override {51        if (swiglu_) {52            Var gate = ops::silu(w1_->forward(x));53            Var up = w3_->forward(x);54            return w2_->forward(ops::mul(gate, up));55        }56        Var h = w1_->forward(x);57        return w2_->forward(relu2_ ? ops::relu2(h) : ops::gelu(h));58    }5960private:61    bool swiglu_, relu2_;62    std::unique_ptr<Linear> w1_, w2_, w3_;63};6465// Mixture-of-experts MLP. Router scoring is softmax (classic) or sigmoid66// (DeepSeek V3/K2); top-k selection may be steered by a per-expert balance67// bias (aux-loss-free balancing: bias_update() nudges it ±gamma by observed68// load after each optimizer step), gates are renormalized or raw, and the69// routed sum can be rescaled (routed_scaling_factor). Experts may use their70// own d_ff (moe_d_ff). v1 computes EVERY expert densely and weights by the71// (mostly zero) gates — correctness first; token gather/scatter sparsity is72// a later optimization. The optional differentiable aux loss73// E · Σ_e (mean_i score)² remains available (moe_aux_weight).74class MoEMLP : public MLPBase {75public:76    MoEMLP(const ModelConfig& cfg, float proj_std, std::mt19937_64& rng)77        : n_experts_(cfg.n_experts), top_k_(cfg.moe_top_k),78          aux_weight_(cfg.moe_aux_weight),79          sigmoid_(cfg.moe_scoring == "sigmoid"),80          norm_topk_(cfg.moe_norm_topk),81          routed_scale_(cfg.routed_scaling_factor),82          track_load_(cfg.moe_bias_gamma > 0.0f) {83        router_ = std::make_unique<Linear>(cfg.d_model, n_experts_, false, 0.02f, rng);84        absorb("router", *router_);85        ModelConfig ecfg = cfg;86        if (cfg.moe_d_ff > 0) ecfg.d_ff = cfg.moe_d_ff;87        experts_.reserve(size_t(n_experts_));88        for (int64_t e = 0; e < n_experts_; ++e) {89            experts_.push_back(std::make_unique<MLP>(ecfg, proj_std, rng));90            absorb("experts." + std::to_string(e), *experts_.back());91        }92        // Always-active experts, added ungated (DeepSeek/Kimi style): they93        // absorb common knowledge so routed experts can specialize.94        for (int64_t s = 0; s < cfg.n_shared_experts; ++s) {95            shared_.push_back(std::make_unique<MLP>(ecfg, proj_std, rng));96            absorb("shared." + std::to_string(s), *shared_.back());97        }98        bias_ = Tensor::zeros({n_experts_});99        counts_ = Tensor::zeros({n_experts_});100        // Registered (grad-free) so checkpoints and .forge snapshots carry101        // the balance state; the optimizer skips non-grad params.102        register_param("router.balance_bias", Var(bias_, /*requires_grad=*/false));103    }104105    Var forward(const Var& x) const override {106        const int64_t N = x.value().size(0);107        Var logits = router_->forward(x); // [N, E]108        Var probs = sigmoid_ ? ops::sigmoid(logits) : ops::softmax(logits);109        Var gates = ops::topk_renorm(probs, bias_, top_k_, norm_topk_);110111        Var y = ops::row_scale(experts_[0]->forward(x), gates, 0);112        for (int64_t e = 1; e < n_experts_; ++e)113            y = ops::add(y, ops::row_scale(experts_[size_t(e)]->forward(x), gates, e));114        if (routed_scale_ != 1.0f) y = ops::scale(y, routed_scale_);115        for (const auto& s : shared_) y = ops::add(y, s->forward(x));116117        if (track_load_) {118            if (ops::backend() == ops::Backend::Metal)119                metal::expert_counts(gates.value(), counts_);120            else121                cpu::expert_counts(gates.value(), counts_);122        }123124        if (aux_weight_ > 0.0f) {125            // Column means of probs via a constant 1/N row — every step is an126            // existing autograd op, so the router gets balance gradients.127            Var mean_row(Tensor::full({1, N}, 1.0f / float(N)), /*requires_grad=*/false);128            Var col_mean = ops::matmul(mean_row, probs);            // [1, E]129            Var sq = ops::mul(col_mean, col_mean);                  // [1, E]130            Var ones_e(Tensor::full({1, n_experts_}, 1.0f), /*requires_grad=*/false);131            Var sum = ops::matmul(sq, ones_e, false, true);         // [1, 1]132            aux_ = ops::scale(sum.reshaped({1}), aux_weight_ * float(n_experts_));133        } else {134            aux_ = Var();135        }136        return y;137    }138139    Var aux() const override { return aux_; }140141    // V3 noaux rule: overloaded experts (above mean load) lose selection142    // bias, underloaded gain. Called by the trainer AFTER the optimizer's143    // final sync — CPU reads/writes are safe there.144    void bias_update(float gamma) override {145        float* c = counts_.data<float>();146        float* b = bias_.data<float>();147        double total = 0.0;148        for (int64_t e = 0; e < n_experts_; ++e) total += c[e];149        const float mean = float(total / double(n_experts_));150        for (int64_t e = 0; e < n_experts_; ++e) {151            if (c[e] > mean) b[e] -= gamma;152            else if (c[e] < mean) b[e] += gamma;153            c[e] = 0.0f;154        }155    }156157private:158    int64_t n_experts_, top_k_;159    float aux_weight_;160    bool sigmoid_, norm_topk_;161    float routed_scale_;162    bool track_load_;163    std::unique_ptr<Linear> router_;164    std::vector<std::unique_ptr<MLP>> experts_;165    std::vector<std::unique_ptr<MLP>> shared_;166    Tensor bias_;           // [E] selection bias (noaux balancing state)167    mutable Tensor counts_; // [E] load accumulator between bias updates168    mutable Var aux_; // set by the last forward; consumed by Transformer::loss169};170171} // namespace forge::nn172