138 lines
6.5 KiB
Plaintext
138 lines
6.5 KiB
Plaintext
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/*
|
|
* test_fullmodel_cuda.cu — Phase-2 completion gate. Parity of the WHOLE
|
|
* device-resident RWKV-x training step (cuda_kernels/fullmodel.cu) vs the CPU
|
|
* oracle rx_pretrain_forward / rx_pretrain_backward (rwkvx/rx_pretrain.c) on
|
|
* GPU 0: embedding + interleaved RWKV-7 / MoBA-attention blocks + ln_out + LM
|
|
* head + cross-entropy, forward AND backward.
|
|
*
|
|
* Checks the scalar loss AND every parameter gradient (embedding, every block
|
|
* param, head) at inf-norm rel-err = max|gpu-cpu|/max|cpu| < 1e-3. Then measures
|
|
* the full fwd+bwd training-step throughput with WEIGHTS RESIDENT on-device
|
|
* (uploaded once, loop without re-uploading) — tokens/sec on one K80 die.
|
|
*/
|
|
#include <cuda_runtime.h>
|
|
#include <cublas_v2.h>
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <stdint.h>
|
|
#include <string.h>
|
|
#include <math.h>
|
|
#include <time.h>
|
|
|
|
#include "cuda_kernels/fullmodel.cu"
|
|
|
|
static uint64_t sm_state;
|
|
static uint64_t sm_next(void){ uint64_t z=(sm_state+=0x9E3779B97F4A7C15ULL);
|
|
z=(z^(z>>30))*0xBF58476D1CE4E5B9ULL; z=(z^(z>>27))*0x94D049BB133111EBULL; return z^(z>>31); }
|
|
static int sm_tok(int V){ return (int)(sm_next()%(uint64_t)V); }
|
|
|
|
static int g_fail=0;
|
|
static double chk(const char*tname,const float*gpu,const float*cpu,size_t n,double bar,int verbose){
|
|
double maxabs=0.0,scale=0.0;
|
|
for(size_t i=0;i<n;i++){ double d=fabs((double)gpu[i]-(double)cpu[i]); if(d>maxabs)maxabs=d;
|
|
double c=fabs((double)cpu[i]); if(c>scale)scale=c; }
|
|
double rel=maxabs/(scale+1e-30); const char*st=(rel<bar)?"PASS":"FAIL"; if(rel>=bar)g_fail=1;
|
|
if(verbose||rel>=bar)
|
|
printf(" %-22s rel=%.3e maxabs=%.3e scale=%.3e [%s]\n",tname,rel,maxabs,scale,st);
|
|
return rel;
|
|
}
|
|
|
|
static void run_parity(const char*name, rx_pt_config pc, int T, uint64_t seed, int ignore_index){
|
|
printf("\n=== parity [%s] C=%d H=%d N=%d rwkv=%d attn=%d V=%d ffn=%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,T,pc.chunk_size,pc.moba_top_k);
|
|
rx_ptmodel*cm=rx_model_new_random(&pc,seed);
|
|
if(!cm){ printf(" model build FAIL\n"); g_fail=1; return; }
|
|
int*tokens=(int*)malloc((size_t)T*sizeof(int)),*targets=(int*)malloc((size_t)T*sizeof(int));
|
|
sm_state=seed^0xABCDEF;
|
|
for(int t=0;t<T;t++){ tokens[t]=sm_tok(pc.vocab_size); targets[t]=sm_tok(pc.vocab_size); }
|
|
|
|
rx_pt_acts*acts=rx_pt_acts_new(cm,T);
|
|
if(rx_pretrain_forward(cm,tokens,T,acts)){ printf(" CPU fwd FAIL\n"); g_fail=1; return; }
|
|
double loss_c=rx_pretrain_backward(cm,acts,targets,T,ignore_index);
|
|
|
|
fm_model*m=fm_build(cm,T);
|
|
fm_forward(m,tokens,T);
|
|
double loss_g=fm_backward(m,targets,T,ignore_index);
|
|
|
|
double lrel=fabs(loss_g-loss_c)/(fabs(loss_c)+1e-30);
|
|
printf(" loss cpu=%.8f gpu=%.8f rel=%.3e [%s]\n",loss_c,loss_g,lrel,(lrel<1e-3)?"PASS":"FAIL");
|
|
if(lrel>=1e-3) g_fail=1;
|
|
|
|
/* compare every param grad by name */
|
|
double worst=0.0; char worstn[48]="";
|
|
int np=rx_ptmodel_num_params(cm);
|
|
for(int i=0;i<np;i++){
|
|
const char*nm; float*data,*grad; int n;
|
|
rx_ptmodel_param(cm,i,&nm,&data,&grad,&n);
|
|
int gn; float*dg=fm_gfind(m,nm,&gn);
|
|
if(!dg){ printf(" %-22s NO GPU GRAD (name mismatch!)\n",nm); g_fail=1; continue; }
|
|
if(gn!=n){ printf(" %-22s SIZE MISMATCH cpu=%d gpu=%d\n",nm,n,gn); g_fail=1; continue; }
|
|
float*hg=(float*)malloc((size_t)n*sizeof(float));
|
|
cudaMemcpy(hg,dg,(size_t)n*sizeof(float),cudaMemcpyDeviceToHost);
|
|
double rel=chk(nm,hg,grad,(size_t)n,1e-3,0);
|
|
if(rel>worst){ worst=rel; snprintf(worstn,48,"%s",nm); }
|
|
free(hg);
|
|
}
|
|
printf(" worst param-grad rel=%.3e (%s) [%d params checked]\n",worst,worstn,np);
|
|
|
|
free(tokens);free(targets);rx_pt_acts_free(acts);rx_ptmodel_free(cm);
|
|
/* (device model left allocated; process exits soon) */
|
|
}
|
|
|
|
static void run_throughput(rx_pt_config pc, int T, int iters, uint64_t seed){
|
|
printf("\n=== throughput (full fwd+bwd, WEIGHTS RESIDENT, one K80 die) ===\n");
|
|
printf(" config: C=%d H=%d N=%d rwkv=%d attn=%d V=%d ffn=%d T=%d B=1 chunk=%d topk=%d\n",
|
|
pc.n_embd,pc.n_head,pc.head_size,pc.n_rwkv_layer,pc.n_attn_layer,pc.vocab_size,
|
|
pc.ffn_dim,T,pc.chunk_size,pc.moba_top_k);
|
|
rx_ptmodel*cm=rx_model_new_random(&pc,seed);
|
|
int*tokens=(int*)malloc((size_t)T*sizeof(int)),*targets=(int*)malloc((size_t)T*sizeof(int));
|
|
sm_state=seed;
|
|
for(int t=0;t<T;t++){ tokens[t]=sm_tok(pc.vocab_size); targets[t]=sm_tok(pc.vocab_size); }
|
|
fm_model*m=fm_build(cm,T); /* upload ALL weights ONCE (resident) */
|
|
/* warmup */
|
|
fm_forward(m,tokens,T); fm_backward(m,targets,T,-1);
|
|
cudaDeviceSynchronize();
|
|
struct timespec t0,t1; clock_gettime(CLOCK_MONOTONIC,&t0);
|
|
for(int it=0;it<iters;it++){ /* NO re-upload of weights inside the loop */
|
|
fm_forward(m,tokens,T);
|
|
fm_backward(m,targets,T,-1);
|
|
}
|
|
cudaDeviceSynchronize(); clock_gettime(CLOCK_MONOTONIC,&t1);
|
|
double sec=(t1.tv_sec-t0.tv_sec)+(t1.tv_nsec-t0.tv_nsec)*1e-9;
|
|
double toks=(double)T*iters;
|
|
printf(" %d steps in %.4fs -> %.1f tokens/sec (full fwd+bwd, resident weights)\n",
|
|
iters,sec,toks/sec);
|
|
printf(" (per-step %.2f ms; %d params)\n",sec/iters*1e3,rx_ptmodel_num_params(cm));
|
|
free(tokens);free(targets);rx_ptmodel_free(cm);
|
|
}
|
|
|
|
int main(void){
|
|
int dev=0; cudaSetDevice(dev); cudaDeviceProp pr; cudaGetDeviceProperties(&pr,dev);
|
|
printf("=== FULL RWKV-x training step — GPU parity vs CPU oracle (Phase-2 gate) ===\n");
|
|
printf("device %d: %s cc %d.%d\n",dev,pr.name,pr.major,pr.minor);
|
|
|
|
/* small config: 3 RWKV + 1 attention (interval=3), N=64, small vocab/T. */
|
|
rx_pt_config pc=(rx_pt_config){ .n_embd=128,.n_head=2,.head_size=64,
|
|
.n_rwkv_layer=3,.n_attn_layer=1,.vocab_size=64,
|
|
.Dw=32,.Da=32,.Dg=32,.Dv=32,.ffn_dim=256,.chunk_size=4,.moba_top_k=0 };
|
|
run_parity("dense causal", pc, 16, 12345, -1);
|
|
|
|
/* genuine MoBA-sparse selection active (topk<ncand) */
|
|
pc.moba_top_k=1;
|
|
run_parity("moba sparse", pc, 16, 22222, -1);
|
|
|
|
/* ignore-index in the CE (mask some targets) — dense */
|
|
pc.moba_top_k=0;
|
|
/* larger throughput config (still one die, B=1) */
|
|
rx_pt_config tc=(rx_pt_config){ .n_embd=512,.n_head=8,.head_size=64,
|
|
.n_rwkv_layer=6,.n_attn_layer=2,.vocab_size=8000,
|
|
.Dw=64,.Da=64,.Dg=128,.Dv=64,.ffn_dim=2048,.chunk_size=64,.moba_top_k=0 };
|
|
run_throughput(tc, 512, 30, 777);
|
|
|
|
printf("\n%s\n", g_fail?"RESULT: FAIL":"RESULT: PASS");
|
|
return g_fail;
|
|
}
|