156 lines
7.3 KiB
Plaintext
156 lines
7.3 KiB
Plaintext
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/*
|
|
* test_optim_cuda.cu — parity harness for the GPU optimizer-step kernels
|
|
* (cuda_kernels/optim.cu) vs the CPU optimizer, on GPU 0. Phase 3, Chunk A.
|
|
*
|
|
* SR4 : ternary_optim_step (ternary_optim/ternary_optim.c) — BIT-EXACT.
|
|
* Require identical int8 q integers (rel-err 0 / exact equality).
|
|
* AdamW : fo_adamw / adam_step (optim/optim.c, linked via a minimal tensor
|
|
* shim) — inf-norm max|gpu-cpu|/max|cpu| < 1e-3 (expect ~1e-6).
|
|
*
|
|
* Covers multiple shapes and several (seed, counter) / step values. The CPU
|
|
* oracle code is LINKED (not reimplemented): ternary_optim.o and optim.o +
|
|
* tensor_shim.o.
|
|
*/
|
|
#include <cuda_runtime.h>
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <stdint.h>
|
|
#include <string.h>
|
|
#include <math.h>
|
|
|
|
#include "cuda_kernels/optim.cu"
|
|
|
|
/* ---- CPU oracle: SR4 (self-contained C lib) ---- */
|
|
#include "ternary_optim/ternary_optim.h"
|
|
|
|
/* ---- CPU oracle: AdamW (real optim.c + tensor shim) ---- */
|
|
#include "tensor/tensor.h"
|
|
#include "optim/optim.h"
|
|
extern "C" tm_tensor *shim_tensor_f32(const float *src, int64_t n);
|
|
|
|
/* ---- deterministic fixture RNG (SplitMix64) ---- */
|
|
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 float sm_uf(void){ return (float)((sm_next()>>11)*(1.0/9007199254740992.0)); }
|
|
static float sm_sym(float s){ return (2.0f*sm_uf()-1.0f)*s; }
|
|
static float* rnd(size_t n,float s){ float*p=(float*)malloc(n*sizeof(float));
|
|
for(size_t i=0;i<n;i++)p[i]=sm_sym(s); return p; }
|
|
|
|
static int g_fail=0;
|
|
|
|
/* =====================================================================
|
|
* SR4 — bit-exact q comparison
|
|
* ===================================================================== */
|
|
static void sr4_case(const char*name,int rows,int cols,uint64_t seed,
|
|
int64_t idx_base,int batch_size,uint64_t counter,
|
|
float target_step,int use_dither,float grad_scale){
|
|
int numel=rows*cols;
|
|
ternary_optim_config cfg; ternary_optim_config_default(&cfg);
|
|
cfg.seed=seed; cfg.idx_base=idx_base; cfg.use_dither=use_dither;
|
|
ternary_optim *opt=ternary_optim_create(rows,cols,&cfg);
|
|
if(!opt){printf(" %s CREATE ERR\n",name);g_fail=1;return;}
|
|
|
|
/* snapshot the identical initial q for the GPU side */
|
|
const int8_t *q0=ternary_optim_q(opt);
|
|
int8_t *q_gpu=(int8_t*)malloc((size_t)numel);
|
|
memcpy(q_gpu,q0,(size_t)numel);
|
|
|
|
sm_state=seed*0x1234567ULL+counter+7;
|
|
float *grad=rnd((size_t)numel,grad_scale);
|
|
|
|
/* CPU step (the oracle) */
|
|
ternary_optim_step(opt,grad,batch_size,counter,target_step);
|
|
const int8_t *q_cpu=ternary_optim_q(opt);
|
|
|
|
/* GPU step from the same initial q + same inputs */
|
|
int rc=sr4_step_cuda(q_gpu,grad,numel,batch_size,counter,target_step,
|
|
cfg.seed,cfg.idx_base,cfg.q_clamp,cfg.use_dither);
|
|
if(rc){printf(" %s GPU ERR rc=%d\n",name,rc);g_fail=1;free(grad);free(q_gpu);ternary_optim_free(opt);return;}
|
|
|
|
int mism=0,first=-1;
|
|
for(int i=0;i<numel;i++) if(q_cpu[i]!=q_gpu[i]){ if(first<0)first=i; mism++; }
|
|
const char*st=mism?"FAIL":"PASS"; if(mism)g_fail=1;
|
|
printf(" %-22s rows=%d cols=%d bs=%d ctr=%llu dith=%d ts=%.4g : mism=%d/%d %s",
|
|
name,rows,cols,batch_size,(unsigned long long)counter,use_dither,target_step,mism,numel,st);
|
|
if(mism)printf(" (first@%d cpu=%d gpu=%d)",first,q_cpu[first],q_gpu[first]);
|
|
printf("\n");
|
|
|
|
free(grad);free(q_gpu);ternary_optim_free(opt);
|
|
}
|
|
|
|
/* =====================================================================
|
|
* AdamW — inf-norm parity over several steps (exercises bias correction)
|
|
* ===================================================================== */
|
|
static void adamw_case(const char*name,int64_t n,float lr,float beta1,float beta2,
|
|
float eps,float wd,int nsteps,uint64_t seed){
|
|
sm_state=seed;
|
|
float *init=rnd((size_t)n,0.5f);
|
|
|
|
/* CPU oracle: real fo_adamw over nsteps */
|
|
tm_tensor *p=shim_tensor_f32(init,n);
|
|
p->grad=shim_tensor_f32(NULL,n);
|
|
p->requires_grad=true;
|
|
tm_tensor *params[1]={p};
|
|
fo_adam_config ac=fo_adam_default(lr);
|
|
ac.beta1=beta1; ac.beta2=beta2; ac.eps=eps; ac.weight_decay=wd;
|
|
fo_optimizer *opt=fo_adamw_new(params,1,&ac);
|
|
|
|
/* GPU: mirror with host-resident m/v, one launch per step */
|
|
float *data_g=(float*)malloc((size_t)n*sizeof(float));
|
|
memcpy(data_g,init,(size_t)n*sizeof(float));
|
|
float *m_g=(float*)calloc((size_t)n,sizeof(float));
|
|
float *v_g=(float*)calloc((size_t)n,sizeof(float));
|
|
|
|
for(int s=0;s<nsteps;s++){
|
|
float *g=rnd((size_t)n,0.3f);
|
|
memcpy(tm_data_f32(p->grad),g,(size_t)n*sizeof(float)); /* same grad both sides */
|
|
fo_step(opt);
|
|
int rc=adamw_step_cuda(data_g,g,m_g,v_g,n,s+1,lr,beta1,beta2,eps,wd,1);
|
|
if(rc){printf(" %s GPU ERR rc=%d\n",name,rc);g_fail=1;free(g);break;}
|
|
free(g);
|
|
}
|
|
|
|
const float *data_c=tm_data_f32(p);
|
|
double maxabs=0.0,scale=0.0;
|
|
for(int64_t i=0;i<n;i++){ double d=fabs((double)data_g[i]-(double)data_c[i]);
|
|
if(d>maxabs)maxabs=d; double c=fabs((double)data_c[i]); if(c>scale)scale=c; }
|
|
double rel=maxabs/(scale+1e-30); const char*st=(rel<1e-3)?"PASS":"FAIL"; if(rel>=1e-3)g_fail=1;
|
|
printf(" %-22s n=%lld steps=%d wd=%.3g : rel=%.3e maxabs=%.3e [%s]\n",
|
|
name,(long long)n,nsteps,wd,rel,maxabs,st);
|
|
|
|
free(init);free(data_g);free(m_g);free(v_g);
|
|
fo_free(opt); tm_release(p->grad); tm_release(p);
|
|
}
|
|
|
|
int main(void){
|
|
int dev=0; cudaSetDevice(dev); cudaDeviceProp pr; cudaGetDeviceProperties(&pr,dev);
|
|
printf("=== GPU optimizer-step parity vs CPU optimizer (GPU %d: %s cc %d.%d) ===\n\n",
|
|
dev,pr.name,pr.major,pr.minor);
|
|
|
|
printf("--- SR4 (require BIT-EXACT q) ---\n");
|
|
sr4_case("sr4_small", 4, 8, 42, 0, 1, 0, 0.003f, 1, 0.01f);
|
|
sr4_case("sr4_small_ctr", 4, 8, 42, 0, 1, 5, 0.003f, 1, 0.01f);
|
|
sr4_case("sr4_stratum2", 4, 8, 42, 0, 1, 513, 0.003f, 1, 0.01f); /* counter>>8 != 0 */
|
|
sr4_case("sr4_nodither", 4, 8, 42, 0, 1, 5, 0.003f, 0, 0.01f);
|
|
sr4_case("sr4_idxbase", 16, 32, 123, 4096, 1, 17, 0.003f, 1, 0.02f);
|
|
sr4_case("sr4_bs4", 16, 32, 42, 0, 4, 3, 0.003f, 1, 0.05f);
|
|
sr4_case("sr4_bigstep", 32, 64, 7, 0, 1, 9, 0.05f, 1, 0.001f); /* large prop -> saturate */
|
|
sr4_case("sr4_mid", 128, 256, 99, 8192, 1, 40, 0.003f, 1, 0.02f);
|
|
sr4_case("sr4_real768", 768, 768, 55, 0, 1, 128, 0.003f, 1, 0.015f);
|
|
sr4_case("sr4_real_ctr300",768, 768, 55, 0, 1, 300, 0.003f, 1, 0.015f);
|
|
|
|
printf("\n--- AdamW (inf-norm < 1e-3) ---\n");
|
|
adamw_case("adamw_small", 64, 1e-3f, 0.9f, 0.999f, 1e-8f, 0.0f, 1, 11);
|
|
adamw_case("adamw_t5", 64, 1e-3f, 0.9f, 0.999f, 1e-8f, 0.0f, 5, 12);
|
|
adamw_case("adamw_t20", 4096, 1e-3f, 0.9f, 0.999f, 1e-8f, 0.0f, 20, 13);
|
|
adamw_case("adamw_wd", 4096, 1e-3f, 0.9f, 0.999f, 1e-8f, 0.01f,10, 14);
|
|
adamw_case("adamw_beta", 4096, 3e-4f, 0.95f,0.98f, 1e-6f, 0.0f, 10, 15);
|
|
adamw_case("adamw_big", 131072, 1e-3f, 0.9f, 0.999f, 1e-8f, 0.0f, 8, 16);
|
|
|
|
printf("\n%s\n", g_fail?"RESULT: FAIL":"RESULT: PASS");
|
|
return g_fail;
|
|
}
|