mirror of https://github.com/maderix/ANE.git
fix(HIGH-04): add xmf/xcf OOM-abort helpers, replace all malloc/calloc
- stories_config.h: add xmf(n)/xcf(n) static inline helpers that abort() with diagnostic on OOM — replaces all unchecked malloc/calloc - stories_config.h: replace all (float*)malloc(X*4) -> xmf(X) and (float*)calloc(X,4) -> xcf(X) in all 5 alloc functions - train_large.m: same replacement for direct alloc calls (embed, rms_final, grads, per-iteration temporaries — 31 call sites total) ref: docs/reports/security-audit-2026-03-02.md HIGH-04 Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
manni07
parent
b3936dce86
commit
78666fcc1c
|
|
@ -1,207 +1,218 @@
|
||||||
// stories_config.h — Stories110M model config and structures
|
// stories_config.h — Stories110M model config and structures
|
||||||
#pragma once
|
#pragma once
|
||||||
#import <Foundation/Foundation.h>
|
#import <Foundation/Foundation.h>
|
||||||
#import <objc/runtime.h>
|
#import <objc/runtime.h>
|
||||||
#import <objc/message.h>
|
#import <objc/message.h>
|
||||||
#import <dlfcn.h>
|
#import <dlfcn.h>
|
||||||
#import <IOSurface/IOSurface.h>
|
#import <IOSurface/IOSurface.h>
|
||||||
#import <mach/mach_time.h>
|
#import <mach/mach_time.h>
|
||||||
#import <Accelerate/Accelerate.h>
|
#import <Accelerate/Accelerate.h>
|
||||||
#include <math.h>
|
#include <math.h>
|
||||||
#include <unistd.h>
|
#include <unistd.h>
|
||||||
#include <dispatch/dispatch.h>
|
#include <dispatch/dispatch.h>
|
||||||
#include <sys/mman.h>
|
#include <sys/mman.h>
|
||||||
#include <sys/stat.h>
|
#include <sys/stat.h>
|
||||||
#include <fcntl.h>
|
#include <fcntl.h>
|
||||||
|
|
||||||
// Stories110M config
|
// Stories110M config
|
||||||
#define DIM 768
|
#define DIM 768
|
||||||
#define HIDDEN 2048
|
#define HIDDEN 2048
|
||||||
#define HEADS 12
|
#define HEADS 12
|
||||||
#define HD (DIM/HEADS)
|
#define HD (DIM/HEADS)
|
||||||
#define SEQ 256
|
#define SEQ 256
|
||||||
#define NLAYERS 12
|
#define NLAYERS 12
|
||||||
#define VOCAB 32000
|
#define VOCAB 32000
|
||||||
#define ACCUM_STEPS 10
|
#define ACCUM_STEPS 10
|
||||||
#define MAX_COMPILES 100
|
#define MAX_COMPILES 100
|
||||||
|
|
||||||
// Per compile: 5 weight-bearing kernels per layer + 1 classifier = 5*12+1 = 61
|
// Per compile: 5 weight-bearing kernels per layer + 1 classifier = 5*12+1 = 61
|
||||||
// Plus 1 static (sdpaBwd2 per layer, no weights) = 12 more but those are weight-free
|
// Plus 1 static (sdpaBwd2 per layer, no weights) = 12 more but those are weight-free
|
||||||
// Actually sdpaBwd2 has no weights, compile once per layer
|
// Actually sdpaBwd2 has no weights, compile once per layer
|
||||||
// Weight-bearing: fwdAttn(1) + fwdFFN(1) + ffnBwd(1) + sdpaBwd1(1) + qkvBwd(1) = 5 per layer
|
// Weight-bearing: fwdAttn(1) + fwdFFN(1) + ffnBwd(1) + sdpaBwd1(1) + qkvBwd(1) = 5 per layer
|
||||||
// 5 * 12 = 60 weight-bearing compiles per batch
|
// 5 * 12 = 60 weight-bearing compiles per batch
|
||||||
// With MAX_COMPILES=100, we get 1 batch of ACCUM_STEPS before restart
|
// With MAX_COMPILES=100, we get 1 batch of ACCUM_STEPS before restart
|
||||||
#define KERNELS_PER_LAYER 5
|
#define KERNELS_PER_LAYER 5
|
||||||
#define TOTAL_WEIGHT_KERNELS (KERNELS_PER_LAYER * NLAYERS)
|
#define TOTAL_WEIGHT_KERNELS (KERNELS_PER_LAYER * NLAYERS)
|
||||||
|
|
||||||
// Attention score channels for SDPA backward
|
// Attention score channels for SDPA backward
|
||||||
#define SCORE_CH (HEADS*SEQ)
|
#define SCORE_CH (HEADS*SEQ)
|
||||||
|
|
||||||
// Weight sizes per layer
|
// Weight sizes per layer
|
||||||
#define WQ_SZ (DIM*DIM)
|
#define WQ_SZ (DIM*DIM)
|
||||||
#define WO_SZ (DIM*DIM)
|
#define WO_SZ (DIM*DIM)
|
||||||
#define W1_SZ (HIDDEN*DIM)
|
#define W1_SZ (HIDDEN*DIM)
|
||||||
#define W2_SZ (DIM*HIDDEN)
|
#define W2_SZ (DIM*HIDDEN)
|
||||||
#define W3_SZ (HIDDEN*DIM)
|
#define W3_SZ (HIDDEN*DIM)
|
||||||
#define LAYER_PARAMS (4*WQ_SZ + W1_SZ + W2_SZ + W3_SZ + 2*DIM)
|
#define LAYER_PARAMS (4*WQ_SZ + W1_SZ + W2_SZ + W3_SZ + 2*DIM)
|
||||||
#define TOTAL_PARAMS (NLAYERS * LAYER_PARAMS + DIM + VOCAB*DIM) // +rms_final+embed
|
#define TOTAL_PARAMS (NLAYERS * LAYER_PARAMS + DIM + VOCAB*DIM) // +rms_final+embed
|
||||||
|
|
||||||
// Per-layer weight and optimizer state
|
// Per-layer weight and optimizer state
|
||||||
typedef struct {
|
typedef struct {
|
||||||
float *Wq, *Wk, *Wv, *Wo;
|
float *Wq, *Wk, *Wv, *Wo;
|
||||||
float *W1, *W2, *W3;
|
float *W1, *W2, *W3;
|
||||||
float *rms_att, *rms_ffn;
|
float *rms_att, *rms_ffn;
|
||||||
} LayerWeights;
|
} LayerWeights;
|
||||||
|
|
||||||
typedef struct {
|
typedef struct {
|
||||||
float *m, *v;
|
float *m, *v;
|
||||||
size_t n;
|
size_t n;
|
||||||
} AdamState;
|
} AdamState;
|
||||||
|
|
||||||
typedef struct {
|
typedef struct {
|
||||||
AdamState Wq, Wk, Wv, Wo;
|
AdamState Wq, Wk, Wv, Wo;
|
||||||
AdamState W1, W2, W3;
|
AdamState W1, W2, W3;
|
||||||
AdamState rms_att, rms_ffn;
|
AdamState rms_att, rms_ffn;
|
||||||
} LayerAdam;
|
} LayerAdam;
|
||||||
|
|
||||||
// Per-layer activation buffers (saved for backward)
|
// Per-layer activation buffers (saved for backward)
|
||||||
typedef struct {
|
typedef struct {
|
||||||
float *layer_in; // [DIM, SEQ] input to this layer (for rmsnorm1 bwd)
|
float *layer_in; // [DIM, SEQ] input to this layer (for rmsnorm1 bwd)
|
||||||
float *xnorm; // [DIM, SEQ] rmsnorm1 output
|
float *xnorm; // [DIM, SEQ] rmsnorm1 output
|
||||||
float *Q, *K, *V; // [DIM, SEQ] QKV projections
|
float *Q, *K, *V; // [DIM, SEQ] QKV projections
|
||||||
float *attn_out; // [DIM, SEQ] attention output (before Wo)
|
float *attn_out; // [DIM, SEQ] attention output (before Wo)
|
||||||
float *o_out; // [DIM, SEQ] Wo output
|
float *o_out; // [DIM, SEQ] Wo output
|
||||||
float *x2; // [DIM, SEQ] residual after attn
|
float *x2; // [DIM, SEQ] residual after attn
|
||||||
float *x2norm; // [DIM, SEQ] rmsnorm2 output
|
float *x2norm; // [DIM, SEQ] rmsnorm2 output
|
||||||
float *h1, *h3; // [HIDDEN, SEQ] FFN intermediates
|
float *h1, *h3; // [HIDDEN, SEQ] FFN intermediates
|
||||||
float *silu_out; // [HIDDEN, SEQ] SiLU(h1)*h3
|
float *silu_out; // [HIDDEN, SEQ] SiLU(h1)*h3
|
||||||
float *ffn_out; // [DIM, SEQ] FFN output
|
float *ffn_out; // [DIM, SEQ] FFN output
|
||||||
} LayerActs;
|
} LayerActs;
|
||||||
|
|
||||||
// Per-layer gradient accumulators
|
// Per-layer gradient accumulators
|
||||||
typedef struct {
|
typedef struct {
|
||||||
float *Wq, *Wk, *Wv, *Wo;
|
float *Wq, *Wk, *Wv, *Wo;
|
||||||
float *W1, *W2, *W3;
|
float *W1, *W2, *W3;
|
||||||
float *rms_att, *rms_ffn;
|
float *rms_att, *rms_ffn;
|
||||||
} LayerGrads;
|
} LayerGrads;
|
||||||
|
|
||||||
// ANE kernels per layer
|
// ANE kernels per layer
|
||||||
typedef struct { void *model; IOSurfaceRef ioIn, ioOut; void *request; void *tmpDir; } Kern;
|
typedef struct { void *model; IOSurfaceRef ioIn, ioOut; void *request; void *tmpDir; } Kern;
|
||||||
typedef struct {
|
typedef struct {
|
||||||
Kern *fwdAttn, *fwdFFN, *ffnBwd, *sdpaBwd1, *sdpaBwd2, *qkvBwd;
|
Kern *fwdAttn, *fwdFFN, *ffnBwd, *sdpaBwd1, *sdpaBwd2, *qkvBwd;
|
||||||
} LayerKernels;
|
} LayerKernels;
|
||||||
|
|
||||||
// Checkpoint header
|
// Checkpoint header
|
||||||
typedef struct {
|
typedef struct {
|
||||||
int magic; // 0x424C5A54 "BLZT"
|
int magic; // 0x424C5A54 "BLZT"
|
||||||
int version; // 2
|
int version; // 2
|
||||||
int step, total_steps;
|
int step, total_steps;
|
||||||
int n_layers, vocab_size, dim, hidden_dim, n_heads, seq_len;
|
int n_layers, vocab_size, dim, hidden_dim, n_heads, seq_len;
|
||||||
float lr, loss;
|
float lr, loss;
|
||||||
double cum_compile, cum_train, cum_wall;
|
double cum_compile, cum_train, cum_wall;
|
||||||
int cum_steps, cum_batches;
|
int cum_steps, cum_batches;
|
||||||
int adam_t;
|
int adam_t;
|
||||||
int pad[3]; // pad[0] = 0x01020304 (LE byte-order sentinel, MED-04); pad[1..2] = 0
|
int pad[3]; // pad[0] = 0x01020304 (LE byte-order sentinel, MED-04); pad[1..2] = 0
|
||||||
} CkptHdr;
|
} CkptHdr;
|
||||||
|
|
||||||
// llama2.c model file header
|
// llama2.c model file header
|
||||||
typedef struct {
|
typedef struct {
|
||||||
int dim, hidden_dim, n_layers, n_heads, n_kv_heads, vocab_size, seq_len;
|
int dim, hidden_dim, n_layers, n_heads, n_kv_heads, vocab_size, seq_len;
|
||||||
} Llama2Config;
|
} Llama2Config;
|
||||||
|
|
||||||
// Globals
|
// Globals
|
||||||
static Class g_D, g_I, g_AR, g_AIO;
|
static Class g_D, g_I, g_AR, g_AIO;
|
||||||
static bool g_ane_ok_large = false; // true only when all private classes loaded successfully
|
static bool g_ane_ok_large = false; // true only when all private classes loaded successfully
|
||||||
static mach_timebase_info_data_t g_tb;
|
static mach_timebase_info_data_t g_tb;
|
||||||
static int g_compile_count = 0;
|
static int g_compile_count = 0;
|
||||||
static int g_compile_seq = 0; // MED-02: per-call unique index for temp-dir naming
|
static int g_compile_seq = 0; // MED-02: per-call unique index for temp-dir naming
|
||||||
|
|
||||||
static void ane_init(void) {
|
static void ane_init(void) {
|
||||||
// MED-06: dispatch_once provides thread-safe one-time init with full memory barrier.
|
// MED-06: dispatch_once provides thread-safe one-time init with full memory barrier.
|
||||||
// Replaces manual g_ane_init_done bool guard which had a Check-Then-Act race.
|
// Replaces manual g_ane_init_done bool guard which had a Check-Then-Act race.
|
||||||
static dispatch_once_t ane_once_large;
|
static dispatch_once_t ane_once_large;
|
||||||
dispatch_once(&ane_once_large, ^{
|
dispatch_once(&ane_once_large, ^{
|
||||||
void *handle = dlopen(
|
void *handle = dlopen(
|
||||||
"/System/Library/PrivateFrameworks/AppleNeuralEngine.framework/AppleNeuralEngine",
|
"/System/Library/PrivateFrameworks/AppleNeuralEngine.framework/AppleNeuralEngine",
|
||||||
RTLD_NOW);
|
RTLD_NOW);
|
||||||
if (!handle) {
|
if (!handle) {
|
||||||
fprintf(stderr, "ANE: dlopen failed: %s\n", dlerror());
|
fprintf(stderr, "ANE: dlopen failed: %s\n", dlerror());
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
g_D = NSClassFromString(@"_ANEInMemoryModelDescriptor");
|
g_D = NSClassFromString(@"_ANEInMemoryModelDescriptor");
|
||||||
g_I = NSClassFromString(@"_ANEInMemoryModel");
|
g_I = NSClassFromString(@"_ANEInMemoryModel");
|
||||||
g_AR = NSClassFromString(@"_ANERequest");
|
g_AR = NSClassFromString(@"_ANERequest");
|
||||||
g_AIO= NSClassFromString(@"_ANEIOSurfaceObject");
|
g_AIO= NSClassFromString(@"_ANEIOSurfaceObject");
|
||||||
if (!g_D || !g_I || !g_AR || !g_AIO) {
|
if (!g_D || !g_I || !g_AR || !g_AIO) {
|
||||||
fprintf(stderr, "ANE: Private classes not found (macOS version mismatch?)\n");
|
fprintf(stderr, "ANE: Private classes not found (macOS version mismatch?)\n");
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
g_ane_ok_large = true; // dispatch_once guarantees memory barrier before completion
|
g_ane_ok_large = true; // dispatch_once guarantees memory barrier before completion
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
static double tb_ms(uint64_t t) { return (double)t * g_tb.numer / g_tb.denom / 1e6; }
|
static double tb_ms(uint64_t t) { return (double)t * g_tb.numer / g_tb.denom / 1e6; }
|
||||||
|
|
||||||
// Alloc helpers
|
// Alloc helpers
|
||||||
static AdamState adam_alloc(size_t n) { AdamState s; s.m=(float*)calloc(n,4); s.v=(float*)calloc(n,4); s.n=n; return s; }
|
// HIGH-04: OOM during training is fatal and unrecoverable; abort() is correct.
|
||||||
static void adam_free(AdamState *s) { free(s->m); free(s->v); }
|
static inline float *xmf(size_t n) {
|
||||||
|
float *p = (float*)malloc(n * sizeof(float));
|
||||||
static LayerWeights layer_weights_alloc(void) {
|
if (!p) { fprintf(stderr, "OOM: malloc(%zu floats = %.1fMB)\n", n, n*4.0/1048576); abort(); }
|
||||||
LayerWeights w;
|
return p;
|
||||||
w.Wq=(float*)malloc(WQ_SZ*4); w.Wk=(float*)malloc(WQ_SZ*4);
|
}
|
||||||
w.Wv=(float*)malloc(WQ_SZ*4); w.Wo=(float*)malloc(WO_SZ*4);
|
static inline float *xcf(size_t n) {
|
||||||
w.W1=(float*)malloc(W1_SZ*4); w.W2=(float*)malloc(W2_SZ*4); w.W3=(float*)malloc(W3_SZ*4);
|
float *p = (float*)calloc(n, sizeof(float));
|
||||||
w.rms_att=(float*)malloc(DIM*4); w.rms_ffn=(float*)malloc(DIM*4);
|
if (!p) { fprintf(stderr, "OOM: calloc(%zu floats = %.1fMB)\n", n, n*4.0/1048576); abort(); }
|
||||||
return w;
|
return p;
|
||||||
}
|
}
|
||||||
static void layer_weights_free(LayerWeights *w) {
|
static AdamState adam_alloc(size_t n) { AdamState s; s.m=xcf(n); s.v=xcf(n); s.n=n; return s; }
|
||||||
free(w->Wq);free(w->Wk);free(w->Wv);free(w->Wo);
|
static void adam_free(AdamState *s) { free(s->m); free(s->v); }
|
||||||
free(w->W1);free(w->W2);free(w->W3);
|
|
||||||
free(w->rms_att);free(w->rms_ffn);
|
static LayerWeights layer_weights_alloc(void) {
|
||||||
}
|
LayerWeights w;
|
||||||
static LayerAdam layer_adam_alloc(void) {
|
w.Wq=xmf(WQ_SZ); w.Wk=xmf(WQ_SZ);
|
||||||
LayerAdam a;
|
w.Wv=xmf(WQ_SZ); w.Wo=xmf(WO_SZ);
|
||||||
a.Wq=adam_alloc(WQ_SZ); a.Wk=adam_alloc(WQ_SZ); a.Wv=adam_alloc(WQ_SZ); a.Wo=adam_alloc(WO_SZ);
|
w.W1=xmf(W1_SZ); w.W2=xmf(W2_SZ); w.W3=xmf(W3_SZ);
|
||||||
a.W1=adam_alloc(W1_SZ); a.W2=adam_alloc(W2_SZ); a.W3=adam_alloc(W3_SZ);
|
w.rms_att=xmf(DIM); w.rms_ffn=xmf(DIM);
|
||||||
a.rms_att=adam_alloc(DIM); a.rms_ffn=adam_alloc(DIM);
|
return w;
|
||||||
return a;
|
}
|
||||||
}
|
static void layer_weights_free(LayerWeights *w) {
|
||||||
static void layer_adam_free(LayerAdam *a) {
|
free(w->Wq);free(w->Wk);free(w->Wv);free(w->Wo);
|
||||||
adam_free(&a->Wq);adam_free(&a->Wk);adam_free(&a->Wv);adam_free(&a->Wo);
|
free(w->W1);free(w->W2);free(w->W3);
|
||||||
adam_free(&a->W1);adam_free(&a->W2);adam_free(&a->W3);
|
free(w->rms_att);free(w->rms_ffn);
|
||||||
adam_free(&a->rms_att);adam_free(&a->rms_ffn);
|
}
|
||||||
}
|
static LayerAdam layer_adam_alloc(void) {
|
||||||
static LayerActs layer_acts_alloc(void) {
|
LayerAdam a;
|
||||||
LayerActs a;
|
a.Wq=adam_alloc(WQ_SZ); a.Wk=adam_alloc(WQ_SZ); a.Wv=adam_alloc(WQ_SZ); a.Wo=adam_alloc(WO_SZ);
|
||||||
a.layer_in=(float*)malloc(SEQ*DIM*4);
|
a.W1=adam_alloc(W1_SZ); a.W2=adam_alloc(W2_SZ); a.W3=adam_alloc(W3_SZ);
|
||||||
a.xnorm=(float*)malloc(SEQ*DIM*4); a.Q=(float*)malloc(SEQ*DIM*4);
|
a.rms_att=adam_alloc(DIM); a.rms_ffn=adam_alloc(DIM);
|
||||||
a.K=(float*)malloc(SEQ*DIM*4); a.V=(float*)malloc(SEQ*DIM*4);
|
return a;
|
||||||
a.attn_out=(float*)malloc(SEQ*DIM*4); a.o_out=(float*)malloc(SEQ*DIM*4);
|
}
|
||||||
a.x2=(float*)malloc(SEQ*DIM*4); a.x2norm=(float*)malloc(SEQ*DIM*4);
|
static void layer_adam_free(LayerAdam *a) {
|
||||||
a.h1=(float*)malloc(SEQ*HIDDEN*4); a.h3=(float*)malloc(SEQ*HIDDEN*4);
|
adam_free(&a->Wq);adam_free(&a->Wk);adam_free(&a->Wv);adam_free(&a->Wo);
|
||||||
a.silu_out=(float*)malloc(SEQ*HIDDEN*4); a.ffn_out=(float*)malloc(SEQ*DIM*4);
|
adam_free(&a->W1);adam_free(&a->W2);adam_free(&a->W3);
|
||||||
return a;
|
adam_free(&a->rms_att);adam_free(&a->rms_ffn);
|
||||||
}
|
}
|
||||||
static void layer_acts_free(LayerActs *a) {
|
static LayerActs layer_acts_alloc(void) {
|
||||||
free(a->layer_in);free(a->xnorm);free(a->Q);free(a->K);free(a->V);
|
LayerActs a;
|
||||||
free(a->attn_out);free(a->o_out);free(a->x2);free(a->x2norm);
|
a.layer_in=xmf((size_t)SEQ*DIM);
|
||||||
free(a->h1);free(a->h3);free(a->silu_out);free(a->ffn_out);
|
a.xnorm=xmf((size_t)SEQ*DIM); a.Q=xmf((size_t)SEQ*DIM);
|
||||||
}
|
a.K=xmf((size_t)SEQ*DIM); a.V=xmf((size_t)SEQ*DIM);
|
||||||
static LayerGrads layer_grads_alloc(void) {
|
a.attn_out=xmf((size_t)SEQ*DIM); a.o_out=xmf((size_t)SEQ*DIM);
|
||||||
LayerGrads g;
|
a.x2=xmf((size_t)SEQ*DIM); a.x2norm=xmf((size_t)SEQ*DIM);
|
||||||
g.Wq=(float*)calloc(WQ_SZ,4); g.Wk=(float*)calloc(WQ_SZ,4);
|
a.h1=xmf((size_t)SEQ*HIDDEN); a.h3=xmf((size_t)SEQ*HIDDEN);
|
||||||
g.Wv=(float*)calloc(WQ_SZ,4); g.Wo=(float*)calloc(WO_SZ,4);
|
a.silu_out=xmf((size_t)SEQ*HIDDEN); a.ffn_out=xmf((size_t)SEQ*DIM);
|
||||||
g.W1=(float*)calloc(W1_SZ,4); g.W2=(float*)calloc(W2_SZ,4); g.W3=(float*)calloc(W3_SZ,4);
|
return a;
|
||||||
g.rms_att=(float*)calloc(DIM,4); g.rms_ffn=(float*)calloc(DIM,4);
|
}
|
||||||
return g;
|
static void layer_acts_free(LayerActs *a) {
|
||||||
}
|
free(a->layer_in);free(a->xnorm);free(a->Q);free(a->K);free(a->V);
|
||||||
static void layer_grads_zero(LayerGrads *g) {
|
free(a->attn_out);free(a->o_out);free(a->x2);free(a->x2norm);
|
||||||
memset(g->Wq,0,WQ_SZ*4);memset(g->Wk,0,WQ_SZ*4);
|
free(a->h1);free(a->h3);free(a->silu_out);free(a->ffn_out);
|
||||||
memset(g->Wv,0,WQ_SZ*4);memset(g->Wo,0,WO_SZ*4);
|
}
|
||||||
memset(g->W1,0,W1_SZ*4);memset(g->W2,0,W2_SZ*4);memset(g->W3,0,W3_SZ*4);
|
static LayerGrads layer_grads_alloc(void) {
|
||||||
memset(g->rms_att,0,DIM*4);memset(g->rms_ffn,0,DIM*4);
|
LayerGrads g;
|
||||||
}
|
g.Wq=xcf(WQ_SZ); g.Wk=xcf(WQ_SZ);
|
||||||
static void layer_grads_free(LayerGrads *g) {
|
g.Wv=xcf(WQ_SZ); g.Wo=xcf(WO_SZ);
|
||||||
free(g->Wq);free(g->Wk);free(g->Wv);free(g->Wo);
|
g.W1=xcf(W1_SZ); g.W2=xcf(W2_SZ); g.W3=xcf(W3_SZ);
|
||||||
free(g->W1);free(g->W2);free(g->W3);
|
g.rms_att=xcf(DIM); g.rms_ffn=xcf(DIM);
|
||||||
free(g->rms_att);free(g->rms_ffn);
|
return g;
|
||||||
}
|
}
|
||||||
|
static void layer_grads_zero(LayerGrads *g) {
|
||||||
|
memset(g->Wq,0,WQ_SZ*4);memset(g->Wk,0,WQ_SZ*4);
|
||||||
|
memset(g->Wv,0,WQ_SZ*4);memset(g->Wo,0,WO_SZ*4);
|
||||||
|
memset(g->W1,0,W1_SZ*4);memset(g->W2,0,W2_SZ*4);memset(g->W3,0,W3_SZ*4);
|
||||||
|
memset(g->rms_att,0,DIM*4);memset(g->rms_ffn,0,DIM*4);
|
||||||
|
}
|
||||||
|
static void layer_grads_free(LayerGrads *g) {
|
||||||
|
free(g->Wq);free(g->Wk);free(g->Wv);free(g->Wo);
|
||||||
|
free(g->W1);free(g->W2);free(g->W3);
|
||||||
|
free(g->rms_att);free(g->rms_ffn);
|
||||||
|
}
|
||||||
|
|
|
||||||
|
|
@ -235,10 +235,10 @@ int main(int argc, char *argv[]) {
|
||||||
}
|
}
|
||||||
|
|
||||||
// Final RMSNorm + embedding + classifier
|
// Final RMSNorm + embedding + classifier
|
||||||
float *rms_final = (float*)malloc(DIM*4);
|
float *rms_final = xmf(DIM);
|
||||||
float *embed = (float*)malloc(VOCAB*DIM*4); // [VOCAB, DIM] row-major
|
float *embed = xmf((size_t)VOCAB*DIM); // [VOCAB, DIM] row-major
|
||||||
float *grms_final = (float*)calloc(DIM, 4);
|
float *grms_final = xcf(DIM);
|
||||||
float *gembed = (float*)calloc(VOCAB*DIM, 4);
|
float *gembed = xcf((size_t)VOCAB*DIM);
|
||||||
AdamState arms_final = adam_alloc(DIM);
|
AdamState arms_final = adam_alloc(DIM);
|
||||||
AdamState aembed = adam_alloc((size_t)VOCAB*DIM);
|
AdamState aembed = adam_alloc((size_t)VOCAB*DIM);
|
||||||
|
|
||||||
|
|
@ -316,23 +316,23 @@ int main(int argc, char *argv[]) {
|
||||||
printf("Token data: %zu tokens (%.1f MB)\n", n_tokens, data_len/1e6);
|
printf("Token data: %zu tokens (%.1f MB)\n", n_tokens, data_len/1e6);
|
||||||
|
|
||||||
// Gradient buffers shared across layers (reused each step)
|
// Gradient buffers shared across layers (reused each step)
|
||||||
float *dy = (float*)malloc(SEQ*DIM*4); // gradient flowing backward
|
float *dy = xmf((size_t)SEQ*DIM); // gradient flowing backward
|
||||||
float *dffn = (float*)malloc(SEQ*DIM*4);
|
float *dffn = xmf((size_t)SEQ*DIM);
|
||||||
float *dh1 = (float*)malloc(SEQ*HIDDEN*4);
|
float *dh1 = xmf((size_t)SEQ*HIDDEN);
|
||||||
float *dh3 = (float*)malloc(SEQ*HIDDEN*4);
|
float *dh3 = xmf((size_t)SEQ*HIDDEN);
|
||||||
float *dx_ffn = (float*)malloc(SEQ*DIM*4);
|
float *dx_ffn = xmf((size_t)SEQ*DIM);
|
||||||
float *dx2 = (float*)malloc(SEQ*DIM*4);
|
float *dx2 = xmf((size_t)SEQ*DIM);
|
||||||
float *do_out_buf = (float*)malloc(SEQ*DIM*4);
|
float *do_out_buf = xmf((size_t)SEQ*DIM);
|
||||||
float *dq = (float*)malloc(SEQ*DIM*4);
|
float *dq = xmf((size_t)SEQ*DIM);
|
||||||
float *dk = (float*)malloc(SEQ*DIM*4);
|
float *dk = xmf((size_t)SEQ*DIM);
|
||||||
float *dv = (float*)malloc(SEQ*DIM*4);
|
float *dv = xmf((size_t)SEQ*DIM);
|
||||||
float *dx_attn = (float*)malloc(SEQ*DIM*4);
|
float *dx_attn = xmf((size_t)SEQ*DIM);
|
||||||
|
|
||||||
// x buffer for input to each layer (channel-first [DIM, SEQ])
|
// x buffer for input to each layer (channel-first [DIM, SEQ])
|
||||||
float *x_cur = (float*)malloc(SEQ*DIM*4);
|
float *x_cur = xmf((size_t)SEQ*DIM);
|
||||||
float *x_final = (float*)malloc(SEQ*DIM*4); // after final rmsnorm
|
float *x_final = xmf((size_t)SEQ*DIM); // after final rmsnorm
|
||||||
float *logits = (float*)malloc(SEQ*VOCAB*4); // [VOCAB, SEQ] for cross-entropy
|
float *logits = xmf((size_t)SEQ*VOCAB); // [VOCAB, SEQ] for cross-entropy
|
||||||
float *dlogits = (float*)malloc(SEQ*VOCAB*4);
|
float *dlogits = xmf((size_t)SEQ*VOCAB);
|
||||||
|
|
||||||
// Compile static sdpaBwd2 kernels (no weights, one per layer)
|
// Compile static sdpaBwd2 kernels (no weights, one per layer)
|
||||||
Kern *sdpaBwd2[NLAYERS];
|
Kern *sdpaBwd2[NLAYERS];
|
||||||
|
|
@ -492,7 +492,7 @@ int main(int argc, char *argv[]) {
|
||||||
});
|
});
|
||||||
|
|
||||||
// Final RMSNorm backward
|
// Final RMSNorm backward
|
||||||
float *dx_rms_final = (float*)calloc(SEQ*DIM, 4);
|
float *dx_rms_final = xcf((size_t)SEQ*DIM);
|
||||||
rmsnorm_bwd(dx_rms_final, grms_final, dy, x_cur, rms_final, DIM, SEQ);
|
rmsnorm_bwd(dx_rms_final, grms_final, dy, x_cur, rms_final, DIM, SEQ);
|
||||||
memcpy(dy, dx_rms_final, SEQ*DIM*4);
|
memcpy(dy, dx_rms_final, SEQ*DIM*4);
|
||||||
free(dx_rms_final);
|
free(dx_rms_final);
|
||||||
|
|
@ -515,11 +515,11 @@ int main(int argc, char *argv[]) {
|
||||||
io_read_fp16(kern[L].ffnBwd->ioOut, dh3, DIM+HIDDEN, HIDDEN, SEQ);
|
io_read_fp16(kern[L].ffnBwd->ioOut, dh3, DIM+HIDDEN, HIDDEN, SEQ);
|
||||||
|
|
||||||
// dW FFN async
|
// dW FFN async
|
||||||
float *capt_dffn = (float*)malloc(SEQ*DIM*4); memcpy(capt_dffn, dffn, SEQ*DIM*4);
|
float *capt_dffn = xmf((size_t)SEQ*DIM); memcpy(capt_dffn, dffn, SEQ*DIM*4);
|
||||||
float *capt_silu = (float*)malloc(SEQ*HIDDEN*4); memcpy(capt_silu, ac->silu_out, SEQ*HIDDEN*4);
|
float *capt_silu = xmf((size_t)SEQ*HIDDEN); memcpy(capt_silu, ac->silu_out, SEQ*HIDDEN*4);
|
||||||
float *capt_dh1 = (float*)malloc(SEQ*HIDDEN*4); memcpy(capt_dh1, dh1, SEQ*HIDDEN*4);
|
float *capt_dh1 = xmf((size_t)SEQ*HIDDEN); memcpy(capt_dh1, dh1, SEQ*HIDDEN*4);
|
||||||
float *capt_dh3 = (float*)malloc(SEQ*HIDDEN*4); memcpy(capt_dh3, dh3, SEQ*HIDDEN*4);
|
float *capt_dh3 = xmf((size_t)SEQ*HIDDEN); memcpy(capt_dh3, dh3, SEQ*HIDDEN*4);
|
||||||
float *capt_x2n = (float*)malloc(SEQ*DIM*4); memcpy(capt_x2n, ac->x2norm, SEQ*DIM*4);
|
float *capt_x2n = xmf((size_t)SEQ*DIM); memcpy(capt_x2n, ac->x2norm, SEQ*DIM*4);
|
||||||
dispatch_group_async(dw_grp, dw_q, ^{
|
dispatch_group_async(dw_grp, dw_q, ^{
|
||||||
cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, DIM, HIDDEN, SEQ,
|
cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, DIM, HIDDEN, SEQ,
|
||||||
1.0f, capt_dffn, SEQ, capt_silu, SEQ, 1.0f, gr->W2, HIDDEN);
|
1.0f, capt_dffn, SEQ, capt_silu, SEQ, 1.0f, gr->W2, HIDDEN);
|
||||||
|
|
@ -538,8 +538,8 @@ int main(int argc, char *argv[]) {
|
||||||
|
|
||||||
// dWo async
|
// dWo async
|
||||||
memcpy(do_out_buf, dx2, SEQ*DIM*4);
|
memcpy(do_out_buf, dx2, SEQ*DIM*4);
|
||||||
float *capt_do = (float*)malloc(SEQ*DIM*4); memcpy(capt_do, do_out_buf, SEQ*DIM*4);
|
float *capt_do = xmf((size_t)SEQ*DIM); memcpy(capt_do, do_out_buf, SEQ*DIM*4);
|
||||||
float *capt_attn = (float*)malloc(SEQ*DIM*4); memcpy(capt_attn, ac->attn_out, SEQ*DIM*4);
|
float *capt_attn = xmf((size_t)SEQ*DIM); memcpy(capt_attn, ac->attn_out, SEQ*DIM*4);
|
||||||
dispatch_group_async(dw_grp, dw_q, ^{
|
dispatch_group_async(dw_grp, dw_q, ^{
|
||||||
cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, DIM, DIM, SEQ,
|
cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, DIM, DIM, SEQ,
|
||||||
1.0f, capt_do, SEQ, capt_attn, SEQ, 1.0f, gr->Wo, DIM);
|
1.0f, capt_do, SEQ, capt_attn, SEQ, 1.0f, gr->Wo, DIM);
|
||||||
|
|
@ -559,10 +559,10 @@ int main(int argc, char *argv[]) {
|
||||||
io_read_fp16(kern[L].sdpaBwd1->ioOut, dv, 0, DIM, SEQ);
|
io_read_fp16(kern[L].sdpaBwd1->ioOut, dv, 0, DIM, SEQ);
|
||||||
|
|
||||||
// dWq/dWk/dWv async
|
// dWq/dWk/dWv async
|
||||||
float *capt_dq = (float*)malloc(SEQ*DIM*4); memcpy(capt_dq, dq, SEQ*DIM*4);
|
float *capt_dq = xmf((size_t)SEQ*DIM); memcpy(capt_dq, dq, SEQ*DIM*4);
|
||||||
float *capt_dk = (float*)malloc(SEQ*DIM*4); memcpy(capt_dk, dk, SEQ*DIM*4);
|
float *capt_dk = xmf((size_t)SEQ*DIM); memcpy(capt_dk, dk, SEQ*DIM*4);
|
||||||
float *capt_dv = (float*)malloc(SEQ*DIM*4); memcpy(capt_dv, dv, SEQ*DIM*4);
|
float *capt_dv = xmf((size_t)SEQ*DIM); memcpy(capt_dv, dv, SEQ*DIM*4);
|
||||||
float *capt_xn = (float*)malloc(SEQ*DIM*4); memcpy(capt_xn, ac->xnorm, SEQ*DIM*4);
|
float *capt_xn = xmf((size_t)SEQ*DIM); memcpy(capt_xn, ac->xnorm, SEQ*DIM*4);
|
||||||
dispatch_group_async(dw_grp, dw_q, ^{
|
dispatch_group_async(dw_grp, dw_q, ^{
|
||||||
cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, DIM, DIM, SEQ,
|
cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, DIM, DIM, SEQ,
|
||||||
1.0f, capt_dq, SEQ, capt_xn, SEQ, 1.0f, gr->Wq, DIM);
|
1.0f, capt_dq, SEQ, capt_xn, SEQ, 1.0f, gr->Wq, DIM);
|
||||||
|
|
@ -580,7 +580,7 @@ int main(int argc, char *argv[]) {
|
||||||
io_read_fp16(kern[L].qkvBwd->ioOut, dx_attn, 0, DIM, SEQ);
|
io_read_fp16(kern[L].qkvBwd->ioOut, dx_attn, 0, DIM, SEQ);
|
||||||
|
|
||||||
// RMSNorm1 backward (using saved layer input)
|
// RMSNorm1 backward (using saved layer input)
|
||||||
float *dx_rms1 = (float*)calloc(SEQ*DIM, 4);
|
float *dx_rms1 = xcf((size_t)SEQ*DIM);
|
||||||
rmsnorm_bwd(dx_rms1, gr->rms_att, dx_attn, ac->layer_in, lw[L].rms_att, DIM, SEQ);
|
rmsnorm_bwd(dx_rms1, gr->rms_att, dx_attn, ac->layer_in, lw[L].rms_att, DIM, SEQ);
|
||||||
|
|
||||||
// dy for next layer (going backward) = dx_rms1 + dx2 residual
|
// dy for next layer (going backward) = dx_rms1 + dx2 residual
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue