sampling: guard against non-finite logits (was silent token-0 spew)
On the default serve path (TEMP>0, 0<NUCLEUS<1) a single NaN or +Inf in the logits — a bad streamed expert tile, or an fp overflow in the matmul at a low-RAM eviction boundary — poisoned softmax: g_pbuf became all-NaN, dist_sample never satisfied cum>=u, and the fallback returned token 0. The engine then emitted an unbroken run of token 0 with NO error. The greedy path was equally blind: argmax_v started bv=lo[0] and `lo[i]>NaN` is always false, so a NaN at index 0 pinned the argmax to 0. - argmax_v: skip NaN (x==x) and seed from -inf, so it returns the max finite/+Inf entry instead of being NaN-pinned to 0. Covers greedy decode and the speculative-verify argmax path. - dist_build: after the softmax sum, if s is non-finite or <=0, collapse g_pbuf to a one-hot over the finite argmax and warn once, instead of dividing every entry into NaN. Covers the nucleus and verify paths. Both are O(1)/free on the happy path (one branch after the existing loop; one extra comparison inside the existing argmax loop). Degrade + diagnose, never silently corrupt. test_logit_nan (wired into TEST_BINS): asserts argmax_v skips NaN/picks +Inf, dist_build yields a finite normalized one-hot on the max finite logit, dist_sample emits that token (not 0), and clean logits still give a valid distribution. Fails on stock dev, passes with this change. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
KingIcyCreamProjects
+4
-1
@@ -166,7 +166,7 @@ else
|
||||
PYTHON ?= python3
|
||||
endif
|
||||
CUDA_OBJ =
|
||||
TEST_BINS = tests/test_json$(EXE) tests/test_st$(EXE) tests/test_st_pread$(EXE) tests/test_tier$(EXE) tests/test_grammar$(EXE) tests/test_schema_gbnf$(EXE) tests/test_decode_batch$(EXE) tests/test_idot$(EXE) tests/test_i4_grouped$(EXE) tests/test_stops$(EXE) tests/test_topp$(EXE) tests/test_kv_alloc$(EXE) tests/test_i4_acc512$(EXE) tests/test_compat_direct$(EXE) tests/test_dsa_select$(EXE)
|
||||
TEST_BINS = tests/test_json$(EXE) tests/test_st$(EXE) tests/test_st_pread$(EXE) tests/test_tier$(EXE) tests/test_grammar$(EXE) tests/test_schema_gbnf$(EXE) tests/test_decode_batch$(EXE) tests/test_idot$(EXE) tests/test_i4_grouped$(EXE) tests/test_stops$(EXE) tests/test_topp$(EXE) tests/test_kv_alloc$(EXE) tests/test_i4_acc512$(EXE) tests/test_compat_direct$(EXE) tests/test_dsa_select$(EXE) tests/test_logit_nan$(EXE)
|
||||
ifneq (,$(LINUX))
|
||||
TEST_BINS += tests/test_uring$(EXE)
|
||||
endif
|
||||
@@ -331,6 +331,9 @@ tests/bench_topp$(EXE): tests/bench_topp.c glm.c st.h uring.h json.h tok.h tok_u
|
||||
tests/test_kv_alloc$(EXE): tests/test_kv_alloc.c glm.c st.h json.h tok.h tok_unicode.h compat.h grammar.h tier.h
|
||||
$(CC) $(CFLAGS) $< -o $@ $(LDFLAGS)
|
||||
|
||||
tests/test_logit_nan$(EXE): tests/test_logit_nan.c glm.c st.h uring.h json.h tok.h tok_unicode.h compat.h grammar.h tier.h
|
||||
$(CC) $(CFLAGS) $< -o $@ $(LDFLAGS)
|
||||
|
||||
tests/test_i4_acc512$(EXE): tests/test_i4_acc512.c
|
||||
$(CC) $(CFLAGS) $< -o $@ $(LDFLAGS)
|
||||
|
||||
|
||||
@@ -4151,7 +4151,11 @@ static void mtp_absorb(Model *m, const int *next_ids, const float *x, int S, int
|
||||
}
|
||||
|
||||
static inline int argmax_v(const float *lo, int V){
|
||||
int b=0; float bv=lo[0]; for(int i=1;i<V;i++) if(lo[i]>bv){bv=lo[i];b=i;} return b;
|
||||
/* skip NaN (x==x is false for NaN) so a poisoned logit can't pin the argmax
|
||||
* to index 0 — pick the max finite/+Inf entry instead. */
|
||||
int b=-1; float bv=-INFINITY;
|
||||
for(int i=0;i<V;i++){ float x=lo[i]; if(x==x && x>bv){ bv=x; b=i; } }
|
||||
return b<0?0:b;
|
||||
}
|
||||
|
||||
/* ---- METODO F: draft grammaticale (#48) ----
|
||||
@@ -4258,6 +4262,15 @@ static void dist_build(const float *lo, int V){
|
||||
float mx=lo[0]; for(int i=1;i<V;i++) if(lo[i]>mx) mx=lo[i];
|
||||
double s=0; float invt=1.f/(g_temp>1e-4f?g_temp:1e-4f);
|
||||
for(int i=0;i<V;i++){ g_pbuf[i]=expf((lo[i]-mx)*invt); s+=g_pbuf[i]; }
|
||||
/* A single NaN/+Inf logit (bad streamed expert tile, matmul overflow at an
|
||||
* eviction boundary) makes s non-finite; the old code then divided every entry
|
||||
* into NaN and dist_sample silently returned token 0 forever. Collapse to a
|
||||
* one-hot over the max finite logit and warn once — degrade, don't corrupt. */
|
||||
if(!(s==s) || s<=0.0 || s>1e300){
|
||||
static int warned=0; if(!warned){ warned=1; fprintf(stderr,"[sample] non-finite logits; emitting argmax of finite entries\n"); }
|
||||
int b=argmax_v(lo,V); for(int i=0;i<V;i++) g_pbuf[i]=(i==b)?1.f:0.f;
|
||||
return;
|
||||
}
|
||||
for(int i=0;i<V;i++) g_pbuf[i]/=(float)s;
|
||||
if(g_nuc>0 && g_nuc<1.f){
|
||||
for(int i=0;i<V;i++) g_pidx[i]=i;
|
||||
|
||||
@@ -0,0 +1,55 @@
|
||||
/* Regression for non-finite-logit poisoning of sampling.
|
||||
*
|
||||
* A single NaN or +Inf in the logits (a bad streamed expert tile, or an fp
|
||||
* overflow in the matmul at a low-RAM eviction boundary) used to make softmax
|
||||
* produce an all-NaN g_pbuf; dist_sample then never satisfied cum>=u and fell
|
||||
* through to return token 0 — so the engine silently emitted an unbroken run of
|
||||
* token 0 with no error, on the DEFAULT serve path (TEMP>0, 0<NUCLEUS<1).
|
||||
*
|
||||
* Fix under test: argmax_v() skips NaN (picks the max finite/+Inf entry instead
|
||||
* of being pinned to index 0), and dist_build() detects a non-finite softmax sum
|
||||
* and collapses to a one-hot over the finite argmax (warning once) rather than
|
||||
* dividing every entry into NaN. Degrade + diagnose, never silently corrupt.
|
||||
*
|
||||
* No model file needed: exercises argmax_v / dist_build / dist_sample directly. */
|
||||
#include <assert.h>
|
||||
#include <math.h>
|
||||
#define main coli_glm_main_unused
|
||||
#include "../glm.c"
|
||||
#undef main
|
||||
|
||||
static int approx1(double x){ return x > 0.999 && x < 1.001; }
|
||||
|
||||
int main(void){
|
||||
/* --- argmax_v must skip NaN (greedy decode + speculative-verify paths) --- */
|
||||
{ float lo[8]={NAN,1.f,5.f,2.f,NAN,-3.f,4.f,0.f};
|
||||
assert(argmax_v(lo,8)==2 && "pick max finite (idx2=5.0), not NaN-pinned idx0"); }
|
||||
{ float lo[8]={3.f,INFINITY,1.f,2.f,0.f,-1.f,2.5f,1.5f};
|
||||
assert(argmax_v(lo,8)==1 && "pick the +Inf position"); }
|
||||
{ float lo[8]; for(int i=0;i<8;i++) lo[i]=NAN;
|
||||
assert(argmax_v(lo,8)==0 && "all-NaN: no crash, defined fallback"); }
|
||||
|
||||
g_temp=0.7f; g_nuc=0.9f; /* the default serve/chat sampling path */
|
||||
|
||||
/* --- dist_build: a NaN logit must yield a finite one-hot, not all-NaN --- */
|
||||
{ float lo[8]={0.5f,1.f,NAN,8.f,0.2f,-1.f,0.f,0.3f}; /* max finite = idx3 (8.0) */
|
||||
dist_build(lo,8);
|
||||
double sum=0; int nan=0;
|
||||
for(int i=0;i<8;i++){ if(!(g_pbuf[i]==g_pbuf[i])) nan=1; sum+=g_pbuf[i]; }
|
||||
assert(!nan && "g_pbuf must be finite after a NaN logit");
|
||||
assert(approx1(sum) && "g_pbuf must normalize to 1");
|
||||
assert(approx1(g_pbuf[3]) && "mass must land on the max finite logit (idx3)");
|
||||
assert(dist_sample(8,-1)==3 && "sampler emits the finite argmax, not token 0"); }
|
||||
|
||||
/* --- regression: clean logits still produce a valid distribution --- */
|
||||
{ float lo[8]={0.1f,0.2f,3.0f,0.4f,0.5f,0.6f,0.7f,0.8f}; /* peak = idx2 */
|
||||
dist_build(lo,8);
|
||||
double sum=0; int nan=0;
|
||||
for(int i=0;i<8;i++){ if(!(g_pbuf[i]==g_pbuf[i])) nan=1; sum+=g_pbuf[i]; }
|
||||
assert(!nan && "clean softmax stays finite");
|
||||
assert(approx1(sum) && "clean softmax must sum to 1");
|
||||
assert(g_pbuf[2]>=g_pbuf[0] && "peak token keeps the most mass"); }
|
||||
|
||||
printf("OK test_logit_nan: argmax_v NaN-skip + dist_build finite-collapse\n");
|
||||
return 0;
|
||||
}
|
||||
Reference in New Issue
Block a user