David
/
Foundry
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Foundry/rwkvx/tests/test_overfit.c

363 lines
16 KiB
C
Raw Permalink Blame History

/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
/*
* test_overfit.c — Training-loop smoke test for the RWKV-x from-scratch trainer,
* with the optimizer SWAPPABLE between AdamW (all-continuous baseline) and SR4
* (ternary projection matrices via the latent 4-bit accumulator).
*
* This is the FIRST proof that either optimizer drives an actual language model:
* we overfit a single FIXED random token sequence (pure memorization, no corpus)
* and watch the cross-entropy fall.
*
* - AdamW arm : every parameter is continuous, trained by foundry optim
* (fo_adamw_new + fo_step). Expected to drive CE -> ~0
* (near-perfect memorization) => proves the loop is correct.
* - SR4 arm : the big linear projection matrices (attention + time-mix
* receptance/key/value/output, channel-mix ffn_key/ffn_value)
* are TERNARY {-1,0,+1} materialized from a per-weight int4
* accumulator `q` (ternary_optim / SR4), each with a per-OUTPUT-
* COLUMN fp32 scale restoring dynamic range:
* W_eff[i,j] = scale[j] * ternaryW[i,j]
* Z = W_eff^T x == scale (*) (ternaryW^T x) (folded)
* Everything else (embedding, head, LayerNorm gamma/beta, the
* RWKV-7 LoRA gate matrices, biases, lerps, AND the per-column
* scales themselves) stays continuous and trains with AdamW.
*
* Param split is by registry NAME suffix (see is_ternary_name).
*
* Ternary-scale coupling (the one real subtlety):
* The model forward/backward only ever sees W_eff (we overwrite the model's
* own weight buffer with scale*ternaryW before each forward), so the analytic
* grad the backward produces is g = dL/dW_eff. We then split it:
* grad wrt ternary weight (drives q) : gq[i,j] = g[i,j] * scale[j]
* grad wrt scale : sg[j] = sum_i g[i,j]*ternaryW[i,j]
* SR4-step q with gq; AdamW-step the scale vector with sg (as a normal
* continuous param).
*
* Build: see Makefile target `test-overfit` (standalone; NOT in the default
* `test` aggregate — it is a slow optimization loop).
*/
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <stdint.h>
#include "../rx_pretrain.h"
#include "../../ternary_optim/ternary_optim.h"
#include "../../optim/optim.h"
#include "../../tensor/tensor.h"
/* ------------------------------------------------------------------ */
/* Tiny config for the overfit smoke. */
/* ------------------------------------------------------------------ */
#define CFG_C 64 /* n_embd */
#define CFG_HEAD 4
#define CFG_HS 16 /* head_size (HEAD*HS == C) */
#define CFG_VOCAB 64
#define CFG_T 16 /* sequence length */
#define CFG_FFN 128
#define CFG_LORA 16
static rx_pt_config make_cfg(void){
rx_pt_config pc;
memset(&pc,0,sizeof pc);
pc.n_embd=CFG_C; pc.n_head=CFG_HEAD; pc.head_size=CFG_HS;
pc.n_rwkv_layer=1; pc.n_attn_layer=1;
pc.vocab_size=CFG_VOCAB;
pc.Dw=CFG_LORA; pc.Da=CFG_LORA; pc.Dg=CFG_LORA; pc.Dv=CFG_LORA;
pc.ffn_dim=CFG_FFN;
pc.chunk_size=CFG_T; /* one chunk => full causal attention */
pc.moba_top_k=0; /* keep all */
return pc;
}
/* Fixed random token sequence (input) + shifted target. length T+1. */
static void make_fixed_seq(int *tokens, int *targets, int T, int vocab){
uint64_t s = 0xC0FFEEULL;
int seq[CFG_T+1];
for (int t=0;t<=T;t++){
s = s*6364136223846793005ULL + 1442695040888963407ULL;
seq[t] = (int)((s>>33) % (uint64_t)vocab);
}
for (int t=0;t<T;t++){ tokens[t]=seq[t]; targets[t]=seq[t+1]; }
}
/* ------------------------------------------------------------------ */
/* Param classification: the big projection matrices go ternary. */
/* ------------------------------------------------------------------ */
static const char *suffix_after_dot(const char *name){
const char *dot = strrchr(name,'.');
return dot ? dot+1 : name;
}
static int is_ternary_name(const char *name){
const char *s = suffix_after_dot(name);
return !strcmp(s,"receptance_w") || !strcmp(s,"key_w") ||
!strcmp(s,"value_w") || !strcmp(s,"output_w") ||
!strcmp(s,"ffn_key_w") || !strcmp(s,"ffn_value_w");
}
/* rows of a [in,out] row-major projection so cols(=out) = numel/rows. */
static int rows_of(const char *name, int numel){
const char *s = suffix_after_dot(name);
if (!strcmp(s,"ffn_key_w")) return CFG_C; /* [C, ffn] */
if (!strcmp(s,"ffn_value_w")) return CFG_FFN; /* [ffn, C] */
(void)numel;
return CFG_C; /* [C, C] square projections */
}
/* ------------------------------------------------------------------ */
/* Wrap a model float buffer + its grad buffer as a foundry tm_tensor. */
/* ------------------------------------------------------------------ */
static tm_tensor *wrap_param(float *data, float *grad, int n){
int64_t shape[1] = { n };
tm_tensor *t = tm_from_data(data, shape, 1, TM_F32);
t->grad = tm_from_data(grad, shape, 1, TM_F32);
t->requires_grad = true;
return t;
}
/* ------------------------------------------------------------------ */
/* Cosine LR schedule. */
/* ------------------------------------------------------------------ */
static float cosine_lr(float lr0, float lr_min, int step, int total){
if (step >= total) return lr_min;
float c = 0.5f*(1.0f + cosf((float)M_PI * (float)step/(float)total));
return lr_min + (lr0-lr_min)*c;
}
/* ================================================================== */
/* AdamW arm: every param continuous. */
/* ================================================================== */
static double run_adamw(const rx_pt_config *pc, const int *tokens,
const int *targets, int T,
int steps, int report_every,
float lr0, float lr_min, int verbose){
rx_ptmodel *m = rx_model_new_random(pc, 1234);
rx_pt_acts *acts = rx_pt_acts_new(m, T);
int nP = rx_ptmodel_num_params(m);
tm_tensor **params = calloc(nP, sizeof(tm_tensor*));
for (int i=0;i<nP;i++){
float *d,*g; int n; const char *nm;
rx_ptmodel_param(m,i,&nm,&d,&g,&n);
params[i] = wrap_param(d,g,n);
}
fo_adam_config cfg = fo_adam_default(lr0);
cfg.weight_decay = 0.0f; /* pure memorization, no decay */
fo_optimizer *opt = fo_adamw_new(params, nP, &cfg);
double loss = 0.0;
if (verbose) printf(" [AdamW] lr0=%.4g -> lr_min=%.4g, %d steps\n", lr0, lr_min, steps);
for (int step=0; step<steps; step++){
float lr = cosine_lr(lr0, lr_min, step, steps);
fo_set_lr(opt, lr);
rx_ptmodel_zero_grad(m);
rx_pretrain_forward(m, tokens, T, acts);
loss = rx_pretrain_backward(m, acts, targets, T, -1);
fo_step(opt);
if (verbose && (step % report_every == 0 || step==steps-1))
printf(" step %4d lr %.4g CE %.6f\n", step, lr, loss);
}
for (int i=0;i<nP;i++){ tm_free(params[i]); } /* releases ->grad too */
free(params);
fo_free(opt);
rx_pt_acts_free(acts);
rx_ptmodel_free(m);
return loss;
}
/* ================================================================== */
/* SR4 arm: ternary projections via latent int4 q + per-column scale. */
/* ================================================================== */
typedef struct {
float *data; /* model weight buffer (overwritten w/ W_eff) */
float *grad; /* dL/dW_eff (from backward) */
int rows, cols, numel;
ternary_optim *opt; /* owns the int4 accumulator q */
float *tern; /* pure ternary {-1,0,+1} [numel] */
float *scale; /* per-output-column fp32 scale [cols] */
float *sgrad; /* scale grad [cols] */
float *gq; /* grad wrt ternary weight = g*scale [numel] */
tm_tensor *scale_t; /* scale as a continuous AdamW param */
} tern_mat;
/* freeze_noncore: when nonzero, ONLY the ternary q accumulators and their
* per-column scales train — emb/head/LayerNorm/LoRA/biases are held at random
* init. This isolates how much of the loss drop is genuinely driven by SR4 on
* the ternary projections vs. carried by the powerful continuous emb/head. */
static double run_sr4(const rx_pt_config *pc, const int *tokens,
const int *targets, int T,
int steps, int report_every,
float target_step, float cont_lr, float scale_lr,
int freeze_noncore, int verbose){
rx_ptmodel *m = rx_model_new_random(pc, 1234);
rx_pt_acts *acts = rx_pt_acts_new(m, T);
int nP = rx_ptmodel_num_params(m);
/* Partition params: ternary matrices vs continuous. */
tern_mat *tmat = calloc(nP, sizeof(tern_mat));
int n_tern = 0;
tm_tensor **cont = calloc(2*nP, sizeof(tm_tensor*)); /* + scale vectors */
int n_cont = 0;
int64_t idx_base = 0;
for (int i=0;i<nP;i++){
float *d,*g; int n; const char *nm;
rx_ptmodel_param(m,i,&nm,&d,&g,&n);
if (is_ternary_name(nm)){
tern_mat *tt = &tmat[n_tern++];
tt->data=d; tt->grad=g; tt->numel=n;
tt->rows = rows_of(nm,n);
tt->cols = n / tt->rows;
ternary_optim_config oc; ternary_optim_config_default(&oc);
oc.seed = 42; oc.idx_base = idx_base; idx_base += n;
tt->opt = ternary_optim_create(tt->rows, tt->cols, &oc);
tt->tern = malloc((size_t)n*sizeof(float));
tt->gq = malloc((size_t)n*sizeof(float));
tt->scale= malloc((size_t)tt->cols*sizeof(float));
tt->sgrad= calloc((size_t)tt->cols,sizeof(float));
for (int c=0;c<tt->cols;c++) tt->scale[c]=1.0f;
/* scale trains as a continuous AdamW param (grad in sgrad) */
tt->scale_t = wrap_param(tt->scale, tt->sgrad, tt->cols);
cont[n_cont++] = tt->scale_t;
} else if (!freeze_noncore){
cont[n_cont++] = wrap_param(d,g,n);
}
}
/* AdamW over all continuous params (incl. per-column scales). */
fo_adam_config cfg = fo_adam_default(cont_lr);
cfg.weight_decay = 0.0f;
fo_optimizer *opt = fo_adamw_new(cont, n_cont, &cfg);
if (verbose)
printf(" [SR4] target_step=%.4g cont_lr=%.4g scale_lr=%.4g "
"(%d ternary mats, %d continuous params)\n",
target_step, cont_lr, scale_lr, n_tern, n_cont);
double loss = 0.0;
for (int step=0; step<steps; step++){
/* 1. materialize W_eff = scale (*) ternaryW into the model buffers */
for (int t=0;t<n_tern;t++){
tern_mat *tt=&tmat[t];
ternary_optim_materialize(tt->opt, tt->tern);
for (int r=0;r<tt->rows;r++)
for (int c=0;c<tt->cols;c++)
tt->data[(size_t)r*tt->cols+c] =
tt->scale[c]*tt->tern[(size_t)r*tt->cols+c];
}
/* 2. forward + backward (grads land in model buffers = dL/dW_eff) */
rx_ptmodel_zero_grad(m);
rx_pretrain_forward(m, tokens, T, acts);
loss = rx_pretrain_backward(m, acts, targets, T, -1);
/* 3. split ternary grad -> {q-grad, scale-grad} */
for (int t=0;t<n_tern;t++){
tern_mat *tt=&tmat[t];
for (int c=0;c<tt->cols;c++){
double sg=0.0;
for (int r=0;r<tt->rows;r++){
size_t k=(size_t)r*tt->cols+c;
sg += (double)tt->grad[k]*(double)tt->tern[k];
tt->gq[k] = tt->grad[k]*tt->scale[c]; /* dL/d ternaryW */
}
tt->sgrad[c]=(float)sg; /* dL/d scale[c] */
}
}
/* 4a. SR4-step the ternary accumulators */
for (int t=0;t<n_tern;t++)
ternary_optim_step(tmat[t].opt, tmat[t].gq, 1,
(uint64_t)step, target_step);
/* 4b. AdamW-step the continuous params (incl. scales) */
fo_set_lr(opt, cont_lr);
/* scale-specific LR: temporarily not separated in foundry optim; we
* fold scale_lr in by pre-multiplying its grad relative to cont_lr. */
if (scale_lr != cont_lr){
float k = scale_lr/cont_lr;
for (int t=0;t<n_tern;t++)
for (int c=0;c<tmat[t].cols;c++) tmat[t].sgrad[c]*=k;
}
fo_step(opt);
if (verbose && (step % report_every == 0 || step==steps-1))
printf(" step %4d CE %.6f\n", step, loss);
}
/* teardown */
for (int i=0;i<n_cont;i++){ tm_free(cont[i]); } /* releases ->grad too */
for (int t=0;t<n_tern;t++){
ternary_optim_free(tmat[t].opt);
free(tmat[t].tern); free(tmat[t].gq);
free(tmat[t].scale); free(tmat[t].sgrad);
}
free(cont); free(tmat);
fo_free(opt);
rx_pt_acts_free(acts);
rx_ptmodel_free(m);
return loss;
}
/* ================================================================== */
int main(int argc, char **argv){
const char *mode = (argc>1)? argv[1] : "both";
rx_pt_config pc = make_cfg();
int tokens[CFG_T], targets[CFG_T];
make_fixed_seq(tokens, targets, CFG_T, CFG_VOCAB);
double ln_vocab = log((double)CFG_VOCAB);
printf("=== RWKV-x overfit smoke ===\n");
printf("config: C=%d heads=%d hs=%d vocab=%d T=%d ffn=%d lora=%d "
"(1 RWKV + 1 attn layer)\n",
CFG_C,CFG_HEAD,CFG_HS,CFG_VOCAB,CFG_T,CFG_FFN,CFG_LORA);
printf("fixed sequence (memorization target), random-init CE ~= ln(vocab) = %.4f\n\n", ln_vocab);
int fail = 0;
if (!strcmp(mode,"adamw") || !strcmp(mode,"both")){
printf("--- AdamW arm (all-continuous baseline) ---\n");
double L = run_adamw(&pc, tokens, targets, CFG_T,
1200, 100, 1e-2f, 1e-4f, 1);
printf(" => AdamW final CE = %.6f (%s memorized)\n\n",
L, (L<0.05)?"YES —":"NOT fully");
if (L >= 0.05) fail = 1;
}
if (!strcmp(mode,"sr4") || !strcmp(mode,"both")){
printf("--- SR4 arm (ternary projections) ---\n");
/* coarse step-size sweep */
float cand[] = {0.003f, 0.01f, 0.03f, 0.1f, 0.3f};
int ncand = (int)(sizeof(cand)/sizeof(cand[0]));
float best_ts = cand[0]; double best_L = 1e9;
printf(" coarse target_step sweep (400 steps each):\n");
for (int i=0;i<ncand;i++){
double L = run_sr4(&pc, tokens, targets, CFG_T,
400, 400, cand[i], 5e-3f, 1e-2f, 0, 0);
printf(" target_step=%-6.3g final CE %.6f\n", cand[i], L);
if (L < best_L){ best_L=L; best_ts=cand[i]; }
}
printf(" best target_step = %.3g (CE %.6f); running full curve "
"(prescribed split: emb/head/LN/LoRA continuous):\n", best_ts, best_L);
double L = run_sr4(&pc, tokens, targets, CFG_T,
1500, 100, best_ts, 5e-3f, 1e-2f, 0, 1);
printf(" => SR4 final CE = %.6f (best target_step %.3g)\n\n", L, best_ts);
/* Isolation probe: freeze emb/head/LN/LoRA at random init; only SR4's
* ternary q + per-column scales may train. Answers "is SR4 actually
* training the projections, or is the continuous emb/head doing it?" */
printf(" --- isolation probe: ONLY ternary q + scales train "
"(emb/head/LN/LoRA frozen at random init) ---\n");
double Lf = run_sr4(&pc, tokens, targets, CFG_T,
1500, 100, best_ts, 5e-3f, 1e-2f, 1, 1);
printf(" => SR4-only-projections final CE = %.6f "
"(start ~%.3f); ternary path %s learning\n\n",
Lf, ln_vocab, (Lf < ln_vocab-0.3)?"IS":"is NOT");
}
if (!strcmp(mode,"both")){
if (fail){ printf("RESULT: FAIL — AdamW did not memorize.\n"); return 1; }
printf("RESULT: PASS — AdamW memorized; SR4 wired and reported.\n");
}
return 0;
}
Reference in New Issue View Git Blame Copy Permalink