402 lines
16 KiB
C
402 lines
16 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
#include "ft_layer_ops.h"
|
|
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <string.h>
|
|
#include <math.h>
|
|
|
|
/* ============================================================= */
|
|
/* RMSNorm */
|
|
/* ============================================================= */
|
|
|
|
void ft_rmsnorm(const float *x, const float *weight, float *out,
|
|
int n, float eps) {
|
|
double sum_sq = 0.0;
|
|
for (int i = 0; i < n; i++)
|
|
sum_sq += (double)x[i] * x[i];
|
|
float rms = sqrtf((float)(sum_sq / n) + eps);
|
|
float inv_rms = 1.0f / rms;
|
|
for (int i = 0; i < n; i++)
|
|
out[i] = weight[i] * x[i] * inv_rms;
|
|
}
|
|
|
|
void ft_rmsnorm_backward(const float *x, const float *weight,
|
|
const float *d_out, float *d_x,
|
|
int n, float eps) {
|
|
/* Forward recomputation */
|
|
double sum_sq = 0.0;
|
|
for (int i = 0; i < n; i++)
|
|
sum_sq += (double)x[i] * x[i];
|
|
float variance = (float)(sum_sq / n);
|
|
float inv_rms = 1.0f / sqrtf(variance + eps);
|
|
|
|
/* d_x[i] = weight[i] * d_out[i] * inv_rms
|
|
* - x[i] * inv_rms^3 * (1/n) * sum_j(weight[j] * d_out[j] * x[j]) */
|
|
double dot = 0.0;
|
|
for (int i = 0; i < n; i++)
|
|
dot += (double)weight[i] * d_out[i] * x[i];
|
|
float coeff = (float)(dot / n) * inv_rms * inv_rms * inv_rms;
|
|
|
|
for (int i = 0; i < n; i++)
|
|
d_x[i] += weight[i] * d_out[i] * inv_rms - x[i] * coeff;
|
|
}
|
|
|
|
/* ============================================================= */
|
|
/* SiLU */
|
|
/* ============================================================= */
|
|
|
|
static float sigmoid_f(float x) {
|
|
return 1.0f / (1.0f + expf(-x));
|
|
}
|
|
|
|
void ft_silu(const float *x, float *out, int n) {
|
|
for (int i = 0; i < n; i++)
|
|
out[i] = x[i] * sigmoid_f(x[i]);
|
|
}
|
|
|
|
void ft_silu_backward(const float *x, const float *d_out, float *d_x, int n) {
|
|
for (int i = 0; i < n; i++) {
|
|
float s = sigmoid_f(x[i]);
|
|
d_x[i] += d_out[i] * (s + x[i] * s * (1.0f - s));
|
|
}
|
|
}
|
|
|
|
/* ============================================================= */
|
|
/* Ternary Matvec (using bit-packed sign+mask) */
|
|
/* ============================================================= */
|
|
|
|
void ft_ternary_matvec(const ft_layer_state *state, int64_t state_offset,
|
|
int rows, int cols,
|
|
const float *x, float *out) {
|
|
/*
|
|
* out[row] = sum_k( ternary[row,k] * scale_group(row*cols+k) * x[k] )
|
|
*
|
|
* ternary value: mask ? (sign ? +1 : -1) : 0
|
|
* scale from per-group array
|
|
*/
|
|
for (int row = 0; row < rows; row++) {
|
|
float sum = 0.0f;
|
|
for (int k = 0; k < cols; k++) {
|
|
int64_t gi = state_offset + (int64_t)row * cols + k;
|
|
int64_t byte_idx = gi / 8;
|
|
int bit = (int)(gi % 8);
|
|
|
|
int m = (state->mask[byte_idx] >> bit) & 1;
|
|
if (!m) continue;
|
|
|
|
int s = (state->sign[byte_idx] >> bit) & 1;
|
|
float w_val = s ? 1.0f : -1.0f;
|
|
|
|
/* Per-group scale */
|
|
int group_idx = (int)(gi / state->group_size);
|
|
float scale = state->scales[group_idx];
|
|
|
|
sum += w_val * scale * x[k];
|
|
}
|
|
out[row] = sum;
|
|
}
|
|
}
|
|
|
|
void ft_ternary_matvec_backward_input(const ft_layer_state *state, int64_t state_offset,
|
|
int rows, int cols,
|
|
const float *d_out, float *d_x) {
|
|
/* d_x[k] += sum_row( ternary[row,k] * scale * d_out[row] ) */
|
|
for (int row = 0; row < rows; row++) {
|
|
if (d_out[row] == 0.0f) continue;
|
|
for (int k = 0; k < cols; k++) {
|
|
int64_t gi = state_offset + (int64_t)row * cols + k;
|
|
int64_t byte_idx = gi / 8;
|
|
int bit = (int)(gi % 8);
|
|
|
|
int m = (state->mask[byte_idx] >> bit) & 1;
|
|
if (!m) continue;
|
|
|
|
int s = (state->sign[byte_idx] >> bit) & 1;
|
|
float w_val = s ? 1.0f : -1.0f;
|
|
|
|
int group_idx = (int)(gi / state->group_size);
|
|
float scale = state->scales[group_idx];
|
|
|
|
d_x[k] += w_val * scale * d_out[row];
|
|
}
|
|
}
|
|
}
|
|
|
|
void ft_ternary_matvec_backward_weight(const ft_layer_state *state, int64_t state_offset,
|
|
int rows, int cols,
|
|
const float *x, const float *d_out,
|
|
ft_grad_stats *stats) {
|
|
(void)state; /* Weight values not needed — we accumulate FP32 gradients */
|
|
|
|
/* For each weight [row, col]:
|
|
* grad = x[col] * d_out[row]
|
|
* accumulate into grad_sum, grad_sq_sum, hessian_diag */
|
|
for (int row = 0; row < rows; row++) {
|
|
for (int k = 0; k < cols; k++) {
|
|
int64_t gi = state_offset + (int64_t)row * cols + k;
|
|
float grad = x[k] * d_out[row];
|
|
|
|
stats->grad_sum[gi] += grad;
|
|
stats->grad_sq_sum[gi] += grad * grad;
|
|
|
|
/* Gauss-Newton Hessian diagonal approximation:
|
|
* h_diag ~= x[k]^2 * d_out[row]^2 (outer product of squared terms) */
|
|
stats->hessian_diag[gi] += x[k] * x[k] * d_out[row] * d_out[row];
|
|
}
|
|
}
|
|
}
|
|
|
|
/* ============================================================= */
|
|
/* Layer Forward (LLaMA-family) */
|
|
/* ============================================================= */
|
|
|
|
void ft_layer_forward(const ft_layer_state *state, const ft_config *cfg,
|
|
const float *input, float *output, int batch_size) {
|
|
int n = cfg->n_embd;
|
|
int q_dim = cfg->n_heads * cfg->head_dim;
|
|
int kv_dim = cfg->n_kv_heads * cfg->head_dim;
|
|
|
|
for (int b = 0; b < batch_size; b++) {
|
|
const float *x = input + (int64_t)b * n;
|
|
float *y = output + (int64_t)b * n;
|
|
|
|
/* --- Attention sub-layer --- */
|
|
|
|
/* RMSNorm */
|
|
float *x_norm = calloc(n, sizeof(float));
|
|
ft_rmsnorm(x, state->norm_attn, x_norm, n, cfg->rms_eps);
|
|
|
|
/* Q/K/V projections */
|
|
float *q = calloc(q_dim, sizeof(float));
|
|
float *k = calloc(kv_dim, sizeof(float));
|
|
float *v = calloc(kv_dim, sizeof(float));
|
|
|
|
ft_ternary_matvec(state, state->weights[0].offset, q_dim, n, x_norm, q);
|
|
ft_ternary_matvec(state, state->weights[1].offset, kv_dim, n, x_norm, k);
|
|
ft_ternary_matvec(state, state->weights[2].offset, kv_dim, n, x_norm, v);
|
|
|
|
/* Simplified attention: for ternarizer calibration, we use a proxy.
|
|
* Full multi-head attention with RoPE is expensive and not needed
|
|
* for gradient estimation. Use V directly as the attention output.
|
|
* This is a simplification — the gradient stats are still meaningful
|
|
* because they capture the sensitivity of each weight to loss. */
|
|
float *attn_out = calloc(q_dim, sizeof(float));
|
|
/* Expand KV if GQA: replicate kv_heads to match n_heads */
|
|
int repeat = cfg->n_heads / cfg->n_kv_heads;
|
|
for (int h = 0; h < cfg->n_heads; h++) {
|
|
int kv_h = h / repeat;
|
|
memcpy(attn_out + h * cfg->head_dim,
|
|
v + kv_h * cfg->head_dim,
|
|
cfg->head_dim * sizeof(float));
|
|
}
|
|
|
|
/* O projection */
|
|
float *o_out = calloc(n, sizeof(float));
|
|
ft_ternary_matvec(state, state->weights[3].offset, n, q_dim, attn_out, o_out);
|
|
|
|
/* Residual */
|
|
float *h = calloc(n, sizeof(float));
|
|
for (int i = 0; i < n; i++)
|
|
h[i] = x[i] + o_out[i];
|
|
|
|
/* --- FFN sub-layer --- */
|
|
|
|
/* RMSNorm */
|
|
float *h_norm = calloc(n, sizeof(float));
|
|
ft_rmsnorm(h, state->norm_ffn, h_norm, n, cfg->rms_eps);
|
|
|
|
/* Gate + Up projections */
|
|
float *gate = calloc(cfg->ffn_dim, sizeof(float));
|
|
float *up = calloc(cfg->ffn_dim, sizeof(float));
|
|
|
|
ft_ternary_matvec(state, state->weights[4].offset, cfg->ffn_dim, n, h_norm, gate);
|
|
ft_ternary_matvec(state, state->weights[5].offset, cfg->ffn_dim, n, h_norm, up);
|
|
|
|
/* SiLU(gate) * up */
|
|
float *ffn_hidden = calloc(cfg->ffn_dim, sizeof(float));
|
|
ft_silu(gate, ffn_hidden, cfg->ffn_dim);
|
|
for (int i = 0; i < cfg->ffn_dim; i++)
|
|
ffn_hidden[i] *= up[i];
|
|
|
|
/* Down projection */
|
|
float *down_out = calloc(n, sizeof(float));
|
|
ft_ternary_matvec(state, state->weights[6].offset, n, cfg->ffn_dim, ffn_hidden, down_out);
|
|
|
|
/* Residual */
|
|
for (int i = 0; i < n; i++)
|
|
y[i] = h[i] + down_out[i];
|
|
|
|
free(x_norm); free(q); free(k); free(v);
|
|
free(attn_out); free(o_out); free(h);
|
|
free(h_norm); free(gate); free(up);
|
|
free(ffn_hidden); free(down_out);
|
|
}
|
|
}
|
|
|
|
/* ============================================================= */
|
|
/* Layer Backward (LLaMA-family) */
|
|
/* ============================================================= */
|
|
|
|
void ft_layer_backward(const ft_layer_state *state, const ft_config *cfg,
|
|
const float *input, const float *d_output,
|
|
float *d_input, ft_grad_stats *stats, int batch_size) {
|
|
int n = cfg->n_embd;
|
|
int q_dim = cfg->n_heads * cfg->head_dim;
|
|
int kv_dim = cfg->n_kv_heads * cfg->head_dim;
|
|
|
|
for (int b = 0; b < batch_size; b++) {
|
|
const float *x = input + (int64_t)b * n;
|
|
const float *dy = d_output + (int64_t)b * n;
|
|
|
|
/* Recompute forward intermediates */
|
|
float *x_norm = calloc(n, sizeof(float));
|
|
ft_rmsnorm(x, state->norm_attn, x_norm, n, cfg->rms_eps);
|
|
|
|
float *v = calloc(kv_dim, sizeof(float));
|
|
ft_ternary_matvec(state, state->weights[2].offset, kv_dim, n, x_norm, v);
|
|
|
|
float *attn_out = calloc(q_dim, sizeof(float));
|
|
int repeat = cfg->n_heads / cfg->n_kv_heads;
|
|
for (int h_i = 0; h_i < cfg->n_heads; h_i++) {
|
|
int kv_h = h_i / repeat;
|
|
memcpy(attn_out + h_i * cfg->head_dim,
|
|
v + kv_h * cfg->head_dim,
|
|
cfg->head_dim * sizeof(float));
|
|
}
|
|
|
|
float *o_out = calloc(n, sizeof(float));
|
|
ft_ternary_matvec(state, state->weights[3].offset, n, q_dim, attn_out, o_out);
|
|
|
|
float *h = calloc(n, sizeof(float));
|
|
for (int i = 0; i < n; i++) h[i] = x[i] + o_out[i];
|
|
|
|
float *h_norm = calloc(n, sizeof(float));
|
|
ft_rmsnorm(h, state->norm_ffn, h_norm, n, cfg->rms_eps);
|
|
|
|
float *gate = calloc(cfg->ffn_dim, sizeof(float));
|
|
float *up = calloc(cfg->ffn_dim, sizeof(float));
|
|
ft_ternary_matvec(state, state->weights[4].offset, cfg->ffn_dim, n, h_norm, gate);
|
|
ft_ternary_matvec(state, state->weights[5].offset, cfg->ffn_dim, n, h_norm, up);
|
|
|
|
float *gate_silu = calloc(cfg->ffn_dim, sizeof(float));
|
|
ft_silu(gate, gate_silu, cfg->ffn_dim);
|
|
|
|
float *ffn_hidden = calloc(cfg->ffn_dim, sizeof(float));
|
|
for (int i = 0; i < cfg->ffn_dim; i++)
|
|
ffn_hidden[i] = gate_silu[i] * up[i];
|
|
|
|
/* --- Backward through FFN sub-layer --- */
|
|
|
|
/* dy flows to both residual paths */
|
|
float *d_down_out = calloc(n, sizeof(float));
|
|
memcpy(d_down_out, dy, n * sizeof(float)); /* d_residual */
|
|
|
|
/* Backward through down_proj */
|
|
float *d_ffn_hidden = calloc(cfg->ffn_dim, sizeof(float));
|
|
ft_ternary_matvec_backward_input(state, state->weights[6].offset,
|
|
n, cfg->ffn_dim, d_down_out, d_ffn_hidden);
|
|
ft_ternary_matvec_backward_weight(state, state->weights[6].offset,
|
|
n, cfg->ffn_dim, ffn_hidden, d_down_out, stats);
|
|
|
|
/* Backward through SiLU(gate) * up */
|
|
float *d_gate_silu = calloc(cfg->ffn_dim, sizeof(float));
|
|
float *d_up = calloc(cfg->ffn_dim, sizeof(float));
|
|
for (int i = 0; i < cfg->ffn_dim; i++) {
|
|
d_gate_silu[i] = d_ffn_hidden[i] * up[i];
|
|
d_up[i] = d_ffn_hidden[i] * gate_silu[i];
|
|
}
|
|
|
|
float *d_gate = calloc(cfg->ffn_dim, sizeof(float));
|
|
ft_silu_backward(gate, d_gate_silu, d_gate, cfg->ffn_dim);
|
|
|
|
/* Backward through gate/up projections */
|
|
float *d_h_norm_ffn = calloc(n, sizeof(float));
|
|
ft_ternary_matvec_backward_input(state, state->weights[4].offset,
|
|
cfg->ffn_dim, n, d_gate, d_h_norm_ffn);
|
|
ft_ternary_matvec_backward_weight(state, state->weights[4].offset,
|
|
cfg->ffn_dim, n, h_norm, d_gate, stats);
|
|
|
|
float *d_h_norm_ffn2 = calloc(n, sizeof(float));
|
|
ft_ternary_matvec_backward_input(state, state->weights[5].offset,
|
|
cfg->ffn_dim, n, d_up, d_h_norm_ffn2);
|
|
ft_ternary_matvec_backward_weight(state, state->weights[5].offset,
|
|
cfg->ffn_dim, n, h_norm, d_up, stats);
|
|
|
|
/* Combine gate+up gradients */
|
|
for (int i = 0; i < n; i++) d_h_norm_ffn[i] += d_h_norm_ffn2[i];
|
|
|
|
/* Backward through FFN RMSNorm */
|
|
float *d_h = calloc(n, sizeof(float));
|
|
ft_rmsnorm_backward(h, state->norm_ffn, d_h_norm_ffn, d_h, n, cfg->rms_eps);
|
|
|
|
/* Add FFN residual gradient */
|
|
for (int i = 0; i < n; i++) d_h[i] += dy[i];
|
|
|
|
/* --- Backward through attention sub-layer --- */
|
|
|
|
float *d_o_out = calloc(n, sizeof(float));
|
|
memcpy(d_o_out, d_h, n * sizeof(float));
|
|
|
|
/* Backward through O projection */
|
|
float *d_attn_out = calloc(q_dim, sizeof(float));
|
|
ft_ternary_matvec_backward_input(state, state->weights[3].offset,
|
|
n, q_dim, d_o_out, d_attn_out);
|
|
ft_ternary_matvec_backward_weight(state, state->weights[3].offset,
|
|
n, q_dim, attn_out, d_o_out, stats);
|
|
|
|
/* Backward through simplified attention (V passthrough) */
|
|
float *d_v = calloc(kv_dim, sizeof(float));
|
|
for (int h_i = 0; h_i < cfg->n_heads; h_i++) {
|
|
int kv_h = h_i / repeat;
|
|
for (int d = 0; d < cfg->head_dim; d++)
|
|
d_v[kv_h * cfg->head_dim + d] += d_attn_out[h_i * cfg->head_dim + d];
|
|
}
|
|
|
|
/* Backward through Q/K/V projections.
|
|
* Q and K don't contribute to output in our simplified attention,
|
|
* but we still accumulate weight gradients for scoring. */
|
|
float *d_x_norm = calloc(n, sizeof(float));
|
|
|
|
/* V backward */
|
|
ft_ternary_matvec_backward_input(state, state->weights[2].offset,
|
|
kv_dim, n, d_v, d_x_norm);
|
|
ft_ternary_matvec_backward_weight(state, state->weights[2].offset,
|
|
kv_dim, n, x_norm, d_v, stats);
|
|
|
|
/* Q backward: use zero d_out for input grad but compute weight grad
|
|
* with a small proxy signal to get meaningful statistics */
|
|
float *d_q_proxy = calloc(q_dim, sizeof(float));
|
|
for (int i = 0; i < q_dim; i++) d_q_proxy[i] = d_attn_out[i] * 0.1f;
|
|
ft_ternary_matvec_backward_weight(state, state->weights[0].offset,
|
|
q_dim, n, x_norm, d_q_proxy, stats);
|
|
|
|
/* K backward: similar proxy */
|
|
float *d_k_proxy = calloc(kv_dim, sizeof(float));
|
|
for (int i = 0; i < kv_dim; i++) d_k_proxy[i] = d_v[i] * 0.1f;
|
|
ft_ternary_matvec_backward_weight(state, state->weights[1].offset,
|
|
kv_dim, n, x_norm, d_k_proxy, stats);
|
|
|
|
/* Backward through attention RMSNorm */
|
|
float *d_x_local = calloc(n, sizeof(float));
|
|
ft_rmsnorm_backward(x, state->norm_attn, d_x_norm, d_x_local, n, cfg->rms_eps);
|
|
|
|
/* Add attention residual gradient */
|
|
for (int i = 0; i < n; i++)
|
|
d_x_local[i] += d_h[i];
|
|
|
|
/* Accumulate into d_input */
|
|
for (int i = 0; i < n; i++)
|
|
d_input[(int64_t)b * n + i] += d_x_local[i];
|
|
|
|
/* Cleanup */
|
|
free(x_norm); free(v); free(attn_out); free(o_out); free(h);
|
|
free(h_norm); free(gate); free(up); free(gate_silu); free(ffn_hidden);
|
|
free(d_down_out); free(d_ffn_hidden); free(d_gate_silu); free(d_up);
|
|
free(d_gate); free(d_h_norm_ffn); free(d_h_norm_ffn2); free(d_h);
|
|
free(d_o_out); free(d_attn_out); free(d_v); free(d_x_norm);
|
|
free(d_q_proxy); free(d_k_proxy); free(d_x_local);
|
|
}
|
|
}
|