SPB Git

spb/forge Public MIT

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

C++ 61.2% C 23% Python 7.6% TeX 7.2% CMake 1.1%

History of src/kernels/moe.metal · clear filter

  1. Add wave 3: DeepSeek-V3-style MoE routing, all config-selected
    - moe_scoring "softmax"|"sigmoid" (new sigmoid op, CPU+Metal+backward)
    - aux-loss-free balancing (V3 "noaux"): top-k selection ranks score+bias
      while gate values stay biasless; per-expert load counted on-GPU each
      forward and the balance bias nudged ±moe_bias_gamma after every
      optimizer step; bias is checkpointed as a grad-free parameter
    - moe_norm_topk (renormalize kept gates or keep raw sigmoid scores),
      routed_scaling_factor (V3: 2.5), moe_d_ff (per-expert width),
      first_k_dense (dense MLPs for the first k layers)
    - parity tests: biased/no-renorm topk, sigmoid, expert_counts, and a
      full V3-style model (sigmoid/noaux/scaled/first-dense) — CPU==GPU
    
    Remaining wave-3 items (MLA, MuonClip QK-clip, MTP) documented in
    ARCHITECTURES.md.
    
    Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
    simon-pierre boucher committed 5 days ago (Aug 5, 2026) · 1 file changed +62 −40
  2. 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) · 1 file changed +120