berkus
/
ANE
mirror of https://github.com/maderix/ANE.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

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 2026-03-02 23:24:16 +01:00
parent b3936dce86
commit 78666fcc1c
2 changed files with 250 additions and 239 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

43
training/stories_config.h
View File

@ -142,15 +142,26 @@ static void ane_init(void) {
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 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 void adam_free(AdamState *s) { free(s->m); free(s->v); }
static LayerWeights layer_weights_alloc(void) { static LayerWeights layer_weights_alloc(void) {
LayerWeights w; LayerWeights w;
w.Wq=(float*)malloc(WQ_SZ*4); w.Wk=(float*)malloc(WQ_SZ*4); w.Wq=xmf(WQ_SZ); w.Wk=xmf(WQ_SZ);
w.Wv=(float*)malloc(WQ_SZ*4); w.Wo=(float*)malloc(WO_SZ*4); w.Wv=xmf(WQ_SZ); w.Wo=xmf(WO_SZ);
w.W1=(float*)malloc(W1_SZ*4); w.W2=(float*)malloc(W2_SZ*4); w.W3=(float*)malloc(W3_SZ*4); w.W1=xmf(W1_SZ); w.W2=xmf(W2_SZ); w.W3=xmf(W3_SZ);
w.rms_att=(float*)malloc(DIM*4); w.rms_ffn=(float*)malloc(DIM*4); w.rms_att=xmf(DIM); w.rms_ffn=xmf(DIM);
return w; return w;
} }
static void layer_weights_free(LayerWeights *w) { static void layer_weights_free(LayerWeights *w) {
@ -172,13 +183,13 @@ static void layer_adam_free(LayerAdam *a) {
} }
static LayerActs layer_acts_alloc(void) { static LayerActs layer_acts_alloc(void) {
LayerActs a; LayerActs a;
a.layer_in=(float*)malloc(SEQ*DIM*4); a.layer_in=xmf((size_t)SEQ*DIM);
a.xnorm=(float*)malloc(SEQ*DIM*4); a.Q=(float*)malloc(SEQ*DIM*4); a.xnorm=xmf((size_t)SEQ*DIM); a.Q=xmf((size_t)SEQ*DIM);
a.K=(float*)malloc(SEQ*DIM*4); a.V=(float*)malloc(SEQ*DIM*4); a.K=xmf((size_t)SEQ*DIM); a.V=xmf((size_t)SEQ*DIM);
a.attn_out=(float*)malloc(SEQ*DIM*4); a.o_out=(float*)malloc(SEQ*DIM*4); a.attn_out=xmf((size_t)SEQ*DIM); a.o_out=xmf((size_t)SEQ*DIM);
a.x2=(float*)malloc(SEQ*DIM*4); a.x2norm=(float*)malloc(SEQ*DIM*4); a.x2=xmf((size_t)SEQ*DIM); a.x2norm=xmf((size_t)SEQ*DIM);
a.h1=(float*)malloc(SEQ*HIDDEN*4); a.h3=(float*)malloc(SEQ*HIDDEN*4); a.h1=xmf((size_t)SEQ*HIDDEN); a.h3=xmf((size_t)SEQ*HIDDEN);
a.silu_out=(float*)malloc(SEQ*HIDDEN*4); a.ffn_out=(float*)malloc(SEQ*DIM*4); a.silu_out=xmf((size_t)SEQ*HIDDEN); a.ffn_out=xmf((size_t)SEQ*DIM);
return a; return a;
} }
static void layer_acts_free(LayerActs *a) { static void layer_acts_free(LayerActs *a) {
@ -188,10 +199,10 @@ static void layer_acts_free(LayerActs *a) {
} }
static LayerGrads layer_grads_alloc(void) { static LayerGrads layer_grads_alloc(void) {
LayerGrads g; LayerGrads g;
g.Wq=(float*)calloc(WQ_SZ,4); g.Wk=(float*)calloc(WQ_SZ,4); g.Wq=xcf(WQ_SZ); g.Wk=xcf(WQ_SZ);
g.Wv=(float*)calloc(WQ_SZ,4); g.Wo=(float*)calloc(WO_SZ,4); g.Wv=xcf(WQ_SZ); g.Wo=xcf(WO_SZ);
g.W1=(float*)calloc(W1_SZ,4); g.W2=(float*)calloc(W2_SZ,4); g.W3=(float*)calloc(W3_SZ,4); g.W1=xcf(W1_SZ); g.W2=xcf(W2_SZ); g.W3=xcf(W3_SZ);
g.rms_att=(float*)calloc(DIM,4); g.rms_ffn=(float*)calloc(DIM,4); g.rms_att=xcf(DIM); g.rms_ffn=xcf(DIM);
return g; return g;
} }
static void layer_grads_zero(LayerGrads *g) { static void layer_grads_zero(LayerGrads *g) {

64
training/train_large.m
View File

@ -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