318 lines
14 KiB
Plaintext
318 lines
14 KiB
Plaintext
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/*
|
|
* optim.cu — CUDA GPU optimizer-step kernels for Foundry (Phase 3, Chunk A).
|
|
*
|
|
* Target: sm_37 (K80, Kepler). Compile: nvcc -ptx -arch=sm_37 ... (or included
|
|
* directly by the parity harness test_optim_cuda.cu and linked with the CPU
|
|
* oracle objects).
|
|
*
|
|
* Two GPU optimizer steps, each with an extern "C" host launcher, parity-checked
|
|
* against the CPU optimizer:
|
|
*
|
|
* 1. sr4_step_cuda — the SR4 4-bit stochastic-rounded latent-accumulator
|
|
* update (ternary_optim_step). BIT-IDENTICAL to the
|
|
* CPU: same SplitMix64 counter-hash keyed on
|
|
* (seed, stratum, idx_base+i), same stratified Weyl
|
|
* low-discrepancy dither, same 8.8 fixed-point grid
|
|
* snap, same saturating add / clamp order.
|
|
* 2. adamw_step_cuda — the continuous-param AdamW update (foundry optim.c
|
|
* adam_step, is_adamw path). Matches beta1/beta2/eps/
|
|
* weight-decay/bias-correction; elementwise so it is
|
|
* order-independent -> parity to fp rounding.
|
|
*
|
|
* The updates are O(params) elementwise (no matmul, no cuBLAS).
|
|
*
|
|
* Bit-exactness note for SR4: the ONLY float->discrete boundary is the per-matrix
|
|
* step scale s = target_step/(mean|g|+1e-12) and the 8.8 grid snap lrintf(prop).
|
|
* mean|g| is a sequential DOUBLE reduction in the CPU code. To be bit-identical
|
|
* we (a) run the reduction on ONE GPU thread in the SAME index order in double,
|
|
* and (b) build the CPU oracle without -ffast-math so gcc cannot reassociate /
|
|
* vectorize that double reduction (-fassociative-math off). Everything after the
|
|
* grid snap is pure integer (the counter-hash, dither, saturating add) and is
|
|
* flag-immune, so it is bit-identical by construction.
|
|
*/
|
|
|
|
#include <cuda_runtime.h>
|
|
#include <stdint.h>
|
|
#include <math.h>
|
|
|
|
/* Stratified low-discrepancy dither constants (ternary_optim.c verbatim). */
|
|
#define TO_LD_STRATUM_BITS 8
|
|
#define TO_LD_STRIDE 159u
|
|
|
|
/* ---------------------------------------------------------------------------
|
|
* SplitMix64-style counter-based hash — byte-identical to ternary_optim.c.
|
|
* ------------------------------------------------------------------------- */
|
|
__device__ static inline uint64_t d_sm64_mix(uint64_t x) {
|
|
x += 0x9E3779B97F4A7C15ULL;
|
|
x = (x ^ (x >> 30)) * 0xBF58476D1CE4E5B9ULL;
|
|
x = (x ^ (x >> 27)) * 0x94D049BB133111EBULL;
|
|
return x ^ (x >> 31);
|
|
}
|
|
__device__ static inline uint64_t d_sr_hash(uint64_t seed, uint64_t batch_counter,
|
|
uint64_t global_idx) {
|
|
uint64_t x = d_sm64_mix(seed ^ 0xD1B54A32D192ED03ULL);
|
|
x = d_sm64_mix(x ^ batch_counter);
|
|
x = d_sm64_mix(x ^ global_idx);
|
|
return x;
|
|
}
|
|
|
|
/* ---------------------------------------------------------------------------
|
|
* SR4 kernel 1: per-matrix step scale s, computed by ONE thread in double in the
|
|
* SAME sequential index order as the CPU (ternary_optim.c lines 163-167). This
|
|
* is the only reduction; the O(params) work is the parallel update below.
|
|
* ------------------------------------------------------------------------- */
|
|
extern "C" __global__ void fc_sr4_scale(const float *grad_sum, int numel,
|
|
float inv_bs, float target_step,
|
|
float *s_out) {
|
|
if (blockIdx.x != 0 || threadIdx.x != 0) return;
|
|
double abs_sum = 0.0;
|
|
for (int i = 0; i < numel; i++)
|
|
abs_sum += fabsf(grad_sum[i] * inv_bs);
|
|
float mean_abs = (float)(abs_sum / numel);
|
|
*s_out = target_step / (mean_abs + 1e-12f);
|
|
}
|
|
|
|
/* ---------------------------------------------------------------------------
|
|
* SR4 kernel 2: the parallel elementwise stochastic-rounded integer update into
|
|
* the int4 accumulator q. Byte-identical to ternary_optim.c lines 169-203.
|
|
* stratum / j / weyl are pure functions of batch_counter (computed here exactly
|
|
* as the CPU does).
|
|
* ------------------------------------------------------------------------- */
|
|
extern "C" __global__ void fc_sr4_update(int8_t *q, const float *grad_sum,
|
|
int numel, float inv_bs,
|
|
const float *s_in, uint64_t seed,
|
|
int64_t idx_base, uint64_t batch_counter,
|
|
int q_clamp, int use_dither) {
|
|
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
|
if (i >= numel) return;
|
|
|
|
float s = *s_in;
|
|
uint64_t stratum = batch_counter >> TO_LD_STRATUM_BITS;
|
|
uint32_t j = (uint32_t)(batch_counter & ((1u << TO_LD_STRATUM_BITS) - 1u));
|
|
uint32_t weyl = (TO_LD_STRIDE * j) & 255u;
|
|
|
|
float g = grad_sum[i] * inv_bs;
|
|
float prop = -g * s * 256.0f;
|
|
if (prop > 4096.0f) prop = 4096.0f;
|
|
if (prop < -4096.0f) prop = -4096.0f;
|
|
int32_t step_fx = (int32_t)lrintf(prop);
|
|
|
|
int32_t k = step_fx >> 8;
|
|
int32_t f = step_fx & 255;
|
|
uint64_t h = d_sr_hash(seed, stratum, (uint64_t)(idx_base + i));
|
|
uint32_t u;
|
|
if (use_dither)
|
|
u = ((uint32_t)(h & 255) + weyl) & 255u; /* stratified Weyl */
|
|
else
|
|
u = (uint32_t)(h & 255); /* plain i.i.d. */
|
|
int32_t sr_bit = ((int32_t)u < f) ? 1 : 0;
|
|
|
|
int32_t nq = (int32_t)q[i] + k + sr_bit;
|
|
if (nq > q_clamp) nq = q_clamp;
|
|
if (nq < -q_clamp) nq = -q_clamp;
|
|
q[i] = (int8_t)nq;
|
|
}
|
|
|
|
/* ---------------------------------------------------------------------------
|
|
* opt#3 graph variant of fc_sr4_update: byte-identical math, but the per-step
|
|
* batch_counter (which MUST advance every replay or the stochastic-rounding
|
|
* dither repeats and every step reuses the same round) is read from a device
|
|
* buffer instead of a by-value arg. The driver refreshes *batch_counter_dev
|
|
* each step before cudaGraphLaunch, so the captured node sees a fresh counter on
|
|
* every replay while its pointers stay fixed. Reading the counter from memory
|
|
* rather than a register changes no arithmetic -> parity is bit-for-bit.
|
|
* ------------------------------------------------------------------------- */
|
|
extern "C" __global__ void fc_sr4_update_g(int8_t *q, const float *grad_sum,
|
|
int numel, float inv_bs,
|
|
const float *s_in, uint64_t seed,
|
|
int64_t idx_base,
|
|
const unsigned long long *batch_counter_dev,
|
|
int q_clamp, int use_dither) {
|
|
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
|
if (i >= numel) return;
|
|
uint64_t batch_counter = (uint64_t)(*batch_counter_dev);
|
|
|
|
float s = *s_in;
|
|
uint64_t stratum = batch_counter >> TO_LD_STRATUM_BITS;
|
|
uint32_t j = (uint32_t)(batch_counter & ((1u << TO_LD_STRATUM_BITS) - 1u));
|
|
uint32_t weyl = (TO_LD_STRIDE * j) & 255u;
|
|
|
|
float g = grad_sum[i] * inv_bs;
|
|
float prop = -g * s * 256.0f;
|
|
if (prop > 4096.0f) prop = 4096.0f;
|
|
if (prop < -4096.0f) prop = -4096.0f;
|
|
int32_t step_fx = (int32_t)lrintf(prop);
|
|
|
|
int32_t k = step_fx >> 8;
|
|
int32_t f = step_fx & 255;
|
|
uint64_t h = d_sr_hash(seed, stratum, (uint64_t)(idx_base + i));
|
|
uint32_t u;
|
|
if (use_dither)
|
|
u = ((uint32_t)(h & 255) + weyl) & 255u;
|
|
else
|
|
u = (uint32_t)(h & 255);
|
|
int32_t sr_bit = ((int32_t)u < f) ? 1 : 0;
|
|
|
|
int32_t nq = (int32_t)q[i] + k + sr_bit;
|
|
if (nq > q_clamp) nq = q_clamp;
|
|
if (nq < -q_clamp) nq = -q_clamp;
|
|
q[i] = (int8_t)nq;
|
|
}
|
|
|
|
/* ---------------------------------------------------------------------------
|
|
* AdamW kernel — elementwise, byte-for-byte the per-element math of optim.c
|
|
* adam_step() on the FO_ADAMW path (decoupled weight decay). bc1/bc2 are passed
|
|
* in (host computes powf, identical to the CPU) so t-dependent bias correction
|
|
* matches exactly.
|
|
* ------------------------------------------------------------------------- */
|
|
extern "C" __global__ void fc_adamw_update(float *data, const float *grad,
|
|
float *m, float *v, int64_t n,
|
|
float lr, float beta1, float beta2,
|
|
float eps, float weight_decay,
|
|
float bc1, float bc2, int is_adamw) {
|
|
int64_t j = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
|
|
if (j >= n) return;
|
|
|
|
float g = grad[j];
|
|
/* For classic Adam (not AdamW), fold weight decay into the gradient. */
|
|
if (!is_adamw && weight_decay != 0.0f)
|
|
g += weight_decay * data[j];
|
|
|
|
float mj = beta1 * m[j] + (1.0f - beta1) * g;
|
|
float vj = beta2 * v[j] + (1.0f - beta2) * g * g;
|
|
m[j] = mj;
|
|
v[j] = vj;
|
|
|
|
float m_hat = mj / bc1;
|
|
float v_hat = vj / bc2;
|
|
float update = lr * m_hat / (sqrtf(v_hat) + eps);
|
|
|
|
float d = data[j];
|
|
if (is_adamw && weight_decay != 0.0f) /* decoupled weight decay */
|
|
d -= lr * weight_decay * d;
|
|
d -= update;
|
|
data[j] = d;
|
|
}
|
|
|
|
/* ---------------------------------------------------------------------------
|
|
* opt#3 graph variant of fc_adamw_update: identical math, but the per-step
|
|
* scalars lr / bc1 / bc2 (cosine LR + t-dependent bias correction, all change
|
|
* every step) are read from a device buffer opt[3] = {lr, bc1, bc2}. The host
|
|
* computes them EXACTLY as before (powf on the CPU) and stages them into that
|
|
* buffer each step before replay, so the float bits are identical — delivering
|
|
* them via memory instead of by-value changes no arithmetic (bit-for-bit parity).
|
|
* ------------------------------------------------------------------------- */
|
|
extern "C" __global__ void fc_adamw_update_g(float *data, const float *grad,
|
|
float *m, float *v, int64_t n,
|
|
const float *opt, float beta1,
|
|
float beta2, float eps,
|
|
float weight_decay, int is_adamw) {
|
|
int64_t j = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
|
|
if (j >= n) return;
|
|
float lr = opt[0], bc1 = opt[1], bc2 = opt[2];
|
|
|
|
float g = grad[j];
|
|
if (!is_adamw && weight_decay != 0.0f)
|
|
g += weight_decay * data[j];
|
|
|
|
float mj = beta1 * m[j] + (1.0f - beta1) * g;
|
|
float vj = beta2 * v[j] + (1.0f - beta2) * g * g;
|
|
m[j] = mj;
|
|
v[j] = vj;
|
|
|
|
float m_hat = mj / bc1;
|
|
float v_hat = vj / bc2;
|
|
float update = lr * m_hat / (sqrtf(v_hat) + eps);
|
|
|
|
float d = data[j];
|
|
if (is_adamw && weight_decay != 0.0f)
|
|
d -= lr * weight_decay * d;
|
|
d -= update;
|
|
data[j] = d;
|
|
}
|
|
|
|
/* ===========================================================================
|
|
* Host launchers
|
|
* ========================================================================= */
|
|
|
|
#define CUDA_OK(call) do { cudaError_t _e = (call); if (_e != cudaSuccess) { \
|
|
fprintf(stderr, "CUDA error %s at %s:%d\n", cudaGetErrorString(_e), \
|
|
__FILE__, __LINE__); return -1; } } while (0)
|
|
|
|
#include <stdio.h>
|
|
|
|
/*
|
|
* One SR4 step on GPU 0. q_io[0..numel) is updated in place (int8, clamped to
|
|
* [-q_clamp,+q_clamp]). Mirrors ternary_optim_step(opt, grad_sum, batch_size,
|
|
* batch_counter, target_step) exactly. Returns 0 on success.
|
|
*/
|
|
extern "C" int sr4_step_cuda(int8_t *q_io, const float *grad_sum, int numel,
|
|
int batch_size, uint64_t batch_counter,
|
|
float target_step, uint64_t seed, int64_t idx_base,
|
|
int q_clamp, int use_dither) {
|
|
if (numel <= 0 || batch_size <= 0) return -1;
|
|
float inv_bs = 1.0f / (float)batch_size;
|
|
|
|
int8_t *d_q = NULL;
|
|
float *d_g = NULL, *d_s = NULL;
|
|
CUDA_OK(cudaMalloc(&d_q, (size_t)numel * sizeof(int8_t)));
|
|
CUDA_OK(cudaMalloc(&d_g, (size_t)numel * sizeof(float)));
|
|
CUDA_OK(cudaMalloc(&d_s, sizeof(float)));
|
|
CUDA_OK(cudaMemcpy(d_q, q_io, (size_t)numel * sizeof(int8_t), cudaMemcpyHostToDevice));
|
|
CUDA_OK(cudaMemcpy(d_g, grad_sum, (size_t)numel * sizeof(float), cudaMemcpyHostToDevice));
|
|
|
|
fc_sr4_scale<<<1, 1>>>(d_g, numel, inv_bs, target_step, d_s);
|
|
|
|
int threads = 256;
|
|
int blocks = (numel + threads - 1) / threads;
|
|
fc_sr4_update<<<blocks, threads>>>(d_q, d_g, numel, inv_bs, d_s, seed,
|
|
idx_base, batch_counter, q_clamp, use_dither);
|
|
CUDA_OK(cudaGetLastError());
|
|
CUDA_OK(cudaDeviceSynchronize());
|
|
|
|
CUDA_OK(cudaMemcpy(q_io, d_q, (size_t)numel * sizeof(int8_t), cudaMemcpyDeviceToHost));
|
|
cudaFree(d_q); cudaFree(d_g); cudaFree(d_s);
|
|
return 0;
|
|
}
|
|
|
|
/*
|
|
* One AdamW step on GPU 0. data[0..n), m[0..n), v[0..n) updated in place.
|
|
* `t` is the 1-based step index for bias correction (CPU: step_count+1).
|
|
* Mirrors optim.c adam_step() on the FO_ADAMW path. Returns 0 on success.
|
|
*/
|
|
extern "C" int adamw_step_cuda(float *data, const float *grad, float *m, float *v,
|
|
int64_t n, int t, float lr, float beta1,
|
|
float beta2, float eps, float weight_decay,
|
|
int is_adamw) {
|
|
if (n <= 0 || t < 1) return -1;
|
|
float bc1 = 1.0f - powf(beta1, (float)t);
|
|
float bc2 = 1.0f - powf(beta2, (float)t);
|
|
|
|
float *d_data = NULL, *d_grad = NULL, *d_m = NULL, *d_v = NULL;
|
|
size_t bytes = (size_t)n * sizeof(float);
|
|
CUDA_OK(cudaMalloc(&d_data, bytes));
|
|
CUDA_OK(cudaMalloc(&d_grad, bytes));
|
|
CUDA_OK(cudaMalloc(&d_m, bytes));
|
|
CUDA_OK(cudaMalloc(&d_v, bytes));
|
|
CUDA_OK(cudaMemcpy(d_data, data, bytes, cudaMemcpyHostToDevice));
|
|
CUDA_OK(cudaMemcpy(d_grad, grad, bytes, cudaMemcpyHostToDevice));
|
|
CUDA_OK(cudaMemcpy(d_m, m, bytes, cudaMemcpyHostToDevice));
|
|
CUDA_OK(cudaMemcpy(d_v, v, bytes, cudaMemcpyHostToDevice));
|
|
|
|
int threads = 256;
|
|
int64_t blocks = (n + threads - 1) / threads;
|
|
fc_adamw_update<<<(unsigned)blocks, threads>>>(d_data, d_grad, d_m, d_v, n, lr,
|
|
beta1, beta2, eps, weight_decay,
|
|
bc1, bc2, is_adamw);
|
|
CUDA_OK(cudaGetLastError());
|
|
CUDA_OK(cudaDeviceSynchronize());
|
|
|
|
CUDA_OK(cudaMemcpy(data, d_data, bytes, cudaMemcpyDeviceToHost));
|
|
CUDA_OK(cudaMemcpy(m, d_m, bytes, cudaMemcpyDeviceToHost));
|
|
CUDA_OK(cudaMemcpy(v, d_v, bytes, cudaMemcpyDeviceToHost));
|
|
cudaFree(d_data); cudaFree(d_grad); cudaFree(d_m); cudaFree(d_v);
|
|
return 0;
|
|
}
|