spb/forge Public MIT

1// Author: Simon-Pierre Boucher — contact@spboucher.ai2#pragma once34#include <algorithm>5#include <cstdint>6#include <fstream>7#include <stdexcept>8#include <string>910#include <nlohmann/json.hpp>1112namespace forge {1314// The entire architecture comes from here — no hardcoded model sizes15// anywhere else in the codebase.16struct ModelConfig {17    std::string name = "model";18    int64_t n_layers = 6;19    int64_t d_model = 384;20    int64_t n_heads = 6;21    int64_t n_kv_heads = 6;      // < n_heads => GQA22    int64_t d_ff = 1024;23    int64_t vocab_size = 4096;24    int64_t context_length = 512;25    bool tied_embeddings = true;26    bool use_rope = true;        // false => learned positional embeddings27    float rope_theta = 10000.0f;28    std::string norm = "rmsnorm";     // "rmsnorm" | "layernorm"29    float norm_eps = 1e-6f;30    std::string activation = "swiglu"; // "swiglu" | "gelu"31    float dropout = 0.0f;32    // Architecture-variant knobs (all off by default = the classic recipe):33    bool qk_norm = false;         // RMSNorm on q/k per head before RoPE (Qwen3/Gemma3)34    float final_softcap = 0.0f;   // logits = cap*tanh(logits/cap) (Gemma2: 30)35    bool scale_embeddings = false;// x *= sqrt(d_model) after embedding (Gemma)36    bool attention_bias = false;  // bias on q/k/v projections only (Qwen2.5)37    int64_t head_dim_override = 0;// decouple head_dim from d_model/n_heads (Qwen3)38    int64_t nope_every = 0;       // skip RoPE every Nth layer (SmolLM3: 4)39    std::string norm_placement = "pre"; // "pre" | "post" (OLMo2) | "sandwich" (Gemma)40    // RoPE scaling, HF "llama3" style (0 = off). Applied to the inv-freq table.41    float rope_scale_factor = 0.0f;      // llama3: 3242    float rope_scale_low = 1.0f;         // low_freq_factor43    float rope_scale_high = 4.0f;        // high_freq_factor44    int64_t rope_scale_orig_ctx = 8192;  // original_max_position_embeddings45    // Sliding-window attention (Mistral/Gemma3). 0 = full attention everywhere.46    int64_t sliding_window = 0;47    int64_t sliding_global_every = 0; // every Nth layer is global (Gemma3: 6); 0 = none48    float rope_theta_global = 0.0f;   // theta for global layers (Gemma3: 1e6); 0 = theta49    float attn_softcap = 0.0f;        // cap on attention scores (Gemma2: 50)50    // Quantization-aware training: linear weights are fake-quantized each51    // forward (per-row scales, straight-through estimator in backward).52    // Master weights and the optimizer stay f32.53    std::string quant = "none";       // "none" | "int8" | "ternary"54    // Mixture-of-experts MLP: n_experts > 0 replaces each block's MLP with55    // `n_experts` expert MLPs (d_ff each) + a top-k softmax router. v156    // computes every expert densely and masks — correctness first.57    int64_t n_experts = 0;            // 0 => dense MLP58    int64_t moe_top_k = 2;59    float moe_aux_weight = 0.01f;     // load-balance loss: E * sum_e mean_gate_e^260    int64_t n_shared_experts = 0;     // always-active experts (DeepSeek/Kimi style)61    // DeepSeek-V3-style routing refinements:62    std::string moe_scoring = "softmax"; // "softmax" | "sigmoid" (V3/K2)63    bool moe_norm_topk = true;        // renormalize kept gates to sum 164    float routed_scaling_factor = 1.0f; // y = shared + factor * sum(gated) (V3: 2.5)65    int64_t moe_d_ff = 0;             // per-expert d_ff; 0 = d_ff (V3: 2048 vs 18432)66    int64_t first_k_dense = 0;        // first k layers keep a dense MLP (V3: 3, K2: 1)67    // Aux-loss-FREE balancing (V3 "noaux"): select top-k on s+b but gate on s;68    // after each step b += gamma for underloaded experts, -= gamma for69    // overloaded. 0 disables (bias stays zero).70    float moe_bias_gamma = 0.0f;      // V3: 0.0017172    int64_t head_dim() const {73        return head_dim_override > 0 ? head_dim_override : d_model / n_heads;74    }7576    // True when layer i (0-based) attends globally; only meaningful with77    // sliding_window > 0.78    bool layer_is_global(int64_t i) const {79        if (sliding_window <= 0) return true;80        return sliding_global_every > 0 && (i + 1) % sliding_global_every == 0;81    }82    bool layer_uses_rope(int64_t i) const {83        return use_rope && !(nope_every > 0 && (i + 1) % nope_every == 0);84    }8586    // Total parameter count (SwiGLU MLP has three mats, GELU/ReLU² have two).87    int64_t num_params() const {88        const int64_t hd = head_dim();89        int64_t attn = d_model * n_heads * hd             // wq90                     + 2 * d_model * n_kv_heads * hd      // wk, wv91                     + n_heads * hd * d_model;            // wo92        if (attention_bias) attn += (n_heads + 2 * n_kv_heads) * hd;93        const int64_t act_mats = (activation == "swiglu") ? 3 : 2;94        const int64_t mlp_dense = act_mats * d_model * d_ff;95        const int64_t expert_dff = (n_experts > 0 && moe_d_ff > 0) ? moe_d_ff : d_ff;96        const int64_t mlp_moe = (n_experts + n_shared_experts) * act_mats * d_model *97                                    expert_dff +98                                n_experts * d_model;       // experts + router99        // first_k_dense layers keep the dense MLP (V3-style).100        const int64_t n_moe_layers =101            n_experts > 0 ? n_layers - std::min(first_k_dense, n_layers) : 0;102        const int64_t mlp_total =103            n_moe_layers * mlp_moe + (n_layers - n_moe_layers) * mlp_dense;104        const int64_t norms_per_layer = norm_placement == "sandwich" ? 4 : 2;105        int64_t norms = (norm == "layernorm" ? 2 : 1) * d_model *106                        (norms_per_layer * n_layers + 1);107        if (qk_norm) norms += 2 * head_dim() * n_layers;108        int64_t total = n_layers * attn + mlp_total + norms + vocab_size * d_model;109        if (!tied_embeddings) total += vocab_size * d_model;110        if (!use_rope) total += context_length * d_model;111        return total;112    }113};114115struct TrainConfig {116    float lr = 6e-4f;117    float min_lr_ratio = 0.1f;       // min_lr = lr * ratio118    int64_t warmup_steps = 2000;119    int64_t max_steps = 100000;120    std::string schedule = "cosine"; // "cosine" | "wsd" (warmup-stable-decay)121    float wsd_decay_frac = 0.15f;    // WSD: fraction of max_steps in the cooldown122    // "muon": Newton-Schulz orthogonalized momentum on 2-D hidden matrices,123    // AdamW on embeddings/head/norms (Keller Jordan's Muon). muon_lr is the124    // peak LR of the Muon group; it follows the same schedule shape as `lr`.125    std::string optimizer = "adamw"; // "adamw" | "muon"126    float muon_lr = 0.02f;127    float muon_momentum = 0.95f;128    float beta1 = 0.9f;129    float beta2 = 0.95f;130    float eps = 1e-8f;131    float weight_decay = 0.1f;       // dim>=2 params only132    float grad_clip = 1.0f;          // 0 disables133    int64_t batch_size = 32;         // sequences per micro-batch134    int64_t grad_accum_steps = 1;135    std::string precision = "f32";   // "f32" | "f16" | "bf16" (compute dtype)136    int64_t checkpoint_every = 1000;137    // Also commit each checkpoint's weights into <out>/model.forge — the138    // git-style zero-copy format (src/core/fmodel.h). The .bin checkpoint139    // stays the resume source of truth (it carries optimizer state).140    bool forge_save = true;141    std::string forge_dtype = "f32"; // "f32" | "f16" | "bf16"142    int64_t eval_every = 500;143    int64_t eval_batches = 20;144    uint64_t seed = 1337;145    bool deterministic = false;146};147148struct Config {149    ModelConfig model;150    TrainConfig train;151};152153inline void from_json(const nlohmann::json& j, ModelConfig& c) {154    c.name = j.value("name", c.name);155    c.n_layers = j.value("n_layers", c.n_layers);156    c.d_model = j.value("d_model", c.d_model);157    c.n_heads = j.value("n_heads", c.n_heads);158    c.n_kv_heads = j.value("n_kv_heads", c.n_heads);159    c.d_ff = j.value("d_ff", c.d_ff);160    c.vocab_size = j.value("vocab_size", c.vocab_size);161    c.context_length = j.value("context_length", c.context_length);162    c.tied_embeddings = j.value("tied_embeddings", c.tied_embeddings);163    c.use_rope = j.value("use_rope", c.use_rope);164    c.rope_theta = j.value("rope_theta", c.rope_theta);165    c.norm = j.value("norm", c.norm);166    c.norm_eps = j.value("norm_eps", c.norm_eps);167    c.activation = j.value("activation", c.activation);168    c.dropout = j.value("dropout", c.dropout);169    c.quant = j.value("quant", c.quant);170    c.n_experts = j.value("n_experts", c.n_experts);171    c.moe_top_k = j.value("moe_top_k", c.moe_top_k);172    c.moe_aux_weight = j.value("moe_aux_weight", c.moe_aux_weight);173    c.n_shared_experts = j.value("n_shared_experts", c.n_shared_experts);174    c.moe_scoring = j.value("moe_scoring", c.moe_scoring);175    c.moe_norm_topk = j.value("moe_norm_topk", c.moe_norm_topk);176    c.routed_scaling_factor = j.value("routed_scaling_factor", c.routed_scaling_factor);177    c.moe_d_ff = j.value("moe_d_ff", c.moe_d_ff);178    c.first_k_dense = j.value("first_k_dense", c.first_k_dense);179    c.moe_bias_gamma = j.value("moe_bias_gamma", c.moe_bias_gamma);180    c.qk_norm = j.value("qk_norm", c.qk_norm);181    c.final_softcap = j.value("final_softcap", c.final_softcap);182    c.scale_embeddings = j.value("scale_embeddings", c.scale_embeddings);183    c.attention_bias = j.value("attention_bias", c.attention_bias);184    c.head_dim_override = j.value("head_dim", c.head_dim_override);185    c.nope_every = j.value("nope_every", c.nope_every);186    c.norm_placement = j.value("norm_placement", c.norm_placement);187    c.rope_scale_factor = j.value("rope_scale_factor", c.rope_scale_factor);188    c.rope_scale_low = j.value("rope_scale_low", c.rope_scale_low);189    c.rope_scale_high = j.value("rope_scale_high", c.rope_scale_high);190    c.rope_scale_orig_ctx = j.value("rope_scale_orig_ctx", c.rope_scale_orig_ctx);191    c.sliding_window = j.value("sliding_window", c.sliding_window);192    c.sliding_global_every = j.value("sliding_global_every", c.sliding_global_every);193    c.rope_theta_global = j.value("rope_theta_global", c.rope_theta_global);194    c.attn_softcap = j.value("attn_softcap", c.attn_softcap);195196    if (c.head_dim_override == 0 && c.d_model % c.n_heads != 0)197        throw std::runtime_error("config: d_model must be divisible by n_heads "198                                 "(or set head_dim explicitly)");199    if (c.n_heads % c.n_kv_heads != 0)200        throw std::runtime_error("config: n_heads must be divisible by n_kv_heads");201    if (c.head_dim() % 2 != 0)202        throw std::runtime_error("config: head_dim must be even (RoPE pairs)");203    if (c.norm != "rmsnorm" && c.norm != "layernorm")204        throw std::runtime_error("config: norm must be rmsnorm or layernorm");205    if (c.activation != "swiglu" && c.activation != "gelu" && c.activation != "relu2")206        throw std::runtime_error("config: activation must be swiglu, gelu or relu2");207    if (c.norm_placement != "pre" && c.norm_placement != "post" &&208        c.norm_placement != "sandwich")209        throw std::runtime_error("config: norm_placement must be pre, post or sandwich");210    if (c.sliding_window < 0 || c.sliding_global_every < 0 || c.nope_every < 0)211        throw std::runtime_error("config: window/pattern values must be >= 0");212    if (c.attn_softcap < 0.0f)213        throw std::runtime_error("config: attn_softcap must be >= 0");214    if (c.quant != "none" && c.quant != "int8" && c.quant != "ternary")215        throw std::runtime_error("config: quant must be none, int8 or ternary");216    if (c.n_experts < 0)217        throw std::runtime_error("config: n_experts must be >= 0");218    if (c.n_experts > 0 && (c.moe_top_k < 1 || c.moe_top_k > c.n_experts))219        throw std::runtime_error("config: moe_top_k must be in [1, n_experts]");220    if (c.n_shared_experts < 0 || (c.n_shared_experts > 0 && c.n_experts == 0))221        throw std::runtime_error("config: n_shared_experts requires n_experts > 0");222    if (c.moe_scoring != "softmax" && c.moe_scoring != "sigmoid")223        throw std::runtime_error("config: moe_scoring must be softmax or sigmoid");224    if (c.first_k_dense < 0 || c.first_k_dense > c.n_layers)225        throw std::runtime_error("config: first_k_dense must be in [0, n_layers]");226    if (c.moe_d_ff < 0 || c.moe_bias_gamma < 0.0f)227        throw std::runtime_error("config: moe_d_ff and moe_bias_gamma must be >= 0");228    if (c.final_softcap < 0.0f)229        throw std::runtime_error("config: final_softcap must be >= 0");230}231232inline void from_json(const nlohmann::json& j, TrainConfig& c) {233    c.lr = j.value("lr", c.lr);234    c.min_lr_ratio = j.value("min_lr_ratio", c.min_lr_ratio);235    c.warmup_steps = j.value("warmup_steps", c.warmup_steps);236    c.max_steps = j.value("max_steps", c.max_steps);237    c.schedule = j.value("schedule", c.schedule);238    c.wsd_decay_frac = j.value("wsd_decay_frac", c.wsd_decay_frac);239    c.optimizer = j.value("optimizer", c.optimizer);240    c.muon_lr = j.value("muon_lr", c.muon_lr);241    c.muon_momentum = j.value("muon_momentum", c.muon_momentum);242    c.beta1 = j.value("beta1", c.beta1);243    c.beta2 = j.value("beta2", c.beta2);244    c.eps = j.value("eps", c.eps);245    c.weight_decay = j.value("weight_decay", c.weight_decay);246    c.grad_clip = j.value("grad_clip", c.grad_clip);247    c.batch_size = j.value("batch_size", c.batch_size);248    c.grad_accum_steps = j.value("grad_accum_steps", c.grad_accum_steps);249    c.precision = j.value("precision", c.precision);250    c.checkpoint_every = j.value("checkpoint_every", c.checkpoint_every);251    c.forge_save = j.value("forge_save", c.forge_save);252    c.forge_dtype = j.value("forge_dtype", c.forge_dtype);253    c.eval_every = j.value("eval_every", c.eval_every);254    c.eval_batches = j.value("eval_batches", c.eval_batches);255    c.seed = j.value("seed", c.seed);256    c.deterministic = j.value("deterministic", c.deterministic);257258    if (c.schedule != "cosine" && c.schedule != "wsd")259        throw std::runtime_error("config: schedule must be cosine or wsd");260    if (c.wsd_decay_frac <= 0.0f || c.wsd_decay_frac >= 1.0f)261        throw std::runtime_error("config: wsd_decay_frac must be in (0, 1)");262    if (c.optimizer != "adamw" && c.optimizer != "muon")263        throw std::runtime_error("config: optimizer must be adamw or muon");264    if (c.forge_dtype != "f32" && c.forge_dtype != "f16" && c.forge_dtype != "bf16")265        throw std::runtime_error("config: forge_dtype must be f32, f16 or bf16");266}267268inline Config load_config(const std::string& path) {269    std::ifstream in(path);270    if (!in) throw std::runtime_error("config: cannot open " + path);271    nlohmann::json j = nlohmann::json::parse(in);272    Config cfg;273    if (j.contains("model")) j.at("model").get_to(cfg.model);274    if (j.contains("train")) j.at("train").get_to(cfg.train);275    return cfg;276}277278} // namespace forge279