spb/forge Public MIT

Forge Studio (the GUI companion) tails log.csv as its metrics channel;
wall-clock elapsed enables its time-axis mode and honest ETA math.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

- 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>

Wave 1 (module-level):
- rope refactor: kernels read a host-precomputed inv-freq table; unlocks
  HF-"llama3" rope scaling (rope_scale_*), per-layer theta, and NoPE
  layers (nope_every, SmolLM3)
- attention_bias (Qwen2.5 QKV bias), head_dim decoupled from
  d_model/n_heads (Qwen3), relu2 activation (nanoGPT-speedrun lineage),
  norm_placement pre|post|sandwich (OLMo2/Gemma)

Wave 2 (attention kernels):
- sliding_window + sliding_global_every (Mistral / Gemma3 local:global
  patterns) in the CPU reference, the unfused Metal kernels, and the
  fused scalar flash kernels — out-of-window KV blocks are skipped, so
  cost scales with the window; window > 0 auto-routes off the MMA kernel
- attn_softcap (Gemma2): cap*tanh on scores pre-softmax, unfused path,
  exact tanh' chain in all backwards
- rope_theta_global for dual-theta local/global layers (Gemma3)

Parity suites cover every knob (fused + unfused paths); gradcheck and
overfit stay green. New demo configs: gpt-50m-mistral, gpt-50m-gemma;
ARCHITECTURES.md documents the per-family config matrix.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

- tools/prepare_hf_data.py: stream any of 13 registered HF datasets
  (FineWeb-Edu, DCLM, Cosmopedia, FineMath, OpenWebMath, Wikipedia, C4,
  SmolTalk, ...) or weighted mixtures/presets (smollm-web, textbooks,
  decay-anneal) straight into train.bin/val.bin — no full downloads
- SMALL_MODELS_RESEARCH.md: how sub-1B models get logical, useful text
  (data quality, deep-and-thin, distillation, test-time compute)
- INFERENCE_RESEARCH.md: Apple Silicon inference speed playbook tied to
  the .forge format (bandwidth math, fused-dequant GEMV, KV cache,
  residency sets, warmup, quant layouts) with a prioritized roadmap
- ARCHITECTURES.md: config matrix to train Llama/Mistral/Qwen/Gemma/
  DeepSeek/Kimi-class variants, with a 3-wave implementation plan
- README: training modes, .forge format, tools

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

A .forge is a model repository: tiny JSON manifests (one commit per save,
with parent links) over content-addressed shards. Tensors are 16KB-page-
aligned inside shards padded to page multiples, so loading is mmap +
newBuffer(bytesNoCopy) — on unified memory the file-cache pages ARE the
GPU memory. Saves are deltas: only tensors whose FNV-1a hash changed since
the parent manifest are written. Shards cap at 95MB (GitHub-pushable).
Store f32 (zero-copy alias at load) or f16/bf16 (half size).

- src/core/fmodel.{h,cpp}: save() + Snapshot zero-copy reader
- Tensor::from_buffer: views over externally-owned MTLBuffers
- forge export CLI; generate/eval accept .forge repos directly
- trainer commits weights natively to <out>/model.forge at each checkpoint
  (forge_save/forge_dtype config keys); .bin keeps optimizer state for resume
- tools/fmodel.py: inspect, log (history), to-safetensors (pure numpy)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

- 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>

- optimizer "adamw" | "muon": Newton-Schulz orthogonalized momentum on 2-D
  hidden matrices (composed from the existing matmul kernels on Metal),
  AdamW kept for embeddings/head/1-D params; muon_lr follows the lr schedule
- schedule "cosine" | "wsd": warmup-stable-decay with 1-sqrt cooldown,
  extendable runs, wsd_decay_frac
- gpt-50m base config + Muon+WSD and MobileLLM-style deep-and-thin variants

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

8 x 32 x 1024 = 262144 tokens/step (effective batch unchanged, so the 3e-4 LR
still applies) x 1553 steps = one epoch over the 407,344,713-token corpus.
Micro-batch 8 rather than 16 halves activation memory, which was sitting at
65 of 77.8 GB working set.

vocab 4096 (matches the proven BPE pipeline; the trainer's pair-count array is
O(V^2) and the encoder O(V*N), which does not scale to 16384 over 400M tokens),
with d_ff 3072 restoring the parameter count to 205.6M. Micro-batch 16 rather
than 4 since 96 GB affords it and it cuts sync points 4x.

A complete transformer training stack with no ML dependencies: tensors,
autograd, hand-written Metal kernels, flash attention (forward and backward),
AdamW, BPE tokenizer, checkpointing and generation. Architecture is fully
config-driven — the same binary trains 12M to 205M parameter models.

Every Metal kernel is validated against a CPU reference (85 parity checks,
<=1e-4, most bit-exact), gradients against central finite differences, and
each optimization was accepted only after the training loss trajectory stayed
numerically unchanged.

Measured findings (M5 Max, documented in RESEARCH.md and paper/forge.tex):

- `constant constexpr` for MSL tile constants declares an address-space
  variable, not a compile-time constant. Loops stop unrolling and every
  matrix accumulator spills: 0.82 -> 10.21 TFLOPS once switched to enums.
- That defect is invisible in the AIR at every -O level, because unrolling
  happens in the driver back end. Benchmark; do not read the IR.
- Register pressure, not bandwidth, dominates attention backward. Guided by
  measured spill counts, three restructurings took it 107 -> 7.05 ms (15.2x).
- On M5, mpp::tensor_ops::matmul2d reaches 51.5 TFLOPS with f16 operands vs
  10.6 for a tuned simdgroup_matrix kernel (4.9x), verified numerically.
  f16 on the simdgroup path alone is worth only +18-22%.
- Concurrent dispatch for the optimizer sweep: +22% on the 100M config.

Trained the 12.2M config for one epoch over 19.14M TinyStories tokens:
loss 8.40 -> 2.99, validation 3.009, perplexity 20.27, ~38.2k tokens/sec.