David
/
Foundry
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Foundry/cuda_kernels/train_gpu.cu

642 lines
35 KiB
Plaintext
Raw Permalink Blame History

/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
/*
* train_gpu.cu — On-device GPU training driver for the RWKV-x SR4-vs-AdamW scale
* test (Phase 3, Chunk B). The GPU counterpart of the scale test's CPU oracle
* training driver: same enwik8 byte data pipeline, same batch
* sampling, same next-byte train loop, same eval + logging (identical CSV), but
* the model forward/backward and the optimizer step run ON THE K80, with weights
* + optimizer state RESIDENT on-device (uploaded once).
*
* REUSE (do not rebuild): cuda_kernels/fullmodel.cu gives the device-resident
* full RWKV-x fwd+bwd (fm_build / fm_forward / fm_backward, parity-green vs the
* CPU oracle) and its grad registry (fm_gfind). cuda_kernels/optim.cu gives the
* GPU SR4 step (fc_sr4_scale + fc_sr4_update, bit-exact vs CPU) and the AdamW
* step (fc_adamw_update). This file adds ONLY: the data/loop/eval/logging glue,
* the SR4 <-> continuous split for the ternary projections done on-device, the
* per-column scale AdamW, and B>1 gradient accumulation. Nothing in the merged
* component files is edited.
*
* Optimizer split (identical policy to src/train.c / test_overfit.c):
* SR4 arm : the big linear projections (attn + time-mix receptance/key/value/
* output, channel-mix ffn_key/ffn_value) are ternary {-1,0,+1} from
* an int4 accumulator q with a per-OUTPUT-COLUMN fp32 scale;
* W_eff[r,c] = scale[c]*ternary(q)[r,c]. Everything else (emb, head,
* LN, biases, lerps, LoRA gates, AND the scales) is continuous AdamW.
* AdamW arm : every parameter continuous, AdamW.
*
* Acceptance: run this and src/train.c (CPU) from the IDENTICAL seed/init/data/
* batch order and compare per-step training loss. Because the SR4 step is
* bit-exact and fwd+bwd is parity-green, the SR4 arm tracks the CPU loss curve
* very tightly (small fp drift only).
*
* Build (K80 host, GPU 0, from foundry root):
* gcc-11 -O2 -I. -c rwkvx/rx_pretrain.c -o rx_pretrain_tg.o
* gcc-11 -O2 -I. -c data/fd_byte_dataset.c -o fd_byte_dataset_tg.o
* gcc-11 -O2 -I. -c ternary_optim/ternary_optim.c -o ternary_optim_tg.o
* nvcc -O2 -arch=sm_37 -ccbin gcc-11 -isystem cuda_kernels/nvcc_shim -I. \
* cuda_kernels/train_gpu.cu rx_pretrain_tg.o fd_byte_dataset_tg.o \
* ternary_optim_tg.o -lcublas -lm -o build/train_gpu
*/
#include <cuda_runtime.h>
#include <cublas_v2.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <stdint.h>
#include <time.h>
#include "cuda_kernels/fullmodel.cu" /* fm_* device model + rx_pretrain.h + block kernels */
#include "cuda_kernels/optim.cu" /* fc_sr4_scale / fc_sr4_update / fc_adamw_update */
#include "ternary_optim/ternary_optim.h" /* host q-init (identical to CPU) */
#include "data/fd_byte_dataset.h"
#define LN2 0.69314718055994530942
#define SR4_Q_CLAMP 7
#define SR4_THETA 2
#define TGK(call) do{ cudaError_t e_=(call); if(e_!=cudaSuccess){ \
fprintf(stderr,"CUDA %s:%d %s\n",__FILE__,__LINE__,cudaGetErrorString(e_)); exit(1);} }while(0)
static inline unsigned tg_nb(size_t n){ return (unsigned)((n+255)/256); }
/* ================================ device glue kernels ================================ */
/* ternarize the int4 accumulator: w = 0 if |q|<theta else sign(q). (== ternary_optim_materialize) */
__global__ void tg_tern_from_q(const int8_t*q,float*tern,int theta,int numel){
int i=blockIdx.x*blockDim.x+threadIdx.x; if(i>=numel)return;
int v=q[i]; tern[i]=(v>=theta)?1.0f:(v<=-theta)?-1.0f:0.0f;
}
/* W_eff[r,c] = scale[c]*tern[r,c] (row-major [rows,cols]) */
__global__ void tg_apply_scale(float*Weff,const float*tern,const float*scale,int rows,int cols){
size_t g=(size_t)blockIdx.x*blockDim.x+threadIdx.x; size_t n=(size_t)rows*cols; if(g>=n)return;
int c=(int)(g%(size_t)cols); Weff[g]=scale[c]*tern[g];
}
/* gq[r,c] = grad[r,c]*scale[c] (dL/d ternaryW) */
__global__ void tg_gq(const float*grad,const float*scale,float*gq,int rows,int cols){
size_t g=(size_t)blockIdx.x*blockDim.x+threadIdx.x; size_t n=(size_t)rows*cols; if(g>=n)return;
int c=(int)(g%(size_t)cols); gq[g]=grad[g]*scale[c];
}
/* per-column scale grad: sgrad[c] = sum_r grad[r,c]*tern[r,c], double-accumulated
* in the SAME index order as the CPU (r ascending) to minimize fp drift. */
__global__ void tg_scale_grad(const float*grad,const float*tern,float*sgrad,int rows,int cols){
int c=blockIdx.x*blockDim.x+threadIdx.x; if(c>=cols)return;
double s=0.0;
for(int r=0;r<rows;r++){ size_t k=(size_t)r*cols+c; s+=(double)grad[k]*(double)tern[k]; }
sgrad[c]=(float)s;
}
__global__ void tg_scale_vec(float*v,float k,int n){ int i=blockIdx.x*blockDim.x+threadIdx.x; if(i<n)v[i]*=k; }
/* opt#3 graph variant: the scale factor (scale_lr/lr, varies with the cosine LR)
* is read from the device optimizer-scalar buffer opt[3]. */
__global__ void tg_scale_vec_g(float*v,const float*opt,int n){ int i=blockIdx.x*blockDim.x+threadIdx.x; if(i<n)v[i]*=opt[3]; }
__global__ void tg_acc_add(float*acc,const float*g,size_t n){ size_t i=(size_t)blockIdx.x*blockDim.x+threadIdx.x; if(i<n)acc[i]+=g[i]; }
__global__ void tg_scale_write(float*dst,const float*acc,float inv,size_t n){ size_t i=(size_t)blockIdx.x*blockDim.x+threadIdx.x; if(i<n)dst[i]=acc[i]*inv; }
__global__ void tg_fill(float*v,float val,int n){ int i=blockIdx.x*blockDim.x+threadIdx.x; if(i<n)v[i]=val; }
/* ================================ resident optimizer launchers ================================ */
static void adamw_dev(float*data,const float*grad,float*mm,float*vv,int64_t n,
int t,float lr,float b1,float b2,float eps,float wd,int is_adamw){
if(n<=0) return;
float bc1=1.0f-powf(b1,(float)t), bc2=1.0f-powf(b2,(float)t);
int th=256; int64_t bl=(n+th-1)/th;
fc_adamw_update<<<(unsigned)bl,th>>>(data,grad,mm,vv,n,lr,b1,b2,eps,wd,bc1,bc2,is_adamw);
}
/* SR4 q-step on device-resident pointers (batch_size==1: grad already averaged). */
static void sr4_dev(int8_t*q,const float*gq,int numel,uint64_t counter,float target_step,
uint64_t seed,int64_t idx_base,int q_clamp,int use_dither,float*d_s){
float inv_bs=1.0f;
fc_sr4_scale<<<1,1>>>(gq,numel,inv_bs,target_step,d_s);
int th=256,bl=(numel+th-1)/th;
fc_sr4_update<<<bl,th>>>(q,gq,numel,inv_bs,d_s,seed,idx_base,counter,q_clamp,use_dither);
}
/* opt#3 graph variants: per-step scalars come from device buffers (d_opt / d_ctr)
* so the captured kernel nodes replay with fresh values while pointers stay fixed. */
static void adamw_dev_g(float*data,const float*grad,float*mm,float*vv,int64_t n,
const float*d_opt,float b1,float b2,float eps,float wd,int is_adamw){
if(n<=0) return;
int th=256; int64_t bl=(n+th-1)/th;
fc_adamw_update_g<<<(unsigned)bl,th>>>(data,grad,mm,vv,n,d_opt,b1,b2,eps,wd,is_adamw);
}
static void sr4_dev_g(int8_t*q,const float*gq,int numel,const unsigned long long*d_ctr,
float target_step,uint64_t seed,int64_t idx_base,int q_clamp,int use_dither,float*d_s){
float inv_bs=1.0f;
fc_sr4_scale<<<1,1>>>(gq,numel,inv_bs,target_step,d_s);
int th=256,bl=(numel+th-1)/th;
fc_sr4_update_g<<<bl,th>>>(q,gq,numel,inv_bs,d_s,seed,idx_base,d_ctr,q_clamp,use_dither);
}
/* ================================ param classification (== train.c) ================================ */
static const char *suffix_after_dot(const char *name){ const char*d=strrchr(name,'.'); return d?d+1:name; }
static int is_ternary_name(const char *name){
const char *s=suffix_after_dot(name);
return !strcmp(s,"receptance_w")||!strcmp(s,"key_w")||!strcmp(s,"value_w")||
!strcmp(s,"output_w")||!strcmp(s,"ffn_key_w")||!strcmp(s,"ffn_value_w");
}
static int rows_of(const char *name,int numel,int C,int ffn){
const char*s=suffix_after_dot(name); (void)numel;
if(!strcmp(s,"ffn_value_w")) return ffn; return C;
}
/* device weight pointer for an oracle param name (mirror of fm_build registration). */
static float* fm_wfind(fm_model*m,const char*name){
if(!strcmp(name,"emb")) return m->emb;
if(!strcmp(name,"head")) return m->head;
if(!strcmp(name,"ln_out_w")) return m->lnout_w;
if(!strcmp(name,"ln_out_b")) return m->lnout_b;
int idx=-1; char suf[48]={0};
if(sscanf(name,"rwkv%d.%47s",&idx,suf)==2 && idx>=0 && idx<m->n_rwkv){
fm_rwkv_blk*B=&m->rw[idx]; tm_dw*D=&B->D;
#define M(s,p) if(!strcmp(suf,s)) return p
M("x_r",D->x_r);M("x_k",D->x_k);M("x_v",D->x_v);M("x_w",D->x_w);M("x_a",D->x_a);M("x_g",D->x_g);
M("receptance_w",D->receptance_w);M("key_w",D->key_w);M("value_w",D->value_w);M("output_w",D->output_w);
M("w0",D->w0);M("w1",D->w1);M("w2",D->w2);
M("a0",D->a0);M("a1",D->a1);M("a2",D->a2);
M("g1",D->g1);M("g2",D->g2);
M("v0",D->v0);M("v1",D->v1);M("v2",D->v2);
M("k_k",D->k_k);M("k_a",D->k_a);M("r_k",D->r_k);M("ln_x_w",D->ln_x_w);M("ln_x_b",D->ln_x_b);
M("ffn_x_k",B->cm.x_k);M("ffn_key_w",B->cm.key_w);M("ffn_value_w",B->cm.value_w);
M("ln1_w",B->ln1_w);M("ln1_b",B->ln1_b);M("ln2_w",B->ln2_w);M("ln2_b",B->ln2_b);
#undef M
}
idx=-1; memset(suf,0,sizeof suf);
if(sscanf(name,"attn%d.%47s",&idx,suf)==2 && idx>=0 && idx<m->n_attn){
fm_attn_blk*A=&m->at[idx];
#define M(s,p) if(!strcmp(suf,s)) return p
M("ln1_w",A->ln1_w);M("ln1_b",A->ln1_b);M("ln2_w",A->ln2_w);M("ln2_b",A->ln2_b);
M("receptance_w",A->Wq);M("key_w",A->Wk);M("value_w",A->Wv);M("output_w",A->Wo);
M("ffn_x_k",A->cm.x_k);M("ffn_key_w",A->cm.key_w);M("ffn_value_w",A->cm.value_w);
#undef M
}
return NULL;
}
/* ================================ config (== train.c) ================================ */
typedef struct {
int C,n_head,head_size,n_rwkv,n_attn,ffn,lora;
int seq_len,batch,steps,eval_every,eval_batches;
uint64_t seed;
int use_sr4;
float target_step;
float lr0,lr_min;
float scale_lr;
const char *logpath,*datapath;
} cfg_t;
static cfg_t default_cfg(void){
cfg_t c; memset(&c,0,sizeof c);
c.C=128; c.n_head=4; c.head_size=32; c.n_rwkv=3; c.n_attn=1;
c.ffn=512; c.lora=32;
c.seq_len=512; c.batch=1; c.steps=1000; c.eval_every=5; c.eval_batches=8;
c.seed=1234ULL;
c.use_sr4=0; c.target_step=0.003f;
c.lr0=1e-3f; c.lr_min=1e-4f; c.scale_lr=1e-2f;
c.logpath=NULL; c.datapath="data/enwik8";
return c;
}
static rx_pt_config to_pt_config(const cfg_t *c){
rx_pt_config pc; memset(&pc,0,sizeof pc);
pc.n_embd=c->C; pc.n_head=c->n_head; pc.head_size=c->head_size;
pc.n_rwkv_layer=c->n_rwkv; pc.n_attn_layer=c->n_attn;
pc.vocab_size=FD_BYTE_VOCAB;
pc.Dw=c->lora; pc.Da=c->lora; pc.Dg=c->lora; pc.Dv=c->lora;
pc.ffn_dim=c->ffn;
pc.chunk_size=c->seq_len; pc.moba_top_k=0;
return pc;
}
static float cosine_lr(float lr0,float lr_min,int step,int total){
if(total<=0||step>=total) return lr_min;
float cs=0.5f*(1.0f+cosf((float)M_PI*(float)step/(float)total));
return lr_min+(lr0-lr_min)*cs;
}
static double now_sec(void){ struct timespec ts; clock_gettime(CLOCK_MONOTONIC,&ts);
return (double)ts.tv_sec+1e-9*(double)ts.tv_nsec; }
static void iso_now(char*buf,size_t n){ time_t tt=time(NULL); struct tm tmv; localtime_r(&tt,&tmv);
strftime(buf,n,"%Y-%m-%dT%H:%M:%S",&tmv); }
/* L2 over ALL model param grads (device registry) — matches train.c grad_l2. */
static double grad_l2_gpu(fm_model*m,float*hbuf){
double s=0.0;
for(int i=0;i<m->ng;i++){ int n=m->gn[i];
TGK(cudaMemcpy(hbuf,m->gptr[i],(size_t)n*sizeof(float),cudaMemcpyDeviceToHost));
for(int k=0;k<n;k++) s+=(double)hbuf[k]*(double)hbuf[k]; }
return sqrt(s);
}
/* held-out val: mean bpb + top-1 acc (identical protocol to train.c eval_metrics). */
static double eval_metrics_gpu(fm_model*m,const fd_byte_split*val,uint64_t seed,int seq_len,
int n_batches,int*in,int*tg,float*hlog,double*out_acc){
double sum_ce=0.0; int n=0; long correct=0,total=0; int V=m->V;
for(int b=0;b<n_batches;b++){
fd_byte_get_batch(val,seed,0,(uint64_t)b,0,1,1,in,tg);
fm_forward(m,in,seq_len);
TGK(cudaMemcpy(hlog,m->logits,(size_t)seq_len*V*sizeof(float),cudaMemcpyDeviceToHost));
double ce=rx_cross_entropy_forward(hlog,tg,seq_len,V,-1);
sum_ce+=ce; n++;
for(int t=0;t<seq_len;t++){ int tgt=tg[t]; if(tgt<0)continue;
const float*row=&hlog[(size_t)t*V]; int am=0; float best=row[0];
for(int v=1;v<V;v++) if(row[v]>best){best=row[v];am=v;}
if(am==tgt)correct++; total++; }
}
if(out_acc)*out_acc= total?(double)correct/(double)total:0.0;
return (n?sum_ce/n:0.0)/LN2;
}
/* ================================ ternary-matrix device bookkeeping ================================ */
typedef struct {
char name[48];
float *Weff,*grad; /* device W_eff buffer + dL/dW_eff (fm registry) */
int rows,cols,numel;
int64_t idx_base;
int8_t *d_q; /* int4 accumulator (device) */
float *d_scale,*d_ms,*d_vs; /* per-column scale + AdamW state [cols] */
float *d_tern,*d_gq; /* scratch [numel] */
float *d_sgrad,*d_s; /* scratch: scale grad [cols], step-scale scalar */
float *snap; /* host prev ternary classification (churn) */
int8_t *hq; /* host q scratch (panel) */
} tern_mat;
typedef struct { float*data,*grad,*mm,*vv; int64_t n; } cont_p;
typedef struct { double churn,pct_neg,pct_zero,pct_pos,rail_pct,dead_pct,scale_mean,scale_min,scale_max; } sr4_panel;
/* compute the SR4 health panel from the POST-step device q + scale (host reduction). */
static void sr4_panel_compute(tern_mat*tm,int n_tern,int snap_valid,sr4_panel*p,float*hscale_max_buf){
(void)hscale_max_buf;
long tot=0,cneg=0,czero=0,cpos=0,crail=0,cdead=0,cchurn=0;
double ssum=0.0,smin=1e30,smax=-1e30; long scols=0;
for(int t=0;t<n_tern;t++){
tern_mat*tt=&tm[t]; int n=tt->numel;
TGK(cudaMemcpy(tt->hq,tt->d_q,(size_t)n,cudaMemcpyDeviceToHost));
float*hs=(float*)malloc((size_t)tt->cols*sizeof(float));
TGK(cudaMemcpy(hs,tt->d_scale,(size_t)tt->cols*sizeof(float),cudaMemcpyDeviceToHost));
for(int k=0;k<n;k++){
int q=tt->hq[k];
float v=(q>=SR4_THETA)?1.0f:(q<=-SR4_THETA)?-1.0f:0.0f;
if(v<0.0f)cneg++; else if(v>0.0f)cpos++; else czero++;
int aq=q<0?-q:q;
if(aq==SR4_Q_CLAMP)crail++;
if(aq<SR4_THETA)cdead++;
if(snap_valid && tt->snap[k]!=v)cchurn++;
tt->snap[k]=v; tot++;
}
for(int cc=0;cc<tt->cols;cc++){ float s=hs[cc]; ssum+=s; if(s<smin)smin=s; if(s>smax)smax=s; scols++; }
free(hs);
}
double dt=tot?(double)tot:1.0;
p->churn=snap_valid?(double)cchurn/dt:0.0;
p->pct_neg=(double)cneg/dt; p->pct_zero=(double)czero/dt; p->pct_pos=(double)cpos/dt;
p->rail_pct=(double)crail/dt; p->dead_pct=(double)cdead/dt;
p->scale_mean=scols?ssum/(double)scols:0.0; p->scale_min=scols?smin:0.0; p->scale_max=scols?smax:0.0;
}
/* ================================================================== */
int main(int argc,char**argv){
setvbuf(stdout,NULL,_IOLBF,0); /* line-buffered for live tails */
cfg_t c=default_cfg();
for(int i=1;i<argc;i++){
const char*a=argv[i];
#define NEXT (++i<argc?argv[i]:"")
if (!strcmp(a,"--opt")){ const char*v=NEXT; c.use_sr4=!strcmp(v,"sr4"); }
else if (!strcmp(a,"--target-step")) c.target_step=(float)atof(NEXT);
else if (!strcmp(a,"--lr")) c.lr0=(float)atof(NEXT);
else if (!strcmp(a,"--lr-min")) c.lr_min=(float)atof(NEXT);
else if (!strcmp(a,"--scale-lr")) c.scale_lr=(float)atof(NEXT);
else if (!strcmp(a,"--seq-len")) c.seq_len=atoi(NEXT);
else if (!strcmp(a,"--batch")) c.batch=atoi(NEXT);
else if (!strcmp(a,"--steps")) c.steps=atoi(NEXT);
else if (!strcmp(a,"--eval-every")) c.eval_every=atoi(NEXT);
else if (!strcmp(a,"--eval-batches"))c.eval_batches=atoi(NEXT);
else if (!strcmp(a,"--seed")) c.seed=(uint64_t)strtoull(NEXT,NULL,10);
else if (!strcmp(a,"--C")) c.C=atoi(NEXT);
else if (!strcmp(a,"--n-head")) c.n_head=atoi(NEXT);
else if (!strcmp(a,"--head-size")) c.head_size=atoi(NEXT);
else if (!strcmp(a,"--n-rwkv")) c.n_rwkv=atoi(NEXT);
else if (!strcmp(a,"--n-attn")) c.n_attn=atoi(NEXT);
else if (!strcmp(a,"--ffn")) c.ffn=atoi(NEXT);
else if (!strcmp(a,"--lora")) c.lora=atoi(NEXT);
else if (!strcmp(a,"--data")) c.datapath=NEXT;
else if (!strcmp(a,"--log")) c.logpath=NEXT;
else if (!strcmp(a,"--device")) { (void)NEXT; } /* accepted; this binary is gpu */
else { fprintf(stderr,"unknown arg: %s\n",a); return 2; }
#undef NEXT
}
if(c.seed==FD_BYTE_SEED_UNSET){ fprintf(stderr,"seed must be nonzero\n"); return 2; }
if(c.head_size*c.n_head!=c.C){ fprintf(stderr,"head_size*n_head (%d) must == C (%d)\n",c.head_size*c.n_head,c.C); return 2; }
if(c.eval_every<1) c.eval_every=1;
int dev=0; cudaSetDevice(dev); cudaDeviceProp pr; cudaGetDeviceProperties(&pr,dev);
FILE*logfp=c.logpath?fopen(c.logpath,"w"):NULL;
#define LOGF(...) do{ printf(__VA_ARGS__); if(logfp){fprintf(logfp,__VA_ARGS__); fflush(logfp);} }while(0)
/* ---- data ---- */
fd_byte_corpus corpus;
if(fd_byte_corpus_open(&corpus,c.datapath)!=0){ fprintf(stderr,"failed to open corpus '%s'\n",c.datapath); return 1; }
fd_byte_split train_s,val_s;
if(fd_byte_split_enwik8_train(&train_s,&corpus,(size_t)c.seq_len)!=0 ||
fd_byte_split_enwik8_val (&val_s, &corpus,(size_t)c.seq_len)!=0){
fprintf(stderr,"failed to build enwik8 splits\n"); return 1; }
/* ---- model (CPU oracle struct, uploaded to device once) ---- */
rx_pt_config pc=to_pt_config(&c);
rx_ptmodel*cm=rx_model_new_random(&pc,c.seed);
if(!cm){ fprintf(stderr,"model init failed\n"); return 1; }
fm_model*m=fm_build(cm,c.seq_len); /* upload ALL weights ONCE (resident) */
int nP=rx_ptmodel_num_params(cm);
long total_params=0,ternary_params=0;
for(int i=0;i<nP;i++){ const char*nm; float*d,*g; int n; rx_ptmodel_param(cm,i,&nm,&d,&g,&n);
total_params+=n; if(is_ternary_name(nm)) ternary_params+=n; }
/* ---- optimizer state (device-resident) ---- */
tern_mat*tmat=(tern_mat*)calloc(nP,sizeof(tern_mat)); int n_tern=0;
cont_p *cont=(cont_p*)calloc(nP,sizeof(cont_p)); int n_cont=0;
int64_t idx_base=0;
for(int i=0;i<nP;i++){
const char*nm; float*d,*g; int n; rx_ptmodel_param(cm,i,&nm,&d,&g,&n);
float*wdev=fm_wfind(m,nm);
int gn; float*gdev=fm_gfind(m,nm,&gn);
if(!wdev||!gdev||gn!=n){ fprintf(stderr,"param map fail: %s (w=%p g=%p gn=%d n=%d)\n",nm,(void*)wdev,(void*)gdev,gn,n); return 1; }
if(c.use_sr4 && is_ternary_name(nm)){
tern_mat*tt=&tmat[n_tern++];
snprintf(tt->name,48,"%s",nm);
tt->Weff=wdev; tt->grad=gdev; tt->numel=n;
tt->rows=rows_of(nm,n,c.C,c.ffn); tt->cols=n/tt->rows;
tt->idx_base=idx_base; idx_base+=n;
/* host ternary_optim purely to reproduce the CPU q-init bit-for-bit */
ternary_optim_config oc; ternary_optim_config_default(&oc);
oc.seed=42; oc.idx_base=tt->idx_base;
ternary_optim*ho=ternary_optim_create(tt->rows,tt->cols,&oc);
const int8_t*qh=ternary_optim_q(ho);
TGK(cudaMalloc(&tt->d_q,(size_t)n)); TGK(cudaMemcpy(tt->d_q,qh,(size_t)n,cudaMemcpyHostToDevice));
ternary_optim_free(ho);
TGK(cudaMalloc(&tt->d_scale,(size_t)tt->cols*sizeof(float)));
tg_fill<<<tg_nb(tt->cols),256>>>(tt->d_scale,1.0f,tt->cols);
TGK(cudaMalloc(&tt->d_ms,(size_t)tt->cols*sizeof(float))); TGK(cudaMemset(tt->d_ms,0,(size_t)tt->cols*sizeof(float)));
TGK(cudaMalloc(&tt->d_vs,(size_t)tt->cols*sizeof(float))); TGK(cudaMemset(tt->d_vs,0,(size_t)tt->cols*sizeof(float)));
TGK(cudaMalloc(&tt->d_tern,(size_t)n*sizeof(float)));
TGK(cudaMalloc(&tt->d_gq,(size_t)n*sizeof(float)));
TGK(cudaMalloc(&tt->d_sgrad,(size_t)tt->cols*sizeof(float)));
TGK(cudaMalloc(&tt->d_s,sizeof(float)));
tt->snap=(float*)malloc((size_t)n*sizeof(float));
tt->hq=(int8_t*)malloc((size_t)n);
} else {
cont_p*p=&cont[n_cont++]; p->data=wdev; p->grad=gdev; p->n=n;
TGK(cudaMalloc(&p->mm,(size_t)n*sizeof(float))); TGK(cudaMemset(p->mm,0,(size_t)n*sizeof(float)));
TGK(cudaMalloc(&p->vv,(size_t)n*sizeof(float))); TGK(cudaMemset(p->vv,0,(size_t)n*sizeof(float)));
}
}
/* B>1 gradient-accumulation shadow (one device buffer per registered grad). */
float**accg=NULL;
if(c.batch>1){ accg=(float**)calloc(m->ng,sizeof(float*));
for(int i=0;i<m->ng;i++){ TGK(cudaMalloc(&accg[i],(size_t)m->gn[i]*sizeof(float))); } }
/* host scratch */
int V=m->V;
int *in =(int*)malloc((size_t)c.batch*c.seq_len*sizeof(int));
int *tg =(int*)malloc((size_t)c.batch*c.seq_len*sizeof(int));
int *in1=(int*)malloc((size_t)c.seq_len*sizeof(int));
int *tg1=(int*)malloc((size_t)c.seq_len*sizeof(int));
int maxgn=0; for(int i=0;i<m->ng;i++) if(m->gn[i]>maxgn)maxgn=m->gn[i];
float*hgrad=(float*)malloc((size_t)maxgn*sizeof(float));
float*hlog =(float*)malloc((size_t)c.seq_len*V*sizeof(float));
double*hloss=(double*)malloc((size_t)c.seq_len*sizeof(double)); /* opt#3 loss readback */
/* ---- banner (mirror train.c) ---- */
char isobuf[32]; iso_now(isobuf,sizeof isobuf);
const double rand_bpb=log((double)FD_BYTE_VOCAB)/LN2;
LOGF("# start=%s\n",isobuf);
LOGF("# device=%s cc %d.%d (GPU driver)\n",pr.name,pr.major,pr.minor);
LOGF("=== RWKV-x 1M smoke on enwik8 (%s arm) [GPU] ===\n",c.use_sr4?"SR4":"AdamW");
LOGF("config: C=%d heads=%d hs=%d ffn=%d lora=%d layers: %d RWKV + %d attn (interleave)\n",
c.C,c.n_head,c.head_size,c.ffn,c.lora,c.n_rwkv,c.n_attn);
LOGF("vocab=%d seq_len=%d batch=%d steps=%d eval_every=%d eval_batches=%d seed=%llu\n",
FD_BYTE_VOCAB,c.seq_len,c.batch,c.steps,c.eval_every,c.eval_batches,(unsigned long long)c.seed);
LOGF("PARAM COUNT: total=%ld (ternary=%ld, continuous=%ld)\n",
total_params,ternary_params,total_params-ternary_params);
if(c.use_sr4) LOGF("SR4 target_step=%.4g scale_lr=%.4g cont_lr=%.4g->%.4g (Q_CLAMP=%d THETA=%d)\n",
c.target_step,c.scale_lr,c.lr0,c.lr_min,SR4_Q_CLAMP,SR4_THETA);
else LOGF("AdamW lr=%.4g->%.4g (cosine)\n",c.lr0,c.lr_min);
LOGF("random-init reference bpb = ln(256)/ln2 = %.4f\n",rand_bpb);
if(c.use_sr4) LOGF("# SR4: %d ternary matrices, %d continuous params\n",n_tern,n_cont);
LOGF("step,train_bpb,val_bpb,val_acc,grad_norm,lr,"
"step_sec,tok_per_sec,elapsed_sec,ETA_sec,"
"churn,pct_neg,pct_zero,pct_pos,rail_pct,dead_pct,"
"scale_mean,scale_min,scale_max\n");
double best_vb=1e30; int best_vb_step=-1;
double best_va=-1.0; int best_va_step=-1;
int snap_valid=0;
double t0=now_sec(); double last_log_wall=t0; int last_log_step=0;
uint64_t counter=0;
const float B1=0.9f,B2=0.999f,EPS=1e-8f,WD=0.0f;
/* ================= opt#3: CUDA-graph capture of the B==1 train step =============
* The trainer fires hundreds of tiny kernels per step; at 1M params the CPU
* cannot issue launches fast enough and the K80 starves between them (~1% of
* peak, launch-bound). Capture the fixed per-step kernel sequence ONCE and
* replay it as a single cudaGraphLaunch -> the kernels run back-to-back with
* ~no per-launch CPU overhead. Opt#2's resident arena already gives the stable
* buffer pointers graphs require. Only B==1 is graphed (the tiny-run case);
* B>1 keeps the proven non-graph accumulation path below. */
float *d_opt=NULL; unsigned long long *d_ctr=NULL;
cudaGraph_t g_graph=NULL; cudaGraphExec_t g_exec=NULL; int graph_ready=0;
TGK(cudaMalloc(&d_opt,4*sizeof(float))); /* {lr, bc1, bc2, scale_lr/lr} */
TGK(cudaMalloc(&d_ctr,sizeof(unsigned long long))); /* SR4 batch_counter */
if(c.batch==1){
/* Every kernel launches on stream 0, which -> the per-thread default stream
* (this TU is built with --default-stream per-thread). Route cuBLAS and
* all per-step memcpy/memset onto that SAME per-thread stream so the entire
* step lives on one capturable stream. */
g_fm_stream = cudaStreamPerThread;
FBK(cublasSetStream(m->H, cudaStreamPerThread));
/* stage a real (step-0) batch into the resident input buffers */
fd_byte_get_batch(&train_s,c.seed,0,0ULL,0,1,1,in,tg);
TGK(cudaMemcpy(m->tok,in,(size_t)c.seq_len*sizeof(int),cudaMemcpyHostToDevice));
TGK(cudaMemcpy(m->dtg,tg,(size_t)c.seq_len*sizeof(int),cudaMemcpyHostToDevice));
/* WARMUP: one fwd+bwd actually EXECUTED (g_fm_capturing==0) to force cuBLAS
* lazy init/workspace + arena hi-water BEFORE capture (a cudaMalloc during
* capture would abort it). Touches only activations/grads (recomputed every
* step), never optimizer state, so step 0 below is still pristine. */
fm_forward(m,in,c.seq_len);
fm_backward(m,tg,c.seq_len,-1);
TGK(cudaDeviceSynchronize());
fmA_reset(); /* drop any overflow allocs + rewind the arena before capture */
/* CAPTURE the full step: materialize + fwd + loss + bwd + optimizer. The
* per-step-varying scalars (cosine lr, bias-correction bc1/bc2, scale_lr/lr,
* and the SR4 batch_counter) are read by the _g kernel variants from the
* d_opt/d_ctr device buffers, which the driver refreshes each step before
* replay -> fresh values every replay while pointers stay fixed. */
g_fm_capturing = 1;
TGK(cudaStreamBeginCapture(g_fm_stream, cudaStreamCaptureModeThreadLocal));
if(c.use_sr4){
for(int t=0;t<n_tern;t++){ tern_mat*tt=&tmat[t];
tg_tern_from_q<<<tg_nb(tt->numel),256>>>(tt->d_q,tt->d_tern,SR4_THETA,tt->numel);
tg_apply_scale<<<tg_nb((size_t)tt->numel),256>>>(tt->Weff,tt->d_tern,tt->d_scale,tt->rows,tt->cols); }
}
fm_forward(m,in,c.seq_len);
fm_backward(m,tg,c.seq_len,-1);
if(c.use_sr4){
for(int t=0;t<n_tern;t++){ tern_mat*tt=&tmat[t];
tg_scale_grad<<<tg_nb(tt->cols),256>>>(tt->grad,tt->d_tern,tt->d_sgrad,tt->rows,tt->cols);
tg_gq<<<tg_nb((size_t)tt->numel),256>>>(tt->grad,tt->d_scale,tt->d_gq,tt->rows,tt->cols); }
for(int t=0;t<n_tern;t++){ tern_mat*tt=&tmat[t];
sr4_dev_g(tt->d_q,tt->d_gq,tt->numel,d_ctr,c.target_step,42ULL,tt->idx_base,SR4_Q_CLAMP,1,tt->d_s); }
for(int t=0;t<n_tern;t++){ tern_mat*tt=&tmat[t];
tg_scale_vec_g<<<tg_nb(tt->cols),256>>>(tt->d_sgrad,d_opt,tt->cols);
adamw_dev_g(tt->d_scale,tt->d_sgrad,tt->d_ms,tt->d_vs,tt->cols,d_opt,B1,B2,EPS,WD,1); }
for(int i=0;i<n_cont;i++){ cont_p*p=&cont[i];
adamw_dev_g(p->data,p->grad,p->mm,p->vv,p->n,d_opt,B1,B2,EPS,WD,1); }
} else {
for(int i=0;i<n_cont;i++){ cont_p*p=&cont[i];
adamw_dev_g(p->data,p->grad,p->mm,p->vv,p->n,d_opt,B1,B2,EPS,WD,1); }
}
cudaError_t cap_err=cudaStreamEndCapture(g_fm_stream,&g_graph);
g_fm_capturing=0;
if(cap_err!=cudaSuccess){ fprintf(stderr,"opt#3 graph capture failed: %s\n",cudaGetErrorString(cap_err)); return 1; }
TGK(cudaGraphInstantiate(&g_exec,g_graph,NULL,NULL,0));
graph_ready=1;
LOGF("# opt#3: CUDA graph captured (full B=1 step on the per-thread stream); replaying per step.\n");
}
for(int step=0; step<c.steps; step++){
float lr=cosine_lr(c.lr0,c.lr_min,step,c.steps);
int last=(step==c.steps-1);
int do_log=(step%c.eval_every==0||last);
double train_ce=0.0, grad_norm=0.0;
if(graph_ready){
/* ---- opt#3 CUDA-graph path (B==1): one launch replays the whole step ---- */
fd_byte_get_batch(&train_s,c.seed,0,(uint64_t)step,0,1,1,in,tg);
TGK(cudaMemcpy(m->tok,in,(size_t)c.seq_len*sizeof(int),cudaMemcpyHostToDevice));
TGK(cudaMemcpy(m->dtg,tg,(size_t)c.seq_len*sizeof(int),cudaMemcpyHostToDevice));
/* refresh per-step scalars (bit-identical to the by-value non-graph path) */
float bc1=1.0f-powf(B1,(float)(step+1)), bc2=1.0f-powf(B2,(float)(step+1));
float h_opt[4]={ lr, bc1, bc2, (lr!=0.0f)?(c.scale_lr/lr):0.0f };
TGK(cudaMemcpy(d_opt,h_opt,4*sizeof(float),cudaMemcpyHostToDevice));
unsigned long long hc=(unsigned long long)step; /* SR4 counter == step */
TGK(cudaMemcpy(d_ctr,&hc,sizeof(hc),cudaMemcpyHostToDevice));
TGK(cudaGraphLaunch(g_exec,g_fm_stream));
TGK(cudaStreamSynchronize(g_fm_stream));
/* reduce resident per-token loss -> train_ce (nv==T since ignore_index=-1) */
TGK(cudaMemcpy(hloss,m->loss_t,(size_t)c.seq_len*sizeof(double),cudaMemcpyDeviceToHost));
double L=0.0; for(int t=0;t<c.seq_len;t++) L+=hloss[t]; train_ce=L/(double)c.seq_len;
/* grads survive the optimizer step unmodified, so grad_l2 post-replay matches */
if(do_log) grad_norm=grad_l2_gpu(m,hgrad);
} else {
/* ---- non-graph path (B>1 accumulation, or fallback) ---- */
if(c.use_sr4){
for(int t=0;t<n_tern;t++){ tern_mat*tt=&tmat[t];
tg_tern_from_q<<<tg_nb(tt->numel),256>>>(tt->d_q,tt->d_tern,SR4_THETA,tt->numel);
tg_apply_scale<<<tg_nb((size_t)tt->numel),256>>>(tt->Weff,tt->d_tern,tt->d_scale,tt->rows,tt->cols);
}
}
fd_byte_get_batch(&train_s,c.seed,0,(uint64_t)step,0,1,(uint32_t)c.batch,in,tg);
if(c.batch==1){
fm_forward(m,in,c.seq_len);
train_ce=fm_backward(m,tg,c.seq_len,-1);
} else {
for(int i=0;i<m->ng;i++) TGK(cudaMemset(accg[i],0,(size_t)m->gn[i]*sizeof(float)));
for(int b=0;b<c.batch;b++){
memcpy(in1,&in[(size_t)b*c.seq_len],(size_t)c.seq_len*sizeof(int));
memcpy(tg1,&tg[(size_t)b*c.seq_len],(size_t)c.seq_len*sizeof(int));
fm_forward(m,in1,c.seq_len);
train_ce+=fm_backward(m,tg1,c.seq_len,-1);
for(int i=0;i<m->ng;i++) tg_acc_add<<<tg_nb((size_t)m->gn[i]),256>>>(accg[i],m->gptr[i],(size_t)m->gn[i]);
}
train_ce/=c.batch;
const float inv=1.0f/(float)c.batch;
for(int i=0;i<m->ng;i++) tg_scale_write<<<tg_nb((size_t)m->gn[i]),256>>>(m->gptr[i],accg[i],inv,(size_t)m->gn[i]);
}
grad_norm=do_log?grad_l2_gpu(m,hgrad):0.0;
if(c.use_sr4){
for(int t=0;t<n_tern;t++){ tern_mat*tt=&tmat[t];
tg_scale_grad<<<tg_nb(tt->cols),256>>>(tt->grad,tt->d_tern,tt->d_sgrad,tt->rows,tt->cols);
tg_gq<<<tg_nb((size_t)tt->numel),256>>>(tt->grad,tt->d_scale,tt->d_gq,tt->rows,tt->cols);
}
for(int t=0;t<n_tern;t++){ tern_mat*tt=&tmat[t];
sr4_dev(tt->d_q,tt->d_gq,tt->numel,counter,c.target_step,42ULL,tt->idx_base,SR4_Q_CLAMP,1,tt->d_s); }
counter++;
for(int t=0;t<n_tern;t++){ tern_mat*tt=&tmat[t];
if(c.scale_lr!=lr) tg_scale_vec<<<tg_nb(tt->cols),256>>>(tt->d_sgrad,c.scale_lr/lr,tt->cols);
adamw_dev(tt->d_scale,tt->d_sgrad,tt->d_ms,tt->d_vs,tt->cols,step+1,lr,B1,B2,EPS,WD,1); }
for(int i=0;i<n_cont;i++){ cont_p*p=&cont[i];
adamw_dev(p->data,p->grad,p->mm,p->vv,p->n,step+1,lr,B1,B2,EPS,WD,1); }
} else {
for(int i=0;i<n_cont;i++){ cont_p*p=&cont[i];
adamw_dev(p->data,p->grad,p->mm,p->vv,p->n,step+1,lr,B1,B2,EPS,WD,1); }
}
TGK(cudaDeviceSynchronize());
}
if(do_log){
double vacc=0.0;
double vb=eval_metrics_gpu(m,&val_s,c.seed,c.seq_len,c.eval_batches,in1,tg1,hlog,&vacc);
double tb=train_ce/LN2;
double now=now_sec(); double elapsed=now-t0;
int dstep=step-last_log_step;
double step_sec=(dstep>0)?(now-last_log_wall)/(double)dstep:elapsed/(double)(step+1);
double avg_step=elapsed/(double)(step+1);
double tok_per_sec=step_sec>0.0?(double)c.batch*(double)c.seq_len/step_sec:0.0;
double eta=(double)(c.steps-1-step)*avg_step;
last_log_wall=now; last_log_step=step;
if(vb<best_vb){best_vb=vb;best_vb_step=step;}
if(vacc>best_va){best_va=vacc;best_va_step=step;}
if(c.use_sr4){
sr4_panel p; sr4_panel_compute(tmat,n_tern,snap_valid,&p,NULL); snap_valid=1;
LOGF("%d,%.5f,%.5f,%.5f,%.4g,%.4g,"
"%.4f,%.1f,%.2f,%.2f,"
"%.5f,%.5f,%.5f,%.5f,%.5f,%.5f,"
"%.4f,%.4f,%.4f\n",
step,tb,vb,vacc,grad_norm,lr,step_sec,tok_per_sec,elapsed,eta,
p.churn,p.pct_neg,p.pct_zero,p.pct_pos,p.rail_pct,p.dead_pct,
p.scale_mean,p.scale_min,p.scale_max);
} else {
LOGF("%d,%.5f,%.5f,%.5f,%.4g,%.4g,"
"%.4f,%.1f,%.2f,%.2f,"
"-,-,-,-,-,-,-,-,-\n",
step,tb,vb,vacc,grad_norm,lr,step_sec,tok_per_sec,elapsed,eta);
}
LOGF("# BEST: val_bpb=%.5f (step %d) val_acc=%.5f (step %d)\n",best_vb,best_vb_step,best_va,best_va_step);
}
}
double wall=now_sec()-t0;
/* final metrics (mirror train.c) */
double final_train_bpb,final_val_bpb,final_val_acc=0.0;
if(c.use_sr4){
for(int t=0;t<n_tern;t++){ tern_mat*tt=&tmat[t];
tg_tern_from_q<<<tg_nb(tt->numel),256>>>(tt->d_q,tt->d_tern,SR4_THETA,tt->numel);
tg_apply_scale<<<tg_nb((size_t)tt->numel),256>>>(tt->Weff,tt->d_tern,tt->d_scale,tt->rows,tt->cols); }
}
fd_byte_get_batch(&train_s,c.seed,0,(uint64_t)c.steps,0,1,1,in1,tg1);
fm_forward(m,in1,c.seq_len);
TGK(cudaMemcpy(hlog,m->logits,(size_t)c.seq_len*V*sizeof(float),cudaMemcpyDeviceToHost));
final_train_bpb=rx_cross_entropy_forward(hlog,tg1,c.seq_len,V,-1)/LN2;
final_val_bpb=eval_metrics_gpu(m,&val_s,c.seed,c.seq_len,c.eval_batches*2,in1,tg1,hlog,&final_val_acc);
LOGF("# FINAL: train_bpb=%.5f val_bpb=%.5f val_acc=%.5f wall=%.2fs (%.3fs/step over %d steps)\n",
final_train_bpb,final_val_bpb,final_val_acc,wall,wall/c.steps,c.steps);
LOGF("# BEST: val_bpb=%.5f (step %d) val_acc=%.5f (step %d)\n",best_vb,best_vb_step,best_va,best_va_step);
{
double avg_step_sec=c.steps>0?wall/(double)c.steps:0.0;
long total_tokens=(long)c.steps*(long)c.batch*(long)c.seq_len;
double tps=wall>0.0?(double)total_tokens/wall:0.0;
char endiso[32]; iso_now(endiso,sizeof endiso);
LOGF("# ===== END-OF-RUN SUMMARY (extrapolate future-run wall-time from these) =====\n");
LOGF("# end=%s\n",endiso);
LOGF("# total_wall_sec=%.2f total_steps=%d avg_step_sec=%.4f tok_per_sec=%.1f total_tokens=%ld (batch=%d seq_len=%d)\n",
wall,c.steps,avg_step_sec,tps,total_tokens,c.batch,c.seq_len);
}
if(logfp) fclose(logfp);
return 0;
}
Reference in New Issue View Git Blame Copy Permalink