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

418 lines
21 KiB
C
Raw Blame History

/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
/*
* test_model_bwd.c — WHOLE-MODEL gradient check for the full RWKV-x
* forward+backward assembly (rx_pretrain_forward / rx_pretrain_backward).
* This is the integration proof: embedding -> interleaved RWKV-7 + sparse-
* attention blocks -> final LayerNorm -> head -> cross-entropy.
*
* Method (same discipline as test_block_bwd.c / test_attn_bwd.c):
* - A DOUBLE-PRECISION mirror of the ENTIRE model forward -> CE loss is the
* finite-difference reference (fp32 forward noise would swamp the small LN
* / decay / a-b grads and false-fail). Every reduction (LN, softmax,
* log-sum-exp, WKV state) is accumulated in double.
* - The analytic backward runs the normal fp32 path (rx_pretrain_backward).
* - Central finite differences, PER-TENSOR normalised infinity-norm metric
* (max|g-fd| / max|fd|). Pass: < 1e-3 for EVERY parameter tensor.
* - Large matrices are element-SUBSAMPLED (evenly spaced, capped) to keep the
* whole-model FD tractable; small tensors are checked exhaustively.
*
* Toy configs carry BOTH block types (RWKV-x interleave) and, in the 2-RWKV
* case, the v_first coupling (block-0 produces it, block-1 consumes it).
* MoBA moba_top_k is set >= #candidate chunks => full causal attention, which
* is also what foundry's rx_forward runs, so the assembly can be compared to it
* (see test_model_forward_matches_rx_forward, guarded by RX_LINK_ENGINE).
*/
#include "lib/test.h"
#include "rwkvx/rx_pretrain.h"
#include <math.h>
#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#ifdef RX_LINK_ENGINE
#include "rwkvx/rwkvx.h"
/* Bridge implemented in rwkvx.c: build a real rx_model from the ptmodel. */
rx_model* rx_model_from_ptmodel(rx_ptmodel *pt);
#endif
/* ======================= fp64 model mirror ======================= */
static double dsig(double z){ return 1.0/(1.0+exp(-z)); }
static void ln_fp64(const double *x, const float *w, const float *b,
double *out, int C){
double mean=0.0; for(int i=0;i<C;i++) mean+=x[i]; mean/=C;
double var=0.0; for(int i=0;i<C;i++){double d=x[i]-mean;var+=d*d;} var/=C;
double rstd=1.0/sqrt(var+1e-5);
for(int i=0;i<C;i++) out[i]=(x[i]-mean)*rstd*(double)w[i]+(double)b[i];
}
/* RWKV-7 block fp64 forward: out[T*C]. vfirst[T*C] WRITTEN at layer 0 (=v),
* READ at layer>0. (Transcribed from test_block_bwd.c loss_fp64.) */
static void block_fwd_fp64(const rx_block_params *P, const double *x,
double *vfirst, double *out, int T){
const rx_tm_params *tm=&P->tm;
int C=tm->c,H=tm->n_head,N=tm->head_size;
int Dw=tm->Dw,Da=tm->Da,Dg=tm->Dg,Dv=tm->Dv,F=P->cm.ffn_dim;
int L0=(tm->layer_id==0);
double *h1=malloc((size_t)T*C*sizeof(double));
double *tmo=malloc((size_t)T*C*sizeof(double));
double *x1=malloc((size_t)T*C*sizeof(double));
double *h2=malloc((size_t)T*C*sizeof(double));
double *S=calloc((size_t)H*N*N,sizeof(double));
double *sh=malloc(C*sizeof(double)),*shv=malloc(C*sizeof(double));
double *r=malloc(C*sizeof(double)),*k=malloc(C*sizeof(double)),*v=malloc(C*sizeof(double));
double *w=malloc(C*sizeof(double)),*a=malloc(C*sizeof(double)),*g=malloc(C*sizeof(double));
double *kk=malloc(C*sizeof(double)),*kmod=malloc(C*sizeof(double)),*vf=malloc(C*sizeof(double));
int maxD=Dw>Da?(Dw>Dg?Dw:Dg):(Da>Dg?Da:Dg); if(Dv>maxD)maxD=Dv; if(maxD<1)maxD=1;
double *lo=malloc((size_t)(maxD+1)*sizeof(double));
double *lo2=malloc((size_t)((Dv>0?Dv:1))*sizeof(double));
double *y=malloc(C*sizeof(double)),*ygn=malloc(C*sizeof(double));
double *xd=malloc(C*sizeof(double));
double *fk=malloc((size_t)F*sizeof(double));
for (int t=0;t<T;t++){ size_t base=(size_t)t*C;
for(int i=0;i<C;i++) xd[i]=x[base+i];
ln_fp64(xd,P->ln1_w,P->ln1_b,&h1[base],C);
}
for (int e=0;e<H*N*N;e++) S[e]=0.0;
for (int t=0;t<T;t++){
size_t base=(size_t)t*C;
const double *xx=&h1[base];
const double *prev=(t==0)?NULL:&h1[base-C];
#define SHIFT(dst,lerp) for(int i=0;i<C;i++){ \
double pv=prev?prev[i]:0.0; dst[i]=xx[i]+(pv-xx[i])*(double)lerp[i]; }
SHIFT(sh,tm->x_r); for(int j=0;j<C;j++){double s=0;for(int i=0;i<C;i++)s+=sh[i]*tm->receptance_w[i*C+j]; r[j]=s;}
SHIFT(sh,tm->x_k); for(int j=0;j<C;j++){double s=0;for(int i=0;i<C;i++)s+=sh[i]*tm->key_w[i*C+j]; k[j]=s;}
SHIFT(shv,tm->x_v);
for(int j=0;j<C;j++){double s=0;for(int i=0;i<C;i++)s+=shv[i]*tm->value_w[i*C+j]; v[j]=s;}
SHIFT(sh,tm->x_w);
for(int d=0;d<Dw;d++){double s=0;for(int i=0;i<C;i++)s+=sh[i]*tm->w1[i*Dw+d]; lo[d]=tanh(s);}
for(int j=0;j<C;j++){double s=0;for(int d=0;d<Dw;d++)s+=lo[d]*tm->w2[d*C+j];
double sg=dsig((double)tm->w0[j]+s); w[j]=exp(-0.606531*sg);}
SHIFT(sh,tm->x_a);
for(int d=0;d<Da;d++){double s=0;for(int i=0;i<C;i++)s+=sh[i]*tm->a1[i*Da+d]; lo[d]=s;}
for(int j=0;j<C;j++){double s=0;for(int d=0;d<Da;d++)s+=lo[d]*tm->a2[d*C+j];
a[j]=dsig((double)tm->a0[j]+s);}
SHIFT(sh,tm->x_g);
for(int d=0;d<Dg;d++){double s=0;for(int i=0;i<C;i++)s+=sh[i]*tm->g1[i*Dg+d]; lo[d]=dsig(s);}
for(int j=0;j<C;j++){double s=0;for(int d=0;d<Dg;d++)s+=lo[d]*tm->g2[d*C+j]; g[j]=s;}
for(int i=0;i<C;i++) kk[i]=k[i]*(double)tm->k_k[i];
for(int h=0;h<H;h++){int off=h*N;double ss=0;for(int i=0;i<N;i++)ss+=kk[off+i]*kk[off+i];
double inv=1.0/(sqrt(ss)+1e-12); for(int i=0;i<N;i++) kk[off+i]*=inv;}
for(int i=0;i<C;i++) kmod[i]=k[i]*(1.0+(a[i]-1.0)*(double)tm->k_a[i]);
if (L0){ for(int i=0;i<C;i++){ vf[i]=v[i]; vfirst[base+i]=v[i]; } }
else {
for(int d=0;d<Dv;d++){double s=0;for(int i=0;i<C;i++)s+=shv[i]*tm->v1[i*Dv+d]; lo2[d]=s;}
for(int j=0;j<C;j++){double s=0;for(int d=0;d<Dv;d++)s+=lo2[d]*tm->v2[d*C+j];
double bl=dsig((double)tm->v0[j]+s);
vf[j]=v[j]+(vfirst[base+j]-v[j])*bl;}
}
for(int h=0;h<H;h++){int off=h*N;
for(int i=0;i<N;i++){
double *Si=&S[(size_t)(h*N+i)*N];
double sa=0; for(int j=0;j<N;j++) sa+=(-kk[off+j])*Si[j];
double vi=vf[off+i], yy=0;
for(int j=0;j<N;j++){ Si[j]=Si[j]*w[off+j]+kmod[off+j]*vi+sa*(kk[off+j]*a[off+j]);
yy+=Si[j]*r[off+j]; }
y[off+i]=yy;
}
}
for(int h=0;h<H;h++){int off=h*N;double mean=0;for(int i=0;i<N;i++)mean+=y[off+i];mean/=N;
double var=0;for(int i=0;i<N;i++){double d=y[off+i]-mean;var+=d*d;}var/=N;
double rstd=1.0/sqrt(var+64e-5);
for(int i=0;i<N;i++) ygn[off+i]=(y[off+i]-mean)*rstd*(double)tm->ln_x_w[off+i]+(double)tm->ln_x_b[off+i];}
for(int h=0;h<H;h++){int off=h*N;double s=0;
for(int j=0;j<N;j++)s+=r[off+j]*kmod[off+j]*(double)tm->r_k[off+j];
for(int j=0;j<N;j++)ygn[off+j]+=s*vf[off+j];}
for(int i=0;i<C;i++) ygn[i]*=g[i]; /* g gate */
for(int j=0;j<C;j++){double s=0;for(int i=0;i<C;i++)s+=ygn[i]*tm->output_w[i*C+j]; tmo[base+j]=s;}
#undef SHIFT
}
for (int t=0;t<T;t++){ size_t base=(size_t)t*C;
for(int i=0;i<C;i++) x1[base+i]=x[base+i]+tmo[base+i];
ln_fp64(&x1[base],P->ln2_w,P->ln2_b,&h2[base],C);
}
for (int t=0;t<T;t++){ size_t base=(size_t)t*C;
const double *xx=&h2[base]; const double *prev=(t==0)?NULL:&h2[base-C];
for(int i=0;i<C;i++){ double pv=prev?prev[i]:0.0;
sh[i]=xx[i]+(pv-xx[i])*(double)P->cm.ffn_x_k[i]; }
for(int mm=0;mm<F;mm++){double s=0;for(int i=0;i<C;i++)s+=sh[i]*P->cm.ffn_key_w[i*F+mm];
double rr=s>0.0?s:0.0; fk[mm]=rr*rr;}
for(int j=0;j<C;j++){double s=0;for(int mm=0;mm<F;mm++)s+=fk[mm]*P->cm.ffn_value_w[mm*C+j];
out[base+j]=x1[base+j]+s;}
}
free(h1);free(tmo);free(x1);free(h2);free(S);free(sh);free(shv);
free(r);free(k);free(v);free(w);free(a);free(g);free(kk);free(kmod);free(vf);
free(lo);free(lo2);free(y);free(ygn);free(xd);free(fk);
}
/* Sparse-attention block fp64 forward: out[T*C]. (Transcribed from
* test_attn_bwd.c loss_fp64; full causal when moba_top_k >= #candidates.) */
static void attn_fwd_fp64(const rx_attn_params *P, const double *x,
double *out, int T){
int C=P->c,H=P->n_head,N=P->head_size,F=P->cm.ffn_dim;
int cs=P->chunk_size<1?1:P->chunk_size;
double scale=1.0/sqrt((double)N);
double *h1=malloc((size_t)T*C*sizeof(double));
double *q=malloc((size_t)T*C*sizeof(double));
double *k=malloc((size_t)T*C*sizeof(double));
double *v=malloc((size_t)T*C*sizeof(double));
double *y=malloc((size_t)T*C*sizeof(double));
double *x1=malloc((size_t)T*C*sizeof(double));
double *h2=malloc((size_t)T*C*sizeof(double));
double *xd=malloc((size_t)C*sizeof(double));
double *sh=malloc((size_t)C*sizeof(double));
double *fk=malloc((size_t)F*sizeof(double));
for (int t=0;t<T;t++){ size_t base=(size_t)t*C;
for(int i=0;i<C;i++) xd[i]=x[base+i];
ln_fp64(xd,P->ln1_w,P->ln1_b,&h1[base],C);
for(int j=0;j<C;j++){ double sq=0,sk=0,sv=0;
for(int i=0;i<C;i++){ double hv=h1[base+i];
sq+=hv*P->receptance_w[i*C+j]; sk+=hv*P->key_w[i*C+j]; sv+=hv*P->value_w[i*C+j]; }
q[base+j]=sq; k[base+j]=sk; v[base+j]=sv; }
}
for (int h=0;h<H;h++){ int off=h*N;
for (int tq=0;tq<T;tq++){
size_t qb=(size_t)tq*C+off;
int local_chunk=tq/cs, ncand=local_chunk;
unsigned char *csel=calloc((size_t)(T>0?T:1),1);
int keepall=(P->moba_top_k<=0 || P->moba_top_k>=ncand);
if (keepall){ for(int cc=0;cc<ncand;cc++) csel[cc]=1; }
else {
double *gate=malloc((size_t)(T>0?T:1)*sizeof(double));
for(int c2=0;c2<ncand;c2++){ double gg=0;
for(int j=0;j<N;j++){ double mkj=0;
for(int s=c2*cs;s<(c2+1)*cs;s++) mkj+=k[(size_t)s*C+off+j];
mkj/=(double)cs; gg+=q[qb+j]*mkj; }
gate[c2]=gg; }
for(int it=0;it<P->moba_top_k;it++){ int arg=-1; double best=-1e300;
for(int c2=0;c2<ncand;c2++) if(!csel[c2]&&gate[c2]>best){best=gate[c2];arg=c2;}
if(arg<0)break; csel[arg]=1; }
free(gate);
}
double *sc=malloc((size_t)(T>0?T:1)*sizeof(double)); double maxs=-1e300;
for(int s=0;s<=tq;s++){ int c2=s/cs;
int inc=(c2==local_chunk)||(c2<ncand&&csel[c2]);
if(!inc){ sc[s]=-1e300; continue; }
double d=0; for(int j=0;j<N;j++) d+=q[qb+j]*k[(size_t)s*C+off+j];
d*=scale; sc[s]=d; if(d>maxs)maxs=d; }
double sum=0;
for(int s=0;s<=tq;s++){ if(sc[s]<=-1e299)continue; double e=exp(sc[s]-maxs); sc[s]=e; sum+=e; }
for(int j=0;j<N;j++){ double acc=0;
for(int s=0;s<=tq;s++){ if(sc[s]<=-1e299)continue;
acc+=(sc[s]/sum)*v[(size_t)s*C+off+j]; }
y[qb+j]=acc; }
free(sc); free(csel);
}
}
for (int t=0;t<T;t++){ size_t base=(size_t)t*C;
for(int j=0;j<C;j++){ double s=0; for(int i=0;i<C;i++) s+=y[base+i]*P->output_w[i*C+j];
x1[base+j]=x[base+j]+s; }
ln_fp64(&x1[base],P->ln2_w,P->ln2_b,&h2[base],C);
}
for (int t=0;t<T;t++){ size_t base=(size_t)t*C;
const double *xx=&h2[base]; const double *prev=(t==0)?NULL:&h2[base-C];
for(int i=0;i<C;i++){ double pv=prev?prev[i]:0.0;
sh[i]=xx[i]+(pv-xx[i])*(double)P->cm.ffn_x_k[i]; }
for(int mm=0;mm<F;mm++){ double s=0; for(int i=0;i<C;i++) s+=sh[i]*P->cm.ffn_key_w[i*F+mm];
double rr=s>0.0?s:0.0; fk[mm]=rr*rr; }
for(int j=0;j<C;j++){ double s=0; for(int mm=0;mm<F;mm++) s+=fk[mm]*P->cm.ffn_value_w[mm*C+j];
out[base+j]=x1[base+j]+s; }
}
free(h1);free(q);free(k);free(v);free(y);free(x1);free(h2);free(xd);free(sh);free(fk);
}
/* Whole-model fp64 forward -> mean CE loss. Reads (possibly perturbed) weights
* straight from the model's parameter views. logits_out optional [T*vocab]. */
static double model_loss_fp64(rx_ptmodel *m, const int *tokens,
const int *targets, int T, int ignore,
double *logits_out){
rx_pt_config pc=rx_ptmodel_config(m);
int C=pc.n_embd, V=pc.vocab_size;
int n_blocks=pc.n_rwkv_layer+pc.n_attn_layer;
/* rebuild the interleave order (same rule as rx_model_new_random) */
int *btype=malloc(n_blocks*sizeof(int)), *bidx=malloc(n_blocks*sizeof(int));
if (pc.n_attn_layer==0){ for(int i=0;i<n_blocks;i++){btype[i]=0;bidx[i]=i;} }
else { int interval=pc.n_rwkv_layer/pc.n_attn_layer, pos=0;
for(int a=0;a<pc.n_attn_layer;a++){ int st=a*interval,en=(a+1)*interval;
for(int j=st;j<en;j++){btype[pos]=0;bidx[pos]=j;pos++;}
btype[pos]=1;bidx[pos]=a;pos++; }
for(int j=pc.n_attn_layer*interval;j<pc.n_rwkv_layer;j++){btype[pos]=0;bidx[pos]=j;pos++;} }
size_t TC=(size_t)T*C;
double *x=malloc(TC*sizeof(double));
double *tmp=malloc(TC*sizeof(double));
double *vfirst=calloc(TC,sizeof(double));
{ /* emb gather — read emb from the registry (name "emb") */
float *emb=NULL;
for (int i=0;i<rx_ptmodel_num_params(m);i++){ const char*nm; float*d;int nn;
if (rx_ptmodel_param(m,i,&nm,&d,NULL,&nn)==0 && strcmp(nm,"emb")==0){emb=d;break;} }
for (int t=0;t<T;t++) for (int i=0;i<C;i++) x[(size_t)t*C+i]=(double)emb[(size_t)tokens[t]*C+i];
}
for (int pos=0;pos<n_blocks;pos++){
if (btype[pos]==0) block_fwd_fp64(rx_ptmodel_rwkv(m,bidx[pos]), x, vfirst, tmp, T);
else attn_fwd_fp64(rx_ptmodel_attn(m,bidx[pos]), x, tmp, T);
memcpy(x,tmp,TC*sizeof(double));
}
/* ln_out + head + CE */
float *lnw=NULL,*lnb=NULL,*head=NULL;
for (int i=0;i<rx_ptmodel_num_params(m);i++){ const char*nm; float*d;int nn;
if (rx_ptmodel_param(m,i,&nm,&d,NULL,&nn)!=0) continue;
if (!strcmp(nm,"ln_out_w")) lnw=d; else if(!strcmp(nm,"ln_out_b")) lnb=d;
else if(!strcmp(nm,"head")) head=d; }
double *h=malloc(C*sizeof(double));
double loss=0.0; int nv=0;
for (int t=0;t<T;t++){
ln_fp64(&x[(size_t)t*C],lnw,lnb,h,C);
double *lg = logits_out ? &logits_out[(size_t)t*V] : malloc(V*sizeof(double));
for (int j=0;j<V;j++){ double s=0; for(int i=0;i<C;i++) s+=h[i]*(double)head[(size_t)i*V+j]; lg[j]=s; }
int tg=targets[t];
if (!(ignore>=0 && tg==ignore)){
double mx=-1e300; for(int j=0;j<V;j++) if(lg[j]>mx)mx=lg[j];
double se=0; for(int j=0;j<V;j++) se+=exp(lg[j]-mx);
loss += (mx+log(se)) - lg[tg]; nv++;
}
if (!logits_out) free(lg);
}
free(x);free(tmp);free(vfirst);free(h);free(btype);free(bidx);
return nv>0?loss/nv:0.0;
}
/* ======================= FD driver ======================= */
#define FD_CAP 40 /* max element samples per tensor */
static float fd_tensor(rx_ptmodel *m, const int *tok, const int *tgt, int T,
int ignore, float *data, const float *grad, int n,
const char *nm, int verbose){
int cap = n<FD_CAP ? n : FD_CAP;
int step = n/cap; if (step<1) step=1;
float maxdiff=0.0f, maxfd=0.0f, maxg=0.0f;
for (int c=0;c<cap;c++){
int e=c*step; if (e>=n) break;
float o=data[e]; float hh=1e-3f*(1.0f+fabsf(o));
data[e]=o+hh; double Lp=model_loss_fp64(m,tok,tgt,T,ignore,NULL);
data[e]=o-hh; double Lm=model_loss_fp64(m,tok,tgt,T,ignore,NULL);
data[e]=o;
float fd=(float)((Lp-Lm)/(2.0*(double)hh));
float d=fabsf(grad[e]-fd);
if (d>maxdiff) maxdiff=d;
if (fabsf(fd)>maxfd) maxfd=fabsf(fd);
if (fabsf(grad[e])>maxg) maxg=fabsf(grad[e]);
}
float re=maxdiff/(maxfd+1e-9f);
if (verbose) printf(" %-22s rel=%.2e (max|g-fd|=%.2e max|fd|=%.2e)\n",
nm,re,maxdiff,maxfd);
return re;
}
static int run_case(const char *name, rx_pt_config pc, uint64_t seed, int T){
printf(" [%s] C=%d H=%d N=%d rwkv=%d attn=%d vocab=%d F=%d "
"Dw=%d Da=%d Dg=%d Dv=%d T=%d chunk=%d topk=%d\n",
name, pc.n_embd,pc.n_head,pc.head_size,pc.n_rwkv_layer,pc.n_attn_layer,
pc.vocab_size,pc.ffn_dim,pc.Dw,pc.Da,pc.Dg,pc.Dv,T,pc.chunk_size,pc.moba_top_k);
rx_ptmodel *m=rx_model_new_random(&pc,seed);
if (!m){ printf(" model_new_random FAILED\n"); return 0; }
/* tokens + targets */
int *tok=malloc(T*sizeof(int)), *tgt=malloc(T*sizeof(int));
srand((unsigned)seed);
for (int t=0;t<T;t++){ tok[t]=rand()%pc.vocab_size; tgt[t]=rand()%pc.vocab_size; }
rx_pt_acts *acts=rx_pt_acts_new(m,T);
if (rx_pretrain_forward(m,tok,T,acts)!=0){ printf(" forward alloc FAIL\n"); return 0; }
/* fp32 forward vs fp64 mirror logits (assembly self-consistency) */
double *lg64=malloc((size_t)T*pc.vocab_size*sizeof(double));
double loss64=model_loss_fp64(m,tok,tgt,T,-1,lg64);
float maxlog=0.0f, reflog=0.0f;
for (int t=0;t<T;t++) for (int j=0;j<pc.vocab_size;j++){
float a=acts->logits[(size_t)t*pc.vocab_size+j];
float d=fabsf(a-(float)lg64[(size_t)t*pc.vocab_size+j]);
if (d>maxlog) maxlog=d;
if (fabsf(a)>reflog) reflog=fabsf(a);
}
printf(" fp32-vs-fp64 logits: rel=%.2e (max|d|=%.2e max|logit|=%.2e)\n",
maxlog/(reflog+1e-9f), maxlog, reflog);
/* sanity: the assembly forward must match foundry's rx_forward on identical
* weights (full-causal attention == topk>=candidates). */
int fwd_ok = 1;
#ifdef RX_LINK_ENGINE
{
rx_model *rm = rx_model_from_ptmodel(m);
rx_state *st = rx_state_new(rm);
float maxd=0.0f, maxref=0.0f;
for (int t=0;t<T;t++){
tm_tensor *lg = rx_forward(rm, (uint32_t)tok[t], st);
const float *ld = tm_data_f32(lg);
for (int j=0;j<pc.vocab_size;j++){
float a=acts->logits[(size_t)t*pc.vocab_size+j];
float d=fabsf(a-ld[j]); if(d>maxd)maxd=d;
if(fabsf(a)>maxref)maxref=fabsf(a);
}
tm_free(lg);
}
rx_state_free(st); rx_model_free(rm);
float rel=maxd/(maxref+1e-9f);
printf(" assembly-vs-rx_forward logits: rel=%.2e (max|d|=%.2e)\n", rel, maxd);
fwd_ok = rel < 5e-3f;
if (!fwd_ok) printf(" rx_forward MISMATCH\n");
}
#else
printf(" assembly-vs-rx_forward: (skipped; build with RX_LINK_ENGINE)\n");
#endif
/* analytic backward */
double loss=rx_pretrain_backward(m,acts,tgt,T,-1);
int finite = isfinite(loss) && isfinite(loss64);
printf(" loss(fp32 CE)=%.6f loss(fp64 mirror)=%.6f finite=%s\n",
loss, loss64, finite?"yes":"NO");
/* per-tensor FD check */
float wmax=0.0f; const char *wname="";
int np=rx_ptmodel_num_params(m);
for (int i=0;i<np;i++){
const char *nm; float *data,*grad; int n;
rx_ptmodel_param(m,i,&nm,&data,&grad,&n);
/* verbose print only for representative tensors to keep output readable */
int rep = (strstr(nm,"emb")||strstr(nm,"head")||strstr(nm,"ln_out")||
!strncmp(nm,"rwkv0.",6)||!strncmp(nm,"attn0.",6));
float e=fd_tensor(m,tok,tgt,T,-1,data,grad,n,nm,rep);
if (e>wmax){ wmax=e; wname=nm; }
}
int pass = wmax<1e-3f && finite && fwd_ok;
printf(" -> worst rel-err = %.2e (%s) %s\n", wmax, wname, pass?"PASS":"FAIL");
free(lg64); free(tok); free(tgt);
rx_pt_acts_free(acts); rx_ptmodel_free(m);
return pass;
}
static void test_model_backward(void){
printf(" Whole-model RWKV-x forward+backward — fp64 finite differences:\n");
int ok=1;
/* Toy A: 1 RWKV + 1 sparse-attention (the minimal RWKV-x interleave). */
ok &= run_case("1R+1A", (rx_pt_config){
.n_embd=64,.n_head=2,.head_size=32,.n_rwkv_layer=1,.n_attn_layer=1,
.vocab_size=48,.Dw=8,.Da=8,.Dg=8,.Dv=8,.ffn_dim=32,
.chunk_size=2,.moba_top_k=999}, 1234, 7);
/* Toy B: 2 RWKV + 1 sparse-attention — exercises v_first coupling
* (block-0 produces, block-1 consumes) AND the interleave [R0,R1,A0]. */
ok &= run_case("2R+1A", (rx_pt_config){
.n_embd=48,.n_head=3,.head_size=16,.n_rwkv_layer=2,.n_attn_layer=1,
.vocab_size=40,.Dw=6,.Da=6,.Dg=6,.Dv=6,.ffn_dim=24,
.chunk_size=2,.moba_top_k=999}, 777, 6);
ASSERT_TRUE(ok);
}
int main(void){
printf("\nWhole-model RWKV-x gradient check (from-scratch trainer)\n");
RUN_TEST(model_backward);
TEST_REPORT();
return __ft_failed>0 ? 1 : 0;
}
Reference in New Issue View Git Blame Copy Permalink