mirror of https://github.com/maderix/ANE.git
257 lines
12 KiB
C
257 lines
12 KiB
C
// model.h — Stories110M model struct + weight loading + ANE kernel compilation
|
|
// Training version: baked-weight conv kernels, recompile when weights update
|
|
#pragma once
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <string.h>
|
|
#include <math.h>
|
|
#include "ane_runtime.h"
|
|
#include "ane_mil_gen.h"
|
|
|
|
#define N_LAYERS 12
|
|
#define DIM 768
|
|
#define HIDDEN_DIM 2048
|
|
#define N_HEADS 12
|
|
#define HEAD_DIM 64
|
|
#define VOCAB_SIZE 32000
|
|
#define MAX_SEQ 1024
|
|
|
|
typedef struct {
|
|
int dim, hidden_dim, n_layers, n_heads, n_kv_heads, vocab_size, seq_len;
|
|
} Config;
|
|
|
|
typedef struct {
|
|
Config cfg;
|
|
int seq_len; // training sequence length
|
|
|
|
// Raw weights (f32)
|
|
float *token_embedding; // [vocab_size, dim]
|
|
float *rms_att_w[N_LAYERS]; // [dim]
|
|
float *wq[N_LAYERS]; // [dim, dim]
|
|
float *wk[N_LAYERS]; // [dim, dim]
|
|
float *wv[N_LAYERS]; // [dim, dim]
|
|
float *wo[N_LAYERS]; // [dim, dim]
|
|
float *rms_ffn_w[N_LAYERS]; // [dim]
|
|
float *w1[N_LAYERS]; // [hidden_dim, dim]
|
|
float *w2[N_LAYERS]; // [dim, hidden_dim]
|
|
float *w3[N_LAYERS]; // [hidden_dim, dim]
|
|
float *rms_final_w; // [dim]
|
|
float *wcls; // [vocab_size, dim]
|
|
|
|
// Per-layer ANE conv kernels (baked weights, recompiled on update)
|
|
ANEKernel *kern_q[N_LAYERS]; // Q projection: dim→dim
|
|
ANEKernel *kern_k[N_LAYERS]; // K projection: dim→dim
|
|
ANEKernel *kern_v[N_LAYERS]; // V projection: dim→dim
|
|
ANEKernel *kern_o[N_LAYERS]; // O projection: dim→dim
|
|
ANEKernel *kern_w1[N_LAYERS]; // FFN w1: dim→hidden
|
|
ANEKernel *kern_w2[N_LAYERS]; // FFN w2: hidden→dim
|
|
ANEKernel *kern_w3[N_LAYERS]; // FFN w3: dim→hidden
|
|
ANEKernel *kern_cls; // Classifier: dim→vocab
|
|
|
|
// Gradient accumulators (f32)
|
|
float *grad_wq[N_LAYERS], *grad_wk[N_LAYERS], *grad_wv[N_LAYERS], *grad_wo[N_LAYERS];
|
|
float *grad_w1[N_LAYERS], *grad_w2[N_LAYERS], *grad_w3[N_LAYERS];
|
|
float *grad_wcls;
|
|
float *grad_emb;
|
|
|
|
// Adam optimizer state
|
|
float *adam_m, *adam_v;
|
|
int adam_step;
|
|
size_t total_params;
|
|
|
|
// Activation cache for backward
|
|
float *act_x[N_LAYERS];
|
|
float *act_xnorm[N_LAYERS];
|
|
float *act_q[N_LAYERS];
|
|
float *act_k[N_LAYERS];
|
|
float *act_v[N_LAYERS];
|
|
float *act_attn_out[N_LAYERS];
|
|
float *act_ffn_in[N_LAYERS];
|
|
float *act_h1[N_LAYERS];
|
|
float *act_h3[N_LAYERS];
|
|
float *act_silu[N_LAYERS];
|
|
float *act_final;
|
|
float *act_pre_final;
|
|
float *logits;
|
|
} Model;
|
|
|
|
static int model_load_weights(Model *m, const char *path) {
|
|
FILE *f = fopen(path, "rb");
|
|
if (!f) { fprintf(stderr, "Cannot open %s\n", path); return -1; }
|
|
fread(&m->cfg, sizeof(Config), 1, f);
|
|
bool shared = m->cfg.vocab_size > 0;
|
|
if (m->cfg.vocab_size < 0) m->cfg.vocab_size = -m->cfg.vocab_size;
|
|
|
|
printf("Model: dim=%d hidden=%d layers=%d heads=%d vocab=%d seq=%d\n",
|
|
m->cfg.dim, m->cfg.hidden_dim, m->cfg.n_layers, m->cfg.n_heads,
|
|
m->cfg.vocab_size, m->cfg.seq_len);
|
|
|
|
int d = m->cfg.dim, hd = m->cfg.hidden_dim, nl = m->cfg.n_layers, vs = m->cfg.vocab_size;
|
|
|
|
m->token_embedding = (float*)malloc(vs * d * sizeof(float));
|
|
fread(m->token_embedding, sizeof(float), vs * d, f);
|
|
|
|
float *rms_att_all = (float*)malloc(nl * d * sizeof(float));
|
|
float *wq_all = (float*)malloc(nl * d * d * sizeof(float));
|
|
float *wk_all = (float*)malloc(nl * d * d * sizeof(float));
|
|
float *wv_all = (float*)malloc(nl * d * d * sizeof(float));
|
|
float *wo_all = (float*)malloc(nl * d * d * sizeof(float));
|
|
float *rms_ffn_all = (float*)malloc(nl * d * sizeof(float));
|
|
float *w1_all = (float*)malloc(nl * hd * d * sizeof(float));
|
|
float *w2_all = (float*)malloc(nl * d * hd * sizeof(float));
|
|
float *w3_all = (float*)malloc(nl * hd * d * sizeof(float));
|
|
|
|
fread(rms_att_all, sizeof(float), nl * d, f);
|
|
fread(wq_all, sizeof(float), nl * d * d, f);
|
|
fread(wk_all, sizeof(float), nl * d * d, f);
|
|
fread(wv_all, sizeof(float), nl * d * d, f);
|
|
fread(wo_all, sizeof(float), nl * d * d, f);
|
|
fread(rms_ffn_all, sizeof(float), nl * d, f);
|
|
fread(w1_all, sizeof(float), nl * hd * d, f);
|
|
fread(w2_all, sizeof(float), nl * d * hd, f);
|
|
fread(w3_all, sizeof(float), nl * hd * d, f);
|
|
|
|
for (int l = 0; l < nl; l++) {
|
|
m->rms_att_w[l] = (float*)malloc(d * sizeof(float));
|
|
memcpy(m->rms_att_w[l], rms_att_all + l*d, d * sizeof(float));
|
|
m->wq[l] = (float*)malloc(d*d*sizeof(float));
|
|
memcpy(m->wq[l], wq_all + l*d*d, d*d*sizeof(float));
|
|
m->wk[l] = (float*)malloc(d*d*sizeof(float));
|
|
memcpy(m->wk[l], wk_all + l*d*d, d*d*sizeof(float));
|
|
m->wv[l] = (float*)malloc(d*d*sizeof(float));
|
|
memcpy(m->wv[l], wv_all + l*d*d, d*d*sizeof(float));
|
|
m->wo[l] = (float*)malloc(d*d*sizeof(float));
|
|
memcpy(m->wo[l], wo_all + l*d*d, d*d*sizeof(float));
|
|
m->rms_ffn_w[l] = (float*)malloc(d * sizeof(float));
|
|
memcpy(m->rms_ffn_w[l], rms_ffn_all + l*d, d * sizeof(float));
|
|
m->w1[l] = (float*)malloc(hd*d*sizeof(float));
|
|
memcpy(m->w1[l], w1_all + l*hd*d, hd*d*sizeof(float));
|
|
m->w2[l] = (float*)malloc(d*hd*sizeof(float));
|
|
memcpy(m->w2[l], w2_all + l*d*hd, d*hd*sizeof(float));
|
|
m->w3[l] = (float*)malloc(hd*d*sizeof(float));
|
|
memcpy(m->w3[l], w3_all + l*hd*d, hd*d*sizeof(float));
|
|
}
|
|
free(rms_att_all); free(wq_all); free(wk_all); free(wv_all); free(wo_all);
|
|
free(rms_ffn_all); free(w1_all); free(w2_all); free(w3_all);
|
|
|
|
m->rms_final_w = (float*)malloc(d * sizeof(float));
|
|
fread(m->rms_final_w, sizeof(float), d, f);
|
|
|
|
if (shared) {
|
|
m->wcls = m->token_embedding;
|
|
} else {
|
|
m->wcls = (float*)malloc(vs * d * sizeof(float));
|
|
fread(m->wcls, sizeof(float), vs * d, f);
|
|
}
|
|
fclose(f);
|
|
return 0;
|
|
}
|
|
|
|
// Compile a single baked-weight conv kernel
|
|
static ANEKernel *compile_conv_kernel(const float *weights, int in_ch, int out_ch, int spatial) {
|
|
NSData *wb = mil_build_weight_blob(weights, out_ch, in_ch);
|
|
NSString *mil = mil_gen_conv(in_ch, out_ch, spatial);
|
|
size_t inBytes = (size_t)in_ch * spatial * 4;
|
|
size_t outBytes = (size_t)out_ch * spatial * 4;
|
|
return ane_compile([mil dataUsingEncoding:NSUTF8StringEncoding], wb, 1, &inBytes, 1, &outBytes);
|
|
}
|
|
|
|
// Compile all per-layer ANE kernels with current weights
|
|
static int model_compile_kernels(Model *m, int seq_len) {
|
|
m->seq_len = seq_len;
|
|
int d = m->cfg.dim, hd = m->cfg.hidden_dim, vs = m->cfg.vocab_size;
|
|
int S = seq_len;
|
|
printf("Compiling %d ANE conv kernels (S=%d)...\n", N_LAYERS * 7 + 1, S);
|
|
|
|
for (int l = 0; l < N_LAYERS; l++) {
|
|
m->kern_q[l] = compile_conv_kernel(m->wq[l], d, d, S);
|
|
m->kern_k[l] = compile_conv_kernel(m->wk[l], d, d, S);
|
|
m->kern_v[l] = compile_conv_kernel(m->wv[l], d, d, S);
|
|
m->kern_o[l] = compile_conv_kernel(m->wo[l], d, d, S);
|
|
m->kern_w1[l] = compile_conv_kernel(m->w1[l], d, hd, S);
|
|
m->kern_w2[l] = compile_conv_kernel(m->w2[l], hd, d, S);
|
|
m->kern_w3[l] = compile_conv_kernel(m->w3[l], d, hd, S);
|
|
if (!m->kern_q[l]) { fprintf(stderr, "L%d kern_q fail\n",l); return -1; }
|
|
if (!m->kern_k[l]) { fprintf(stderr, "L%d kern_k fail\n",l); return -1; }
|
|
if (!m->kern_v[l]) { fprintf(stderr, "L%d kern_v fail\n",l); return -1; }
|
|
if (!m->kern_o[l]) { fprintf(stderr, "L%d kern_o fail\n",l); return -1; }
|
|
if (!m->kern_w1[l]) { fprintf(stderr, "L%d kern_w1 fail\n",l); return -1; }
|
|
if (!m->kern_w2[l]) { fprintf(stderr, "L%d kern_w2 fail\n",l); return -1; }
|
|
if (!m->kern_w3[l]) { fprintf(stderr, "L%d kern_w3 fail\n",l); return -1; }
|
|
printf(" Layer %d OK\n", l);
|
|
}
|
|
m->kern_cls = compile_conv_kernel(m->wcls, d, vs, S);
|
|
if (!m->kern_cls) {
|
|
fprintf(stderr, "Classifier kernel compile failed (dim=%d→vocab=%d too large?), using CPU for cls\n", d, vs);
|
|
}
|
|
printf(" All kernels compiled (%d conv + %s)\n", N_LAYERS * 7, m->kern_cls ? "cls" : "cls=CPU");
|
|
return 0;
|
|
}
|
|
|
|
// Recompile all kernels after weight update — unload all first to avoid ANE model limit
|
|
static int model_recompile_kernels(Model *m) {
|
|
int d = m->cfg.dim, hd = m->cfg.hidden_dim, vs = m->cfg.vocab_size;
|
|
int S = m->seq_len;
|
|
// Phase 1: unload+free all
|
|
for (int l = 0; l < N_LAYERS; l++) {
|
|
ane_free(m->kern_q[l]); ane_free(m->kern_k[l]); ane_free(m->kern_v[l]); ane_free(m->kern_o[l]);
|
|
ane_free(m->kern_w1[l]); ane_free(m->kern_w2[l]); ane_free(m->kern_w3[l]);
|
|
m->kern_q[l]=m->kern_k[l]=m->kern_v[l]=m->kern_o[l]=NULL;
|
|
m->kern_w1[l]=m->kern_w2[l]=m->kern_w3[l]=NULL;
|
|
}
|
|
if (m->kern_cls) { ane_free(m->kern_cls); m->kern_cls=NULL; }
|
|
// Phase 2: recompile all
|
|
for (int l = 0; l < N_LAYERS; l++) {
|
|
m->kern_q[l] = compile_conv_kernel(m->wq[l], d, d, S);
|
|
m->kern_k[l] = compile_conv_kernel(m->wk[l], d, d, S);
|
|
m->kern_v[l] = compile_conv_kernel(m->wv[l], d, d, S);
|
|
m->kern_o[l] = compile_conv_kernel(m->wo[l], d, d, S);
|
|
m->kern_w1[l] = compile_conv_kernel(m->w1[l], d, hd, S);
|
|
m->kern_w2[l] = compile_conv_kernel(m->w2[l], hd, d, S);
|
|
m->kern_w3[l] = compile_conv_kernel(m->w3[l], d, hd, S);
|
|
if (!m->kern_q[l] || !m->kern_k[l] || !m->kern_v[l] || !m->kern_o[l] ||
|
|
!m->kern_w1[l] || !m->kern_w2[l] || !m->kern_w3[l]) return -1;
|
|
}
|
|
m->kern_cls = compile_conv_kernel(m->wcls, d, vs, S);
|
|
// cls may fail for large vocab — that's OK, forward uses CPU fallback
|
|
return 0;
|
|
}
|
|
|
|
static void model_alloc_training(Model *m) {
|
|
int d = m->cfg.dim, hd = m->cfg.hidden_dim, vs = m->cfg.vocab_size, S = m->seq_len;
|
|
for (int l = 0; l < N_LAYERS; l++) {
|
|
m->act_x[l] = (float*)calloc(S * d, sizeof(float));
|
|
m->act_xnorm[l] = (float*)calloc(S * d, sizeof(float));
|
|
m->act_q[l] = (float*)calloc(S * d, sizeof(float));
|
|
m->act_k[l] = (float*)calloc(S * d, sizeof(float));
|
|
m->act_v[l] = (float*)calloc(S * d, sizeof(float));
|
|
m->act_attn_out[l] = (float*)calloc(S * d, sizeof(float));
|
|
m->act_ffn_in[l] = (float*)calloc(S * d, sizeof(float));
|
|
m->act_h1[l] = (float*)calloc(S * hd, sizeof(float));
|
|
m->act_h3[l] = (float*)calloc(S * hd, sizeof(float));
|
|
m->act_silu[l] = (float*)calloc(S * hd, sizeof(float));
|
|
|
|
m->grad_wq[l] = (float*)calloc(d * d, sizeof(float));
|
|
m->grad_wk[l] = (float*)calloc(d * d, sizeof(float));
|
|
m->grad_wv[l] = (float*)calloc(d * d, sizeof(float));
|
|
m->grad_wo[l] = (float*)calloc(d * d, sizeof(float));
|
|
m->grad_w1[l] = (float*)calloc(hd * d, sizeof(float));
|
|
m->grad_w2[l] = (float*)calloc(d * hd, sizeof(float));
|
|
m->grad_w3[l] = (float*)calloc(hd * d, sizeof(float));
|
|
}
|
|
m->act_final = (float*)calloc(S * d, sizeof(float));
|
|
m->act_pre_final = (float*)calloc(S * d, sizeof(float));
|
|
m->logits = (float*)calloc(S * vs, sizeof(float));
|
|
m->grad_wcls = (float*)calloc(vs * d, sizeof(float));
|
|
m->grad_emb = (float*)calloc(vs * d, sizeof(float));
|
|
|
|
m->total_params = 0;
|
|
for (int l = 0; l < N_LAYERS; l++)
|
|
m->total_params += 4*(size_t)d*d + 2*(size_t)hd*d + (size_t)d*hd;
|
|
m->total_params += (size_t)vs * d * 2;
|
|
m->adam_m = (float*)calloc(m->total_params, sizeof(float));
|
|
m->adam_v = (float*)calloc(m->total_params, sizeof(float));
|
|
m->adam_step = 0;
|
|
printf("Total trainable params: %zu (%.1f M)\n", m->total_params, m->total_params/1e6);
|
|
}
|