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

329 lines
16 KiB
C
Raw Blame History

/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
/*
* test_block_bwd.c — Gradient-check for the FULL RWKV-7 BLOCK backward
* (rx_block_backward), the counterpart of rwkvx.c `rwkv7_block_forward_fast`:
* x -> ln1 -> time-mix -> +x -> ln2 -> channel-mix -> +x1
*
* Method (same discipline as test_timemix_bwd.c / docs §5-§6):
* - A DOUBLE-PRECISION mirror of the whole block computes the FD reference
* loss L = sum(cc * out). fp32 forward noise would swamp the small
* a/b/decay/LN grads and false-fail, so the reference is fp64.
* - The analytic backward runs the normal fp32 path (rx_block_backward).
* - Central finite differences, PER-TENSOR normalised infinity-norm metric
* (max|g-fd| / max|fd|). Pass: < 1e-3 for EVERY tensor and the block dx.
*
* Focus of this chunk: the channel-mix params (ffn_x_k, ffn_key_w,
* ffn_value_w), the two LayerNorms (ln1/ln2 gamma/beta) and the whole-block
* dx (residual grad threading). The time-mix + WKV7 params are re-checked
* for free since the block exercises them end-to-end.
*/
#include "lib/test.h"
#include "rwkvx/rx_pretrain.h"
#include <math.h>
#include <stdlib.h>
#include <string.h>
typedef struct {
int B, T, H, N, C, Dw, Da, Dg, Dv, layer_id, F; /* F = ffn_dim */
} cfg;
/* Every learnable tensor in a block, fp32. */
typedef struct {
/* time-mix */
float *x_r,*x_k,*x_v,*x_w,*x_a,*x_g;
float *receptance_w,*key_w,*value_w,*output_w;
float *w0,*w1,*w2, *a0,*a1,*a2, *g1,*g2, *v0,*v1,*v2;
float *k_k,*k_a,*r_k,*ln_x_w,*ln_x_b;
/* channel-mix */
float *ffn_x_k,*ffn_key_w,*ffn_value_w;
/* block LNs */
float *ln1_w,*ln1_b,*ln2_w,*ln2_b;
} BP;
static float frand(void){ return (float)rand()/(float)RAND_MAX - 0.5f; }
static float *rbuf(int n, float scale){
float *p=malloc((size_t)n*sizeof(float));
for (int i=0;i<n;i++) p[i]=scale*frand();
return p;
}
static void bp_alloc(BP *m, const cfg *cf){
int C=cf->C,Dw=cf->Dw,Da=cf->Da,Dg=cf->Dg,Dv=cf->Dv,F=cf->F;
m->x_r=rbuf(C,1.0f); m->x_k=rbuf(C,1.0f); m->x_v=rbuf(C,1.0f);
m->x_w=rbuf(C,1.0f); m->x_a=rbuf(C,1.0f); m->x_g=rbuf(C,1.0f);
m->receptance_w=rbuf(C*C,0.3f); m->key_w=rbuf(C*C,0.3f);
m->value_w=rbuf(C*C,0.3f); m->output_w=rbuf(C*C,0.3f);
m->w0=rbuf(C,0.5f); m->w1=rbuf(C*Dw,0.3f); m->w2=rbuf(Dw*C,0.3f);
m->a0=rbuf(C,0.5f); m->a1=rbuf(C*Da,0.3f); m->a2=rbuf(Da*C,0.3f);
m->g1=rbuf(C*Dg,0.3f); m->g2=rbuf(Dg*C,0.3f);
if (cf->layer_id>0){ m->v0=rbuf(C,0.5f); m->v1=rbuf(C*Dv,0.3f); m->v2=rbuf(Dv*C,0.3f); }
else { m->v0=m->v1=m->v2=NULL; }
m->k_k=rbuf(C,0.5f); m->k_a=rbuf(C,0.5f); m->r_k=rbuf(C,0.5f);
m->ln_x_w=malloc((size_t)C*sizeof(float)); m->ln_x_b=rbuf(C,0.2f);
for (int i=0;i<C;i++) m->ln_x_w[i]=1.0f+0.2f*frand();
/* channel-mix */
m->ffn_x_k=rbuf(C,1.0f);
m->ffn_key_w=rbuf(C*F,0.3f);
m->ffn_value_w=rbuf(F*C,0.3f);
/* block LNs: gamma near 1, beta small */
m->ln1_w=malloc((size_t)C*sizeof(float)); m->ln1_b=rbuf(C,0.2f);
m->ln2_w=malloc((size_t)C*sizeof(float)); m->ln2_b=rbuf(C,0.2f);
for (int i=0;i<C;i++){ m->ln1_w[i]=1.0f+0.2f*frand(); m->ln2_w[i]=1.0f+0.2f*frand(); }
}
static void bp_free(BP *m){
float **a=(float**)m; for (size_t i=0;i<sizeof(BP)/sizeof(float*);i++) free(a[i]);
}
static void bp_to_params(const BP *m, const cfg *cf, rx_block_params *P){
rx_tm_params *tm=&P->tm;
tm->c=cf->C; tm->n_head=cf->H; tm->head_size=cf->N;
tm->Dw=cf->Dw; tm->Da=cf->Da; tm->Dg=cf->Dg; tm->Dv=cf->Dv; tm->layer_id=cf->layer_id;
tm->x_r=m->x_r;tm->x_k=m->x_k;tm->x_v=m->x_v;tm->x_w=m->x_w;tm->x_a=m->x_a;tm->x_g=m->x_g;
tm->receptance_w=m->receptance_w;tm->key_w=m->key_w;tm->value_w=m->value_w;tm->output_w=m->output_w;
tm->w0=m->w0;tm->w1=m->w1;tm->w2=m->w2;tm->a0=m->a0;tm->a1=m->a1;tm->a2=m->a2;
tm->g1=m->g1;tm->g2=m->g2;tm->v0=m->v0;tm->v1=m->v1;tm->v2=m->v2;
tm->k_k=m->k_k;tm->k_a=m->k_a;tm->r_k=m->r_k;tm->ln_x_w=m->ln_x_w;tm->ln_x_b=m->ln_x_b;
P->cm.c=cf->C; P->cm.ffn_dim=cf->F;
P->cm.ffn_x_k=m->ffn_x_k; P->cm.ffn_key_w=m->ffn_key_w; P->cm.ffn_value_w=m->ffn_value_w;
P->ln1_w=m->ln1_w;P->ln1_b=m->ln1_b;P->ln2_w=m->ln2_w;P->ln2_b=m->ln2_b;
}
/* ---------- fp64 reference forward + scalar loss ---------- */
static double dsig(double z){ return 1.0/(1.0+exp(-z)); }
/* fp64 LayerNorm over C (eps=1e-5). */
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];
}
static double loss_fp64(const BP *m, const cfg *cf, const float *x,
const float *tmxp, const float *cmxp, const float *v_first, const float *cc)
{
int B=cf->B,T=cf->T,H=cf->H,N=cf->N,C=cf->C;
int Dw=cf->Dw,Da=cf->Da,Dg=cf->Dg,Dv=cf->Dv,F=cf->F;
double L=0.0;
size_t tot=(size_t)B*T*C;
double *h1=malloc(tot*sizeof(double)); /* ln1(x) */
double *tm=malloc(tot*sizeof(double)); /* time-mix out */
double *x1=malloc(tot*sizeof(double));
double *h2=malloc(tot*sizeof(double)); /* ln2(x1) */
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;
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)),*tg=malloc(C*sizeof(double));
double *xd=malloc(C*sizeof(double));
double *fk=malloc((size_t)F*sizeof(double));
/* h1 = ln1(x), per token */
for (int b=0;b<B;b++) for (int t=0;t<T;t++){
size_t base=((size_t)b*T+t)*C;
for(int i=0;i<C;i++) xd[i]=(double)x[base+i];
ln_fp64(xd,m->ln1_w,m->ln1_b,&h1[base],C);
}
/* time-mix over h1 -> tm (sequential state per batch) */
for (int b=0;b<B;b++){
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)b*T+t)*C;
const double *xx=&h1[base];
const double *prev=(t==0)?NULL:&h1[base-C];
double tmxp_v;
#define SHIFT(dst,lerp) for(int i=0;i<C;i++){ \
double pv = prev?prev[i] : (tmxp?(double)tmxp[(size_t)b*C+i]:0.0); \
dst[i]=xx[i]+(pv-xx[i])*(double)lerp[i]; } (void)tmxp_v
SHIFT(sh,m->x_r); for(int j=0;j<C;j++){double s=0;for(int i=0;i<C;i++)s+=sh[i]*m->receptance_w[i*C+j]; r[j]=s;}
SHIFT(sh,m->x_k); for(int j=0;j<C;j++){double s=0;for(int i=0;i<C;i++)s+=sh[i]*m->key_w[i*C+j]; k[j]=s;}
SHIFT(shv,m->x_v);
for(int j=0;j<C;j++){double s=0;for(int i=0;i<C;i++)s+=shv[i]*m->value_w[i*C+j]; v[j]=s;}
SHIFT(sh,m->x_w);
for(int d=0;d<Dw;d++){double s=0;for(int i=0;i<C;i++)s+=sh[i]*m->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]*m->w2[d*C+j];
double sg=dsig((double)m->w0[j]+s); w[j]=exp(-0.606531*sg);}
SHIFT(sh,m->x_a);
for(int d=0;d<Da;d++){double s=0;for(int i=0;i<C;i++)s+=sh[i]*m->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]*m->a2[d*C+j];
a[j]=dsig((double)m->a0[j]+s);}
SHIFT(sh,m->x_g);
for(int d=0;d<Dg;d++){double s=0;for(int i=0;i<C;i++)s+=sh[i]*m->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]*m->g2[d*C+j]; g[j]=s;}
for(int i=0;i<C;i++) kk[i]=k[i]*(double)m->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)m->k_a[i]);
if (cf->layer_id==0){ for(int i=0;i<C;i++) vf[i]=v[i]; }
else {
for(int d=0;d<Dv;d++){double s=0;for(int i=0;i<C;i++)s+=shv[i]*m->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]*m->v2[d*C+j];
double bl=dsig((double)m->v0[j]+s);
vf[j]=v[j]+((double)v_first[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)m->ln_x_w[off+i]+(double)m->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)m->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++) tg[i]=ygn[i]*g[i];
for(int j=0;j<C;j++){double s=0;for(int i=0;i<C;i++)s+=tg[i]*m->output_w[i*C+j]; tm[base+j]=s;}
#undef SHIFT
}
}
/* x1 = x + tm ; h2 = ln2(x1) */
for (int b=0;b<B;b++) for (int t=0;t<T;t++){
size_t base=((size_t)b*T+t)*C;
for(int i=0;i<C;i++) x1[base+i]=(double)x[base+i]+tm[base+i];
ln_fp64(&x1[base],m->ln2_w,m->ln2_b,&h2[base],C);
}
/* channel-mix over h2 -> cm ; out = x1 + cm ; L += cc*out */
for (int b=0;b<B;b++) for (int t=0;t<T;t++){
size_t base=((size_t)b*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]:(cmxp?(double)cmxp[(size_t)b*C+i]:0.0);
sh[i]=xx[i]+(pv-xx[i])*(double)m->ffn_x_k[i];
}
for(int mm=0;mm<F;mm++){double s=0;for(int i=0;i<C;i++)s+=sh[i]*m->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]*m->ffn_value_w[mm*C+j];
double out=x1[base+j]+s;
L+=(double)cc[base+j]*out;}
}
free(h1);free(tm);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(tg);free(xd);free(fk);
return L;
}
/* FD-check an array (nv elems) vs analytic grad g — per-tensor rel-err. */
static float fd_array(float *arr, const float *g, int nv,
const BP *m, const cfg *cf, const float *x, const float *tmxp,
const float *cmxp, const float *vfst, const float *cc,
const char *nm, int verbose)
{
float maxdiff=0.0f, maxfd=0.0f;
for (int e=0;e<nv;e++){
float o=arr[e]; float h=1e-3f*(1.0f+fabsf(o));
arr[e]=o+h; double Lp=loss_fp64(m,cf,x,tmxp,cmxp,vfst,cc);
arr[e]=o-h; double Lm=loss_fp64(m,cf,x,tmxp,cmxp,vfst,cc);
arr[e]=o;
float fd=(float)((Lp-Lm)/(2.0*(double)h));
float d=fabsf(g[e]-fd);
if (d>maxdiff) maxdiff=d;
if (fabsf(fd)>maxfd) maxfd=fabsf(fd);
}
float re=maxdiff/(maxfd+1e-9f);
if (verbose) printf(" %-14s rel=%.2e (max|g-fd|=%.2e max|fd|=%.2e)\n",
nm,re,maxdiff,maxfd);
return re;
}
static int run_case(const char *name, cfg cf, unsigned seed){
srand(seed);
cf.C = cf.H*cf.N;
int C=cf.C, B=cf.B, T=cf.T;
size_t tot=(size_t)B*T*C;
BP m; bp_alloc(&m,&cf);
rx_block_params P; bp_to_params(&m,&cf,&P);
float *x=rbuf((int)tot,1.0f);
float *vfst=rbuf((int)tot,0.6f); /* external v_first (layer>0) */
float *cc=rbuf((int)tot,1.0f);
float *dy=malloc(tot*sizeof(float));
for (size_t i=0;i<tot;i++) dy[i]=cc[i]; /* L = sum cc*out => dy = cc */
/* analytic grads */
BP gm; memset(&gm,0,sizeof gm);
#define GA(f,n) gm.f=calloc((size_t)(n),sizeof(float))
GA(x_r,C);GA(x_k,C);GA(x_v,C);GA(x_w,C);GA(x_a,C);GA(x_g,C);
GA(receptance_w,C*C);GA(key_w,C*C);GA(value_w,C*C);GA(output_w,C*C);
GA(w0,C);GA(w1,C*cf.Dw);GA(w2,cf.Dw*C);
GA(a0,C);GA(a1,C*cf.Da);GA(a2,cf.Da*C);
GA(g1,C*cf.Dg);GA(g2,cf.Dg*C);
if (cf.layer_id>0){ GA(v0,C);GA(v1,C*cf.Dv);GA(v2,cf.Dv*C); }
GA(k_k,C);GA(k_a,C);GA(r_k,C);GA(ln_x_w,C);GA(ln_x_b,C);
GA(ffn_x_k,C);GA(ffn_key_w,C*cf.F);GA(ffn_value_w,cf.F*C);
GA(ln1_w,C);GA(ln1_b,C);GA(ln2_w,C);GA(ln2_b,C);
#undef GA
rx_block_grads G; memset(&G,0,sizeof G);
rx_tm_grads *tg=&G.tm;
tg->x_r=gm.x_r;tg->x_k=gm.x_k;tg->x_v=gm.x_v;tg->x_w=gm.x_w;tg->x_a=gm.x_a;tg->x_g=gm.x_g;
tg->receptance_w=gm.receptance_w;tg->key_w=gm.key_w;tg->value_w=gm.value_w;tg->output_w=gm.output_w;
tg->w0=gm.w0;tg->w1=gm.w1;tg->w2=gm.w2;tg->a0=gm.a0;tg->a1=gm.a1;tg->a2=gm.a2;
tg->g1=gm.g1;tg->g2=gm.g2;tg->v0=gm.v0;tg->v1=gm.v1;tg->v2=gm.v2;
tg->k_k=gm.k_k;tg->k_a=gm.k_a;tg->r_k=gm.r_k;tg->ln_x_w=gm.ln_x_w;tg->ln_x_b=gm.ln_x_b;
G.cm.ffn_x_k=gm.ffn_x_k;G.cm.ffn_key_w=gm.ffn_key_w;G.cm.ffn_value_w=gm.ffn_value_w;
G.ln1_w=gm.ln1_w;G.ln1_b=gm.ln1_b;G.ln2_w=gm.ln2_w;G.ln2_b=gm.ln2_b;
float *dx=calloc(tot,sizeof(float));
float *dvf=cf.layer_id>0?calloc(tot,sizeof(float)):NULL;
rx_block_backward(&P,x,NULL,NULL,vfst,dy,&G,dx,NULL,NULL,dvf,B,T);
printf(" [%s] B=%d T=%d H=%d N=%d C=%d F=%d layer=%d\n",
name,B,T,cf.H,cf.N,C,cf.F,cf.layer_id);
float wmax=0.0f;
#define CK(f,n) do{ float e=fd_array(m.f,gm.f,(n),&m,&cf,x,NULL,NULL,vfst,cc,#f,1); \
if(e>wmax)wmax=e; }while(0)
/* channel-mix + block LNs (the new pieces) */
CK(ffn_x_k,C);CK(ffn_key_w,C*cf.F);CK(ffn_value_w,cf.F*C);
CK(ln1_w,C);CK(ln1_b,C);CK(ln2_w,C);CK(ln2_b,C);
/* time-mix params (re-checked end-to-end through the block) */
CK(x_r,C);CK(x_k,C);CK(x_v,C);CK(x_w,C);CK(x_a,C);CK(x_g,C);
CK(receptance_w,C*C);CK(key_w,C*C);CK(value_w,C*C);CK(output_w,C*C);
CK(w0,C);CK(w1,C*cf.Dw);CK(w2,cf.Dw*C);
CK(a0,C);CK(a1,C*cf.Da);CK(a2,cf.Da*C);
CK(g1,C*cf.Dg);CK(g2,cf.Dg*C);
if (cf.layer_id>0){ CK(v0,C);CK(v1,C*cf.Dv);CK(v2,cf.Dv*C); }
CK(k_k,C);CK(k_a,C);CK(r_k,C);CK(ln_x_w,C);CK(ln_x_b,C);
#undef CK
{ float e=fd_array(x,dx,(int)tot,&m,&cf,x,NULL,NULL,vfst,cc,"dx(block in)",1);
if(e>wmax)wmax=e; }
printf(" -> worst rel-err = %.2e (%s)\n", wmax, wmax<1e-3f?"PASS":"FAIL");
bp_free(&m); bp_free(&gm);
free(x);free(vfst);free(cc);free(dy);free(dx);free(dvf);
return wmax<1e-3f;
}
static void test_block_backward(void){
printf(" RWKV-7 full block backward — fp64 finite differences:\n");
int ok=1;
ok &= run_case("layer0 base", (cfg){.B=1,.T=5,.H=2,.N=4,.Dw=3,.Da=3,.Dg=3,.Dv=0,.layer_id=0,.F=12}, 11);
ok &= run_case("layer1 vres", (cfg){.B=1,.T=6,.H=2,.N=4,.Dw=3,.Da=3,.Dg=3,.Dv=3,.layer_id=1,.F=16}, 23);
ok &= run_case("batch heads", (cfg){.B=2,.T=4,.H=3,.N=6,.Dw=4,.Da=3,.Dg=5,.Dv=0,.layer_id=0,.F=20}, 37);
ASSERT_TRUE(ok);
}
int main(void){
printf("\nRWKV-7 full block backward gradient check\n");
RUN_TEST(block_backward);
TEST_REPORT();
return __ft_failed>0 ? 1 : 0;
}
Reference in New Issue View Git Blame Copy Permalink