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

232 lines
9.5 KiB
C
Raw Blame History

/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
/*
* test_wkv7_bwd.c — Gradient-check for the RWKV-7 WKV7 backward-through-time.
*
* Two independent gates (docs/wkv7_backward_reference.md §6):
*
* A. Parity vs the verified fp64 PyTorch oracle. The fixed-seed fixtures in
* wkv7_fixture.h (generated by gen_wkv7_fixture.py from wkv7_reference.py)
* carry inputs, y, and reference grads dr,dw,dk,dv,da,db. We run the C
* forward + backward on the SAME inputs and diff. Pass: max rel-err<1e-3
* on every grad (and on y).
*
* B. Central finite differences in pure C (self-contained, no torch). For a
* scalar loss L=sum(c*y) with fixed random c, FD-check every input
* r,w,k,v,a,b AND the initial state carry S_{-1} (-> dS) independently.
* Pass: max rel-err<1e-3 for all seven.
*
* Prints per-grad relative error for both gates.
*/
#include "lib/test.h"
#include "rwkvx/rx_pretrain.h"
#include "rwkvx/tests/wkv7_fixture.h"
#include <math.h>
#include <stdlib.h>
#include <string.h>
/* max relative error: max_e |g[e]-ref[e]| / (|ref[e]| + eps) */
static float rel_err(const float *g, const float *ref, int n) {
float worst = 0.0f;
for (int e = 0; e < n; e++) {
float d = fabsf(g[e] - ref[e]);
float r = d / (fabsf(ref[e]) + 1e-6f);
if (r > worst) worst = r;
}
return worst;
}
/* ================================================================
* Gate A: parity vs oracle fixture
* ================================================================ */
static int gate_a_case(const char *name, int B, int T, int H, int N,
const float *r, const float *w, const float *k, const float *v,
const float *a, const float *b, const float *dy, const float *y_ref,
const float *dr_ref, const float *dw_ref, const float *dk_ref,
const float *dv_ref, const float *da_ref, const float *db_ref)
{
int nv = B * T * H * N;
float *y = (float *)malloc(nv * sizeof(float));
float *dr = (float *)malloc(nv * sizeof(float));
float *dw = (float *)malloc(nv * sizeof(float));
float *dk = (float *)malloc(nv * sizeof(float));
float *dv = (float *)malloc(nv * sizeof(float));
float *da = (float *)malloc(nv * sizeof(float));
float *db = (float *)malloc(nv * sizeof(float));
rx_wkv7_forward_seq(r, w, k, v, a, b, NULL, y, NULL, B, T, H, N);
rx_wkv7_backward(r, w, k, v, a, b, dy, NULL,
dr, dw, dk, dv, da, db, NULL, B, T, H, N);
float ey = rel_err(y, y_ref, nv);
float edr = rel_err(dr, dr_ref, nv);
float edw = rel_err(dw, dw_ref, nv);
float edk = rel_err(dk, dk_ref, nv);
float edv = rel_err(dv, dv_ref, nv);
float eda = rel_err(da, da_ref, nv);
float edb = rel_err(db, db_ref, nv);
printf(" [A %-10s B=%d T=%2d H=%d N=%2d] y=%.2e | dr=%.2e dw=%.2e "
"dk=%.2e dv=%.2e da=%.2e db=%.2e\n",
name, B, T, H, N, ey, edr, edw, edk, edv, eda, edb);
free(y); free(dr); free(dw); free(dk); free(dv); free(da); free(db);
int ok = (ey < 1e-3f) && (edr < 1e-3f) && (edw < 1e-3f) && (edk < 1e-3f)
&& (edv < 1e-3f) && (eda < 1e-3f) && (edb < 1e-3f);
return ok;
}
static void test_gate_a_oracle_parity(void) {
printf(" Gate A — parity vs verified fp64 oracle fixture:\n");
int ok = 1;
ok &= gate_a_case("canonical", C0_B, C0_T, C0_H, C0_N,
c0_r, c0_w, c0_k, c0_v, c0_a, c0_b, c0_dy, c0_y,
c0_dr, c0_dw, c0_dk, c0_dv, c0_da, c0_db);
ok &= gate_a_case("multihead", C1_B, C1_T, C1_H, C1_N,
c1_r, c1_w, c1_k, c1_v, c1_a, c1_b, c1_dy, c1_y,
c1_dr, c1_dw, c1_dk, c1_dv, c1_da, c1_db);
ok &= gate_a_case("batched", C2_B, C2_T, C2_H, C2_N,
c2_r, c2_w, c2_k, c2_v, c2_a, c2_b, c2_dy, c2_y,
c2_dr, c2_dw, c2_dk, c2_dv, c2_da, c2_db);
ASSERT_TRUE(ok);
}
/* ================================================================
* Gate B: pure-C central finite differences (incl. dS carry)
* ================================================================ */
/* fp64 reference forward, mirroring rx_wkv7_forward_seq exactly, so the FD
* reference loss is free of fp32 forward noise (otherwise single-element
* perturbations of small inputs like a/b lose precision to cancellation).
* Loss L = sum(c*y). */
static double loss_of(const float *r, const float *w, const float *k,
const float *v, const float *a, const float *b,
const float *s_init, const float *c,
int B, int T, int H, int N)
{
double *S = (double *)malloc((size_t)B * H * N * N * sizeof(double));
for (int e = 0; e < B*H*N*N; e++) S[e] = s_init ? (double)s_init[e] : 0.0;
double L = 0.0;
for (int bch = 0; bch < B; bch++)
for (int h = 0; h < H; h++)
for (int t = 0; t < T; t++)
for (int i = 0; i < N; i++) {
int sbase = ((bch*H + h)*N + i)*N;
double sa = 0.0;
for (int j = 0; j < N; j++)
sa += (double)a[(((bch*T+t)*H+h)*N)+j] * S[sbase+j];
double vi = (double)v[(((bch*T+t)*H+h)*N)+i];
double yy = 0.0;
for (int j = 0; j < N; j++) {
int vj = (((bch*T+t)*H+h)*N)+j;
S[sbase+j] = S[sbase+j]*(double)w[vj]
+ (double)k[vj]*vi
+ sa*(double)b[vj];
yy += S[sbase+j] * (double)r[vj];
}
L += (double)c[(((bch*T+t)*H+h)*N)+i] * yy;
}
free(S);
return L;
}
/* FD-check one input tensor `x` (nv elements) against analytic grad `g`. */
static float fd_check_input(float *x, const float *g, int nv,
const float *r, const float *w, const float *k, const float *v,
const float *a, const float *b, const float *s_init, const float *c,
int B, int T, int H, int N)
{
float worst = 0.0f;
for (int e = 0; e < nv; e++) {
float orig = x[e];
float hstep = 1e-3f * (1.0f + fabsf(orig));
x[e] = orig + hstep;
double Lp = loss_of(r, w, k, v, a, b, s_init, c, B, T, H, N);
x[e] = orig - hstep;
double Lm = loss_of(r, w, k, v, a, b, s_init, c, B, T, H, N);
x[e] = orig;
float fd = (float)((Lp - Lm) / (2.0 * (double)hstep));
float rerr = fabsf(g[e] - fd) / (fabsf(fd) + 1e-6f);
if (rerr > worst) worst = rerr;
}
return worst;
}
/* FD-check the initial state carry S_{-1} against dS_out (nvs elements). */
static float fd_check_state(float *s_init, const float *dS, int nvs,
const float *r, const float *w, const float *k, const float *v,
const float *a, const float *b, const float *c,
int B, int T, int H, int N)
{
float worst = 0.0f;
for (int e = 0; e < nvs; e++) {
float orig = s_init[e];
float hstep = 1e-3f * (1.0f + fabsf(orig));
s_init[e] = orig + hstep;
double Lp = loss_of(r, w, k, v, a, b, s_init, c, B, T, H, N);
s_init[e] = orig - hstep;
double Lm = loss_of(r, w, k, v, a, b, s_init, c, B, T, H, N);
s_init[e] = orig;
float fd = (float)((Lp - Lm) / (2.0 * (double)hstep));
float rerr = fabsf(dS[e] - fd) / (fabsf(fd) + 1e-6f);
if (rerr > worst) worst = rerr;
}
return worst;
}
static float frand(void) { return (float)rand() / (float)RAND_MAX - 0.5f; }
static void test_gate_b_finite_diff(void) {
printf(" Gate B — central finite differences in C (incl. dS carry):\n");
int B = 1, T = 6, H = 2, N = 6;
int nv = B * T * H * N, nvs = B * H * N * N;
srand(1234);
float *r = malloc(nv*4), *w = malloc(nv*4), *k = malloc(nv*4);
float *v = malloc(nv*4), *a = malloc(nv*4), *b = malloc(nv*4);
float *c = malloc(nv*4), *s_init = malloc(nvs*4);
for (int e = 0; e < nv; e++) {
r[e] = 0.5f*frand(); k[e] = 0.5f*frand(); v[e] = 0.5f*frand();
a[e] = 0.5f*frand(); b[e] = 0.5f*frand(); c[e] = frand();
/* effective decay in (0,1) */
w[e] = 0.3f + 0.6f*((float)rand()/(float)RAND_MAX);
}
/* nonzero initial state carry -> genuinely exercises the reverse scan */
for (int e = 0; e < nvs; e++) s_init[e] = 0.3f*frand();
/* analytic grads: dy = c (since L = sum c*y) */
float *dr = malloc(nv*4), *dw = malloc(nv*4), *dk = malloc(nv*4);
float *dv = malloc(nv*4), *da = malloc(nv*4), *db = malloc(nv*4);
float *dS = malloc(nvs*4);
rx_wkv7_backward(r, w, k, v, a, b, c, s_init,
dr, dw, dk, dv, da, db, dS, B, T, H, N);
float er = fd_check_input(r, dr, nv, r,w,k,v,a,b,s_init,c, B,T,H,N);
float ek = fd_check_input(k, dk, nv, r,w,k,v,a,b,s_init,c, B,T,H,N);
float ev = fd_check_input(v, dv, nv, r,w,k,v,a,b,s_init,c, B,T,H,N);
float ew = fd_check_input(w, dw, nv, r,w,k,v,a,b,s_init,c, B,T,H,N);
float eb = fd_check_input(b, db, nv, r,w,k,v,a,b,s_init,c, B,T,H,N);
float ea = fd_check_input(a, da, nv, r,w,k,v,a,b,s_init,c, B,T,H,N);
float es = fd_check_state(s_init, dS, nvs, r,w,k,v,a,b,c, B,T,H,N);
printf(" [B fd B=%d T=%2d H=%d N=%2d] dr=%.2e dk=%.2e dv=%.2e "
"dw=%.2e db=%.2e da=%.2e | dS=%.2e\n",
B, T, H, N, er, ek, ev, ew, eb, ea, es);
int ok = (er<1e-3f)&&(ek<1e-3f)&&(ev<1e-3f)&&(ew<1e-3f)
&&(eb<1e-3f)&&(ea<1e-3f)&&(es<1e-3f);
free(r);free(w);free(k);free(v);free(a);free(b);free(c);free(s_init);
free(dr);free(dw);free(dk);free(dv);free(da);free(db);free(dS);
ASSERT_TRUE(ok);
}
int main(void) {
printf("\nWKV7 backward-through-time gradient check\n");
RUN_TEST(gate_a_oracle_parity);
RUN_TEST(gate_b_finite_diff);
TEST_REPORT();
return __ft_failed > 0 ? 1 : 0;
}
Reference in New Issue View Git Blame Copy Permalink