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
|
||||
#pragma once
|
||||
#import <Foundation/Foundation.h>
|
||||
#import <objc/runtime.h>
|
||||
#import <objc/message.h>
|
||||
#import <dlfcn.h>
|
||||
#import <IOSurface/IOSurface.h>
|
||||
#import <mach/mach_time.h>
|
||||
#import <Accelerate/Accelerate.h>
|
||||
#include <math.h>
|
||||
#include <unistd.h>
|
||||
#include <dispatch/dispatch.h>
|
||||
#include <sys/mman.h>
|
||||
#include <sys/stat.h>
|
||||
#include <fcntl.h>
|
||||
|
||||
// Stories110M config
|
||||
#define DIM 768
|
||||
#define HIDDEN 2048
|
||||
#define HEADS 12
|
||||
#define HD (DIM/HEADS)
|
||||
#define SEQ 256
|
||||
#define NLAYERS 12
|
||||
#define VOCAB 32000
|
||||
#define ACCUM_STEPS 10
|
||||
#define MAX_COMPILES 100
|
||||
|
||||
// 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
|
||||
// Actually sdpaBwd2 has no weights, compile once 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
|
||||
// With MAX_COMPILES=100, we get 1 batch of ACCUM_STEPS before restart
|
||||
#define KERNELS_PER_LAYER 5
|
||||
#define TOTAL_WEIGHT_KERNELS (KERNELS_PER_LAYER * NLAYERS)
|
||||
|
||||
// Attention score channels for SDPA backward
|
||||
#define SCORE_CH (HEADS*SEQ)
|
||||
|
||||
// Weight sizes per layer
|
||||
#define WQ_SZ (DIM*DIM)
|
||||
#define WO_SZ (DIM*DIM)
|
||||
#define W1_SZ (HIDDEN*DIM)
|
||||
#define W2_SZ (DIM*HIDDEN)
|
||||
#define W3_SZ (HIDDEN*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
|
||||
|
||||
// Per-layer weight and optimizer state
|
||||
typedef struct {
|
||||
float *Wq, *Wk, *Wv, *Wo;
|
||||
float *W1, *W2, *W3;
|
||||
float *rms_att, *rms_ffn;
|
||||
} LayerWeights;
|
||||
|
||||
typedef struct {
|
||||
float *m, *v;
|
||||
size_t n;
|
||||
} AdamState;
|
||||
|
||||
typedef struct {
|
||||
AdamState Wq, Wk, Wv, Wo;
|
||||
AdamState W1, W2, W3;
|
||||
AdamState rms_att, rms_ffn;
|
||||
} LayerAdam;
|
||||
|
||||
// Per-layer activation buffers (saved for backward)
|
||||
typedef struct {
|
||||
float *layer_in; // [DIM, SEQ] input to this layer (for rmsnorm1 bwd)
|
||||
float *xnorm; // [DIM, SEQ] rmsnorm1 output
|
||||
float *Q, *K, *V; // [DIM, SEQ] QKV projections
|
||||
float *attn_out; // [DIM, SEQ] attention output (before Wo)
|
||||
float *o_out; // [DIM, SEQ] Wo output
|
||||
float *x2; // [DIM, SEQ] residual after attn
|
||||
float *x2norm; // [DIM, SEQ] rmsnorm2 output
|
||||
float *h1, *h3; // [HIDDEN, SEQ] FFN intermediates
|
||||
float *silu_out; // [HIDDEN, SEQ] SiLU(h1)*h3
|
||||
float *ffn_out; // [DIM, SEQ] FFN output
|
||||
} LayerActs;
|
||||
|
||||
// Per-layer gradient accumulators
|
||||
typedef struct {
|
||||
float *Wq, *Wk, *Wv, *Wo;
|
||||
float *W1, *W2, *W3;
|
||||
float *rms_att, *rms_ffn;
|
||||
} LayerGrads;
|
||||
|
||||
// ANE kernels per layer
|
||||
typedef struct { void *model; IOSurfaceRef ioIn, ioOut; void *request; void *tmpDir; } Kern;
|
||||
typedef struct {
|
||||
Kern *fwdAttn, *fwdFFN, *ffnBwd, *sdpaBwd1, *sdpaBwd2, *qkvBwd;
|
||||
} LayerKernels;
|
||||
|
||||
// Checkpoint header
|
||||
typedef struct {
|
||||
int magic; // 0x424C5A54 "BLZT"
|
||||
int version; // 2
|
||||
int step, total_steps;
|
||||
int n_layers, vocab_size, dim, hidden_dim, n_heads, seq_len;
|
||||
float lr, loss;
|
||||
double cum_compile, cum_train, cum_wall;
|
||||
int cum_steps, cum_batches;
|
||||
int adam_t;
|
||||
int pad[3]; // pad[0] = 0x01020304 (LE byte-order sentinel, MED-04); pad[1..2] = 0
|
||||
} CkptHdr;
|
||||
|
||||
// llama2.c model file header
|
||||
typedef struct {
|
||||
int dim, hidden_dim, n_layers, n_heads, n_kv_heads, vocab_size, seq_len;
|
||||
} Llama2Config;
|
||||
|
||||
// Globals
|
||||
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 mach_timebase_info_data_t g_tb;
|
||||
static int g_compile_count = 0;
|
||||
static int g_compile_seq = 0; // MED-02: per-call unique index for temp-dir naming
|
||||
|
||||
static void ane_init(void) {
|
||||
// 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.
|
||||
static dispatch_once_t ane_once_large;
|
||||
dispatch_once(&ane_once_large, ^{
|
||||
void *handle = dlopen(
|
||||
"/System/Library/PrivateFrameworks/AppleNeuralEngine.framework/AppleNeuralEngine",
|
||||
RTLD_NOW);
|
||||
if (!handle) {
|
||||
fprintf(stderr, "ANE: dlopen failed: %s\n", dlerror());
|
||||
return;
|
||||
}
|
||||
g_D = NSClassFromString(@"_ANEInMemoryModelDescriptor");
|
||||
g_I = NSClassFromString(@"_ANEInMemoryModel");
|
||||
g_AR = NSClassFromString(@"_ANERequest");
|
||||
g_AIO= NSClassFromString(@"_ANEIOSurfaceObject");
|
||||
if (!g_D || !g_I || !g_AR || !g_AIO) {
|
||||
fprintf(stderr, "ANE: Private classes not found (macOS version mismatch?)\n");
|
||||
return;
|
||||
}
|
||||
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; }
|
||||
|
||||
// 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; }
|
||||
static void adam_free(AdamState *s) { free(s->m); free(s->v); }
|
||||
|
||||
static LayerWeights layer_weights_alloc(void) {
|
||||
LayerWeights w;
|
||||
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);
|
||||
w.W1=(float*)malloc(W1_SZ*4); w.W2=(float*)malloc(W2_SZ*4); w.W3=(float*)malloc(W3_SZ*4);
|
||||
w.rms_att=(float*)malloc(DIM*4); w.rms_ffn=(float*)malloc(DIM*4);
|
||||
return w;
|
||||
}
|
||||
static void layer_weights_free(LayerWeights *w) {
|
||||
free(w->Wq);free(w->Wk);free(w->Wv);free(w->Wo);
|
||||
free(w->W1);free(w->W2);free(w->W3);
|
||||
free(w->rms_att);free(w->rms_ffn);
|
||||
}
|
||||
static LayerAdam layer_adam_alloc(void) {
|
||||
LayerAdam 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.W1=adam_alloc(W1_SZ); a.W2=adam_alloc(W2_SZ); a.W3=adam_alloc(W3_SZ);
|
||||
a.rms_att=adam_alloc(DIM); a.rms_ffn=adam_alloc(DIM);
|
||||
return a;
|
||||
}
|
||||
static void layer_adam_free(LayerAdam *a) {
|
||||
adam_free(&a->Wq);adam_free(&a->Wk);adam_free(&a->Wv);adam_free(&a->Wo);
|
||||
adam_free(&a->W1);adam_free(&a->W2);adam_free(&a->W3);
|
||||
adam_free(&a->rms_att);adam_free(&a->rms_ffn);
|
||||
}
|
||||
static LayerActs layer_acts_alloc(void) {
|
||||
LayerActs a;
|
||||
a.layer_in=(float*)malloc(SEQ*DIM*4);
|
||||
a.xnorm=(float*)malloc(SEQ*DIM*4); a.Q=(float*)malloc(SEQ*DIM*4);
|
||||
a.K=(float*)malloc(SEQ*DIM*4); a.V=(float*)malloc(SEQ*DIM*4);
|
||||
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);
|
||||
a.h1=(float*)malloc(SEQ*HIDDEN*4); a.h3=(float*)malloc(SEQ*HIDDEN*4);
|
||||
a.silu_out=(float*)malloc(SEQ*HIDDEN*4); a.ffn_out=(float*)malloc(SEQ*DIM*4);
|
||||
return a;
|
||||
}
|
||||
static void layer_acts_free(LayerActs *a) {
|
||||
free(a->layer_in);free(a->xnorm);free(a->Q);free(a->K);free(a->V);
|
||||
free(a->attn_out);free(a->o_out);free(a->x2);free(a->x2norm);
|
||||
free(a->h1);free(a->h3);free(a->silu_out);free(a->ffn_out);
|
||||
}
|
||||
static LayerGrads layer_grads_alloc(void) {
|
||||
LayerGrads g;
|
||||
g.Wq=(float*)calloc(WQ_SZ,4); g.Wk=(float*)calloc(WQ_SZ,4);
|
||||
g.Wv=(float*)calloc(WQ_SZ,4); g.Wo=(float*)calloc(WO_SZ,4);
|
||||
g.W1=(float*)calloc(W1_SZ,4); g.W2=(float*)calloc(W2_SZ,4); g.W3=(float*)calloc(W3_SZ,4);
|
||||
g.rms_att=(float*)calloc(DIM,4); g.rms_ffn=(float*)calloc(DIM,4);
|
||||
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);
|
||||
}
|
||||
// stories_config.h — Stories110M model config and structures
|
||||
#pragma once
|
||||
#import <Foundation/Foundation.h>
|
||||
#import <objc/runtime.h>
|
||||
#import <objc/message.h>
|
||||
#import <dlfcn.h>
|
||||
#import <IOSurface/IOSurface.h>
|
||||
#import <mach/mach_time.h>
|
||||
#import <Accelerate/Accelerate.h>
|
||||
#include <math.h>
|
||||
#include <unistd.h>
|
||||
#include <dispatch/dispatch.h>
|
||||
#include <sys/mman.h>
|
||||
#include <sys/stat.h>
|
||||
#include <fcntl.h>
|
||||
|
||||
// Stories110M config
|
||||
#define DIM 768
|
||||
#define HIDDEN 2048
|
||||
#define HEADS 12
|
||||
#define HD (DIM/HEADS)
|
||||
#define SEQ 256
|
||||
#define NLAYERS 12
|
||||
#define VOCAB 32000
|
||||
#define ACCUM_STEPS 10
|
||||
#define MAX_COMPILES 100
|
||||
|
||||
// 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
|
||||
// Actually sdpaBwd2 has no weights, compile once 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
|
||||
// With MAX_COMPILES=100, we get 1 batch of ACCUM_STEPS before restart
|
||||
#define KERNELS_PER_LAYER 5
|
||||
#define TOTAL_WEIGHT_KERNELS (KERNELS_PER_LAYER * NLAYERS)
|
||||
|
||||
// Attention score channels for SDPA backward
|
||||
#define SCORE_CH (HEADS*SEQ)
|
||||
|
||||
// Weight sizes per layer
|
||||
#define WQ_SZ (DIM*DIM)
|
||||
#define WO_SZ (DIM*DIM)
|
||||
#define W1_SZ (HIDDEN*DIM)
|
||||
#define W2_SZ (DIM*HIDDEN)
|
||||
#define W3_SZ (HIDDEN*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
|
||||
|
||||
// Per-layer weight and optimizer state
|
||||
typedef struct {
|
||||
float *Wq, *Wk, *Wv, *Wo;
|
||||
float *W1, *W2, *W3;
|
||||
float *rms_att, *rms_ffn;
|
||||
} LayerWeights;
|
||||
|
||||
typedef struct {
|
||||
float *m, *v;
|
||||
size_t n;
|
||||
} AdamState;
|
||||
|
||||
typedef struct {
|
||||
AdamState Wq, Wk, Wv, Wo;
|
||||
AdamState W1, W2, W3;
|
||||
AdamState rms_att, rms_ffn;
|
||||
} LayerAdam;
|
||||
|
||||
// Per-layer activation buffers (saved for backward)
|
||||
typedef struct {
|
||||
float *layer_in; // [DIM, SEQ] input to this layer (for rmsnorm1 bwd)
|
||||
float *xnorm; // [DIM, SEQ] rmsnorm1 output
|
||||
float *Q, *K, *V; // [DIM, SEQ] QKV projections
|
||||
float *attn_out; // [DIM, SEQ] attention output (before Wo)
|
||||
float *o_out; // [DIM, SEQ] Wo output
|
||||
float *x2; // [DIM, SEQ] residual after attn
|
||||
float *x2norm; // [DIM, SEQ] rmsnorm2 output
|
||||
float *h1, *h3; // [HIDDEN, SEQ] FFN intermediates
|
||||
float *silu_out; // [HIDDEN, SEQ] SiLU(h1)*h3
|
||||
float *ffn_out; // [DIM, SEQ] FFN output
|
||||
} LayerActs;
|
||||
|
||||
// Per-layer gradient accumulators
|
||||
typedef struct {
|
||||
float *Wq, *Wk, *Wv, *Wo;
|
||||
float *W1, *W2, *W3;
|
||||
float *rms_att, *rms_ffn;
|
||||
} LayerGrads;
|
||||
|
||||
// ANE kernels per layer
|
||||
typedef struct { void *model; IOSurfaceRef ioIn, ioOut; void *request; void *tmpDir; } Kern;
|
||||
typedef struct {
|
||||
Kern *fwdAttn, *fwdFFN, *ffnBwd, *sdpaBwd1, *sdpaBwd2, *qkvBwd;
|
||||
} LayerKernels;
|
||||
|
||||
// Checkpoint header
|
||||
typedef struct {
|
||||
int magic; // 0x424C5A54 "BLZT"
|
||||
int version; // 2
|
||||
int step, total_steps;
|
||||
int n_layers, vocab_size, dim, hidden_dim, n_heads, seq_len;
|
||||
float lr, loss;
|
||||
double cum_compile, cum_train, cum_wall;
|
||||
int cum_steps, cum_batches;
|
||||
int adam_t;
|
||||
int pad[3]; // pad[0] = 0x01020304 (LE byte-order sentinel, MED-04); pad[1..2] = 0
|
||||
} CkptHdr;
|
||||
|
||||
// llama2.c model file header
|
||||
typedef struct {
|
||||
int dim, hidden_dim, n_layers, n_heads, n_kv_heads, vocab_size, seq_len;
|
||||
} Llama2Config;
|
||||
|
||||
// Globals
|
||||
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 mach_timebase_info_data_t g_tb;
|
||||
static int g_compile_count = 0;
|
||||
static int g_compile_seq = 0; // MED-02: per-call unique index for temp-dir naming
|
||||
|
||||
static void ane_init(void) {
|
||||
// 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.
|
||||
static dispatch_once_t ane_once_large;
|
||||
dispatch_once(&ane_once_large, ^{
|
||||
void *handle = dlopen(
|
||||
"/System/Library/PrivateFrameworks/AppleNeuralEngine.framework/AppleNeuralEngine",
|
||||
RTLD_NOW);
|
||||
if (!handle) {
|
||||
fprintf(stderr, "ANE: dlopen failed: %s\n", dlerror());
|
||||
return;
|
||||
}
|
||||
g_D = NSClassFromString(@"_ANEInMemoryModelDescriptor");
|
||||
g_I = NSClassFromString(@"_ANEInMemoryModel");
|
||||
g_AR = NSClassFromString(@"_ANERequest");
|
||||
g_AIO= NSClassFromString(@"_ANEIOSurfaceObject");
|
||||
if (!g_D || !g_I || !g_AR || !g_AIO) {
|
||||
fprintf(stderr, "ANE: Private classes not found (macOS version mismatch?)\n");
|
||||
return;
|
||||
}
|
||||
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; }
|
||||
|
||||
// Alloc helpers
|
||||
// HIGH-04: OOM during training is fatal and unrecoverable; abort() is correct.
|
||||
static inline float *xmf(size_t n) {
|
||||
float *p = (float*)malloc(n * sizeof(float));
|
||||
if (!p) { fprintf(stderr, "OOM: malloc(%zu floats = %.1fMB)\n", n, n*4.0/1048576); abort(); }
|
||||
return p;
|
||||
}
|
||||
static inline float *xcf(size_t n) {
|
||||
float *p = (float*)calloc(n, sizeof(float));
|
||||
if (!p) { fprintf(stderr, "OOM: calloc(%zu floats = %.1fMB)\n", n, n*4.0/1048576); abort(); }
|
||||
return p;
|
||||
}
|
||||
static AdamState adam_alloc(size_t n) { AdamState s; s.m=xcf(n); s.v=xcf(n); s.n=n; return s; }
|
||||
static void adam_free(AdamState *s) { free(s->m); free(s->v); }
|
||||
|
||||
static LayerWeights layer_weights_alloc(void) {
|
||||
LayerWeights w;
|
||||
w.Wq=xmf(WQ_SZ); w.Wk=xmf(WQ_SZ);
|
||||
w.Wv=xmf(WQ_SZ); w.Wo=xmf(WO_SZ);
|
||||
w.W1=xmf(W1_SZ); w.W2=xmf(W2_SZ); w.W3=xmf(W3_SZ);
|
||||
w.rms_att=xmf(DIM); w.rms_ffn=xmf(DIM);
|
||||
return w;
|
||||
}
|
||||
static void layer_weights_free(LayerWeights *w) {
|
||||
free(w->Wq);free(w->Wk);free(w->Wv);free(w->Wo);
|
||||
free(w->W1);free(w->W2);free(w->W3);
|
||||
free(w->rms_att);free(w->rms_ffn);
|
||||
}
|
||||
static LayerAdam layer_adam_alloc(void) {
|
||||
LayerAdam 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.W1=adam_alloc(W1_SZ); a.W2=adam_alloc(W2_SZ); a.W3=adam_alloc(W3_SZ);
|
||||
a.rms_att=adam_alloc(DIM); a.rms_ffn=adam_alloc(DIM);
|
||||
return a;
|
||||
}
|
||||
static void layer_adam_free(LayerAdam *a) {
|
||||
adam_free(&a->Wq);adam_free(&a->Wk);adam_free(&a->Wv);adam_free(&a->Wo);
|
||||
adam_free(&a->W1);adam_free(&a->W2);adam_free(&a->W3);
|
||||
adam_free(&a->rms_att);adam_free(&a->rms_ffn);
|
||||
}
|
||||
static LayerActs layer_acts_alloc(void) {
|
||||
LayerActs a;
|
||||
a.layer_in=xmf((size_t)SEQ*DIM);
|
||||
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);
|
||||
a.attn_out=xmf((size_t)SEQ*DIM); a.o_out=xmf((size_t)SEQ*DIM);
|
||||
a.x2=xmf((size_t)SEQ*DIM); a.x2norm=xmf((size_t)SEQ*DIM);
|
||||
a.h1=xmf((size_t)SEQ*HIDDEN); a.h3=xmf((size_t)SEQ*HIDDEN);
|
||||
a.silu_out=xmf((size_t)SEQ*HIDDEN); a.ffn_out=xmf((size_t)SEQ*DIM);
|
||||
return a;
|
||||
}
|
||||
static void layer_acts_free(LayerActs *a) {
|
||||
free(a->layer_in);free(a->xnorm);free(a->Q);free(a->K);free(a->V);
|
||||
free(a->attn_out);free(a->o_out);free(a->x2);free(a->x2norm);
|
||||
free(a->h1);free(a->h3);free(a->silu_out);free(a->ffn_out);
|
||||
}
|
||||
static LayerGrads layer_grads_alloc(void) {
|
||||
LayerGrads g;
|
||||
g.Wq=xcf(WQ_SZ); g.Wk=xcf(WQ_SZ);
|
||||
g.Wv=xcf(WQ_SZ); g.Wo=xcf(WO_SZ);
|
||||
g.W1=xcf(W1_SZ); g.W2=xcf(W2_SZ); g.W3=xcf(W3_SZ);
|
||||
g.rms_att=xcf(DIM); g.rms_ffn=xcf(DIM);
|
||||
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
|
||||
float *rms_final = (float*)malloc(DIM*4);
|
||||
float *embed = (float*)malloc(VOCAB*DIM*4); // [VOCAB, DIM] row-major
|
||||
float *grms_final = (float*)calloc(DIM, 4);
|
||||
float *gembed = (float*)calloc(VOCAB*DIM, 4);
|
||||
float *rms_final = xmf(DIM);
|
||||
float *embed = xmf((size_t)VOCAB*DIM); // [VOCAB, DIM] row-major
|
||||
float *grms_final = xcf(DIM);
|
||||
float *gembed = xcf((size_t)VOCAB*DIM);
|
||||
AdamState arms_final = adam_alloc(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);
|
||||
|
||||
// Gradient buffers shared across layers (reused each step)
|
||||
float *dy = (float*)malloc(SEQ*DIM*4); // gradient flowing backward
|
||||
float *dffn = (float*)malloc(SEQ*DIM*4);
|
||||
float *dh1 = (float*)malloc(SEQ*HIDDEN*4);
|
||||
float *dh3 = (float*)malloc(SEQ*HIDDEN*4);
|
||||
float *dx_ffn = (float*)malloc(SEQ*DIM*4);
|
||||
float *dx2 = (float*)malloc(SEQ*DIM*4);
|
||||
float *do_out_buf = (float*)malloc(SEQ*DIM*4);
|
||||
float *dq = (float*)malloc(SEQ*DIM*4);
|
||||
float *dk = (float*)malloc(SEQ*DIM*4);
|
||||
float *dv = (float*)malloc(SEQ*DIM*4);
|
||||
float *dx_attn = (float*)malloc(SEQ*DIM*4);
|
||||
float *dy = xmf((size_t)SEQ*DIM); // gradient flowing backward
|
||||
float *dffn = xmf((size_t)SEQ*DIM);
|
||||
float *dh1 = xmf((size_t)SEQ*HIDDEN);
|
||||
float *dh3 = xmf((size_t)SEQ*HIDDEN);
|
||||
float *dx_ffn = xmf((size_t)SEQ*DIM);
|
||||
float *dx2 = xmf((size_t)SEQ*DIM);
|
||||
float *do_out_buf = xmf((size_t)SEQ*DIM);
|
||||
float *dq = xmf((size_t)SEQ*DIM);
|
||||
float *dk = xmf((size_t)SEQ*DIM);
|
||||
float *dv = xmf((size_t)SEQ*DIM);
|
||||
float *dx_attn = xmf((size_t)SEQ*DIM);
|
||||
|
||||
// x buffer for input to each layer (channel-first [DIM, SEQ])
|
||||
float *x_cur = (float*)malloc(SEQ*DIM*4);
|
||||
float *x_final = (float*)malloc(SEQ*DIM*4); // after final rmsnorm
|
||||
float *logits = (float*)malloc(SEQ*VOCAB*4); // [VOCAB, SEQ] for cross-entropy
|
||||
float *dlogits = (float*)malloc(SEQ*VOCAB*4);
|
||||
float *x_cur = xmf((size_t)SEQ*DIM);
|
||||
float *x_final = xmf((size_t)SEQ*DIM); // after final rmsnorm
|
||||
float *logits = xmf((size_t)SEQ*VOCAB); // [VOCAB, SEQ] for cross-entropy
|
||||
float *dlogits = xmf((size_t)SEQ*VOCAB);
|
||||
|
||||
// Compile static sdpaBwd2 kernels (no weights, one per layer)
|
||||
Kern *sdpaBwd2[NLAYERS];
|
||||
|
|
@ -492,7 +492,7 @@ int main(int argc, char *argv[]) {
|
|||
});
|
||||
|
||||
// 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);
|
||||
memcpy(dy, dx_rms_final, SEQ*DIM*4);
|
||||
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);
|
||||
|
||||
// dW FFN async
|
||||
float *capt_dffn = (float*)malloc(SEQ*DIM*4); 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_dh1 = (float*)malloc(SEQ*HIDDEN*4); memcpy(capt_dh1, dh1, SEQ*HIDDEN*4);
|
||||
float *capt_dh3 = (float*)malloc(SEQ*HIDDEN*4); 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_dffn = xmf((size_t)SEQ*DIM); memcpy(capt_dffn, dffn, SEQ*DIM*4);
|
||||
float *capt_silu = xmf((size_t)SEQ*HIDDEN); memcpy(capt_silu, ac->silu_out, SEQ*HIDDEN*4);
|
||||
float *capt_dh1 = xmf((size_t)SEQ*HIDDEN); memcpy(capt_dh1, dh1, SEQ*HIDDEN*4);
|
||||
float *capt_dh3 = xmf((size_t)SEQ*HIDDEN); memcpy(capt_dh3, dh3, SEQ*HIDDEN*4);
|
||||
float *capt_x2n = xmf((size_t)SEQ*DIM); memcpy(capt_x2n, ac->x2norm, SEQ*DIM*4);
|
||||
dispatch_group_async(dw_grp, dw_q, ^{
|
||||
cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, DIM, HIDDEN, SEQ,
|
||||
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
|
||||
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_attn = (float*)malloc(SEQ*DIM*4); memcpy(capt_attn, ac->attn_out, SEQ*DIM*4);
|
||||
float *capt_do = xmf((size_t)SEQ*DIM); memcpy(capt_do, do_out_buf, 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, ^{
|
||||
cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, DIM, DIM, SEQ,
|
||||
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);
|
||||
|
||||
// dWq/dWk/dWv async
|
||||
float *capt_dq = (float*)malloc(SEQ*DIM*4); memcpy(capt_dq, dq, SEQ*DIM*4);
|
||||
float *capt_dk = (float*)malloc(SEQ*DIM*4); memcpy(capt_dk, dk, SEQ*DIM*4);
|
||||
float *capt_dv = (float*)malloc(SEQ*DIM*4); 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_dq = xmf((size_t)SEQ*DIM); memcpy(capt_dq, dq, SEQ*DIM*4);
|
||||
float *capt_dk = xmf((size_t)SEQ*DIM); memcpy(capt_dk, dk, SEQ*DIM*4);
|
||||
float *capt_dv = xmf((size_t)SEQ*DIM); memcpy(capt_dv, dv, 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, ^{
|
||||
cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, DIM, DIM, SEQ,
|
||||
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);
|
||||
|
||||
// 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);
|
||||
|
||||
// dy for next layer (going backward) = dx_rms1 + dx2 residual
|
||||
|
|
|
|||
Loading…
Reference in New Issue