spb/forge Public MIT

1// Author: Simon-Pierre Boucher — contact@spboucher.ai2//3// End-to-end sanity (CLAUDE.md testing protocol #3): overfit a single batch4// of 64 sequences to loss < 0.05 within 500 steps. Tiny model so the CPU5// reference path finishes in reasonable time; the same test re-runs on the6// GPU path at M4.7#include "nn/transformer.h"8#include "ops/ops.h"9#include "train/optimizer.h"1011#include <cstdio>12#include <random>1314using namespace forge;1516int main() {17    ModelConfig cfg;18    cfg.n_layers = 2;19    cfg.d_model = 64;20    cfg.n_heads = 2;21    cfg.n_kv_heads = 2;22    cfg.d_ff = 172; // ~8/3 * d23    cfg.vocab_size = 128;24    cfg.context_length = 32;25    cfg.tied_embeddings = true;2627    const int64_t B = 64, T = 32;2829    nn::Transformer model(cfg, 1337);3031    std::mt19937_64 rng(999);32    Tensor ids = Tensor::empty({B, T}, DType::I32);33    cpu::fill_uniform_int(ids, 0, cfg.vocab_size, rng);34    // next-token targets within the fixed batch35    Tensor targets = Tensor::empty({B * T}, DType::I32);36    for (int64_t b = 0; b < B; ++b) {37        for (int64_t t = 0; t < T; ++t) {38            targets.data<int32_t>()[b * T + t] =39                (t + 1 < T) ? ids.data<int32_t>()[b * T + t + 1] : -1;40        }41    }4243    train::AdamW::Options opts;44    opts.weight_decay = 0.0f; // pure memorization task45    train::AdamW opt(model.named_parameters(), opts);4647    const float lr = 3e-3f;48    float loss_val = -1.0f;49    for (int step = 0; step < 500; ++step) {50        opt.zero_grad();51        Var loss = model.loss(ids, targets);52        Tape::get().backward(loss);53        opt.clip_global_norm(1.0f);54        opt.step(lr);55        loss_val = loss.value().data<float>()[0];56        if (step % 50 == 0) std::printf("step %3d loss %.4f\n", step, double(loss_val));57        if (loss_val < 0.05f) {58            std::printf("step %3d loss %.4f — target reached\n", step, double(loss_val));59            break;60        }61    }6263    if (loss_val < 0.05f) {64        std::printf("\noverfit ok (final loss %.4f)\n", double(loss_val));65        return 0;66    }67    std::printf("\noverfit FAILED (final loss %.4f >= 0.05)\n", double(loss_val));68    return 1;69}70