274 lines
13 KiB
C
274 lines
13 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
||
/* Copyright (C) 2026 David (Corey) Herrera */
|
||
/*
|
||
* test_timemix_bwd.c — Gradient-check for the RWKV-7 TIME-MIX BLOCK backward
|
||
* (rx_timemix_backward), the counterpart of rwkvx.c `rwkv7_time_mix_fast`.
|
||
*
|
||
* Method (docs/wkv7_backward_reference.md §5/§6, "use an fp64 reference"):
|
||
* - A DOUBLE-PRECISION mirror of the time-mix forward computes the FD
|
||
* reference loss L = sum(cc * out). fp32 forward noise would swamp the
|
||
* small a/b/decay-related grads, so the reference forward is fp64.
|
||
* - The analytic backward runs in the normal fp32 path (rx_timemix_backward).
|
||
* - Central finite differences per-parameter and for the block input.
|
||
* Pass: max relative error < 1e-3 for EVERY parameter grad and the input.
|
||
*
|
||
* Prints per-parameter rel-err so a miss is localised. Several shapes are
|
||
* exercised (T, heads, head-dim, batch) plus a layer>0 case (value-residual
|
||
* gating path) and a non-trivial T so token-shift + the WKV recurrence run.
|
||
*/
|
||
#include "lib/test.h"
|
||
#include "rwkvx/rx_pretrain.h"
|
||
|
||
#include <math.h>
|
||
#include <stdlib.h>
|
||
#include <string.h>
|
||
|
||
/* ---------- config + parameter bag ---------- */
|
||
typedef struct {
|
||
int B, T, H, N, C, Dw, Da, Dg, Dv, layer_id;
|
||
} cfg;
|
||
|
||
/* Every learnable time-mix tensor, fp32. Sizes derived from cfg. */
|
||
typedef struct {
|
||
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;
|
||
} TMP;
|
||
|
||
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 tmp_alloc(TMP *m, const cfg *cf){
|
||
int C=cf->C,Dw=cf->Dw,Da=cf->Da,Dg=cf->Dg,Dv=cf->Dv;
|
||
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);
|
||
/* group-norm affine: weight near 1, bias small */
|
||
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();
|
||
}
|
||
static void tmp_free(TMP *m){
|
||
float **a=(float**)m; for (size_t i=0;i<sizeof(TMP)/sizeof(float*);i++) free(a[i]);
|
||
}
|
||
static void tmp_to_params(const TMP *m, const cfg *cf, rx_tm_params *P){
|
||
P->c=cf->C; P->n_head=cf->H; P->head_size=cf->N;
|
||
P->Dw=cf->Dw; P->Da=cf->Da; P->Dg=cf->Dg; P->Dv=cf->Dv; P->layer_id=cf->layer_id;
|
||
P->x_r=m->x_r;P->x_k=m->x_k;P->x_v=m->x_v;P->x_w=m->x_w;P->x_a=m->x_a;P->x_g=m->x_g;
|
||
P->receptance_w=m->receptance_w;P->key_w=m->key_w;P->value_w=m->value_w;P->output_w=m->output_w;
|
||
P->w0=m->w0;P->w1=m->w1;P->w2=m->w2;P->a0=m->a0;P->a1=m->a1;P->a2=m->a2;
|
||
P->g1=m->g1;P->g2=m->g2;P->v0=m->v0;P->v1=m->v1;P->v2=m->v2;
|
||
P->k_k=m->k_k;P->k_a=m->k_a;P->r_k=m->r_k;P->ln_x_w=m->ln_x_w;P->ln_x_b=m->ln_x_b;
|
||
}
|
||
|
||
/* ---------- fp64 reference forward + scalar loss ---------- */
|
||
static double dsig(double z){ return 1.0/(1.0+exp(-z)); }
|
||
|
||
static double loss_fp64(const TMP *m, const cfg *cf, const float *x,
|
||
const float *x_prev_init, 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;
|
||
double L=0.0;
|
||
double *S=calloc((size_t)H*N*N,sizeof(double));
|
||
double *sh=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));
|
||
double *lo=malloc((size_t)(Dw>Da?(Dw>Dg?Dw:Dg):(Da>Dg?Da:Dg))*sizeof(double)+64);
|
||
double *lo2=malloc((size_t)(Dv>0?Dv:1)*sizeof(double)+64);
|
||
double *y=malloc(C*sizeof(double)),*ygn=malloc(C*sizeof(double)),*tg=malloc(C*sizeof(double));
|
||
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 float *xx=&x[base];
|
||
const float *prev=(t==0)?(x_prev_init?&x_prev_init[(size_t)b*C]:NULL)
|
||
:&x[(((size_t)b*T)+(t-1))*C];
|
||
#define SHIFT(dst,lerp) for(int i=0;i<C;i++){ double pv=prev?prev[i]:0.0; \
|
||
dst[i]=(double)xx[i]+(pv-(double)xx[i])*(double)lerp[i]; }
|
||
/* R/K/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;}
|
||
double shv[4096]; 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;}
|
||
/* decay */
|
||
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);}
|
||
/* gate a */
|
||
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);}
|
||
/* gate g */
|
||
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;}
|
||
/* key norm (per head) */
|
||
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;}
|
||
/* key mod */
|
||
for(int i=0;i<C;i++) kmod[i]=k[i]*(1.0+(a[i]-1.0)*(double)m->k_a[i]);
|
||
/* value residual */
|
||
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;}
|
||
}
|
||
/* WKV7 recurrence (double) */
|
||
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;
|
||
}
|
||
}
|
||
/* group norm */
|
||
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];}
|
||
/* bonus */
|
||
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];}
|
||
/* gate + output */
|
||
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];
|
||
L+=(double)cc[base+j]*s;}
|
||
#undef SHIFT
|
||
}
|
||
}
|
||
free(S);free(sh);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);
|
||
return L;
|
||
}
|
||
|
||
/* FD-check a parameter/input array (nv elements) vs analytic grad g. */
|
||
static float fd_array(float *arr, const float *g, int nv,
|
||
const TMP *m, const cfg *cf, const float *x, const float *xpi,
|
||
const float *vfst, const float *cc, const char *nm, int verbose)
|
||
{
|
||
/* Per-tensor relative error (standard gradcheck metric): the max absolute
|
||
* analytic-vs-FD mismatch, normalised by the tensor's gradient scale. A
|
||
* per-ELEMENT relative error is meaningless for the many near-zero entries
|
||
* of the C×C weight matrices (a ~1e-7 absolute mismatch over a ~1e-5 entry
|
||
* reads as ~1e-2 while g and fd still agree to 3-4 sig figs). */
|
||
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,xpi,vfst,cc);
|
||
arr[e]=o-h; double Lm=loss_fp64(m,cf,x,xpi,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);
|
||
int C=cf.C, B=cf.B, T=cf.T;
|
||
cf.C = C = cf.H*cf.N;
|
||
size_t tot=(size_t)B*T*C;
|
||
|
||
TMP m; tmp_alloc(&m,&cf);
|
||
rx_tm_params P; tmp_to_params(&m,&cf,&P);
|
||
|
||
float *x=rbuf((int)tot,1.0f);
|
||
float *xpi=rbuf(B*C,1.0f);
|
||
float *vfst=rbuf((int)tot,0.6f); /* external v_first (layer>0) */
|
||
float *cc=rbuf((int)tot,1.0f); /* dL/dout weights */
|
||
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 */
|
||
TMP 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);
|
||
#undef GA
|
||
rx_tm_grads G;
|
||
G.x_r=gm.x_r;G.x_k=gm.x_k;G.x_v=gm.x_v;G.x_w=gm.x_w;G.x_a=gm.x_a;G.x_g=gm.x_g;
|
||
G.receptance_w=gm.receptance_w;G.key_w=gm.key_w;G.value_w=gm.value_w;G.output_w=gm.output_w;
|
||
G.w0=gm.w0;G.w1=gm.w1;G.w2=gm.w2;G.a0=gm.a0;G.a1=gm.a1;G.a2=gm.a2;
|
||
G.g1=gm.g1;G.g2=gm.g2;G.v0=gm.v0;G.v1=gm.v1;G.v2=gm.v2;
|
||
G.k_k=gm.k_k;G.k_a=gm.k_a;G.r_k=gm.r_k;G.ln_x_w=gm.ln_x_w;G.ln_x_b=gm.ln_x_b;
|
||
|
||
float *dx=calloc(tot,sizeof(float));
|
||
float *dvf=cf.layer_id>0?calloc(tot,sizeof(float)):NULL;
|
||
rx_timemix_backward(&P,x,xpi,vfst,dy,&G,dx,NULL,dvf,B,T);
|
||
|
||
printf(" [%s] B=%d T=%d H=%d N=%d C=%d layer=%d\n",
|
||
name,B,T,cf.H,cf.N,C,cf.layer_id);
|
||
float wmax=0.0f;
|
||
#define CK(f,n) do{ float e=fd_array(m.f,gm.f,(n),&m,&cf,x,xpi,vfst,cc,#f,1); \
|
||
if(e>wmax)wmax=e; }while(0)
|
||
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
|
||
/* block input */
|
||
{ float e=fd_array(x,dx,(int)tot,&m,&cf,x,xpi,vfst,cc,"dx(input)",1);
|
||
if(e>wmax)wmax=e; }
|
||
printf(" -> worst rel-err = %.2e (%s)\n", wmax, wmax<1e-3f?"PASS":"FAIL");
|
||
|
||
/* cleanup */
|
||
tmp_free(&m); tmp_free(&gm);
|
||
free(x);free(xpi);free(vfst);free(cc);free(dy);free(dx); free(dvf);
|
||
return wmax<1e-3f;
|
||
}
|
||
|
||
static void test_timemix_backward(void){
|
||
printf(" Time-mix block backward — fp64 finite differences:\n");
|
||
int ok=1;
|
||
/* layer 0, small multi-head, non-trivial T */
|
||
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}, 11);
|
||
/* layer>0: exercises value-residual gating path */
|
||
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}, 23);
|
||
/* batch>1, more heads, bigger head-dim */
|
||
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}, 37);
|
||
ASSERT_TRUE(ok);
|
||
}
|
||
|
||
int main(void){
|
||
printf("\nRWKV-7 time-mix block backward gradient check\n");
|
||
RUN_TEST(timemix_backward);
|
||
TEST_REPORT();
|
||
return __ft_failed>0 ? 1 : 0;
|
||
}
|