sampling: survive non-finite logits instead of emitting token 0 forever (#369)
A single NaN or +Inf logit silently broke the default sampling path. +Inf became `mx`, then `expf((Inf-mx))`/`expf((NaN-mx))` is NaN, the softmax sum went NaN, every probability went NaN — and dist_sample's fallback loop `if(g_pbuf[i]>0)` is false for NaN at every index, so it returned 0. Every subsequent token: 0. No error, no warning. @KingIcyCreamProjects found it. dist_build now takes `mx` over finite logits only, gives a non-finite logit probability 0, and when the distribution is unusable (no finite logit, or a non-finite/zero sum) collapses to a delta on the finite argmax and warns ONCE on stderr — degraded, but a valid token and a visible cause, never a silent stream of zeros. The finite argmax uses the index found during the mx pass (robust even when lo[0] itself is NaN, where argmax_v would wrongly return 0). tests/test_sample_nan.c: healthy logits still sample correctly; NaN/+Inf injected at lo[0], the middle, and the last position all pick the finite argmax; an all-non-finite vocab leaves no NaN in the buffer and doesn't crash. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
JustVugg
+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_sample_nan$(EXE) tests/test_kv_alloc$(EXE) tests/test_i4_acc512$(EXE) tests/test_compat_direct$(EXE) tests/test_dsa_select$(EXE)
|
||||
ifneq (,$(LINUX))
|
||||
TEST_BINS += tests/test_uring$(EXE)
|
||||
endif
|
||||
@@ -328,6 +328,9 @@ tests/test_topp$(EXE): tests/test_topp.c glm.c st.h uring.h json.h tok.h tok_uni
|
||||
tests/bench_topp$(EXE): tests/bench_topp.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_sample_nan$(EXE): tests/test_sample_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_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)
|
||||
|
||||
|
||||
@@ -4296,9 +4296,28 @@ static void topp_siftdown(int *h, int n, int i){
|
||||
* la coda troncata va AZZERATA in g_pbuf (dist_sample la legge direttamente per id). */
|
||||
static void dist_build(const float *lo, int V){
|
||||
if(!g_pbuf){ g_pbuf=falloc(V); g_pidx=malloc(V*sizeof(int)); }
|
||||
float mx=lo[0]; for(int i=1;i<V;i++) if(lo[i]>mx) mx=lo[i];
|
||||
/* Un solo logit NaN/+Inf avvelenava tutto (#369): +Inf diventava mx, NaN/Inf-mx
|
||||
* -> expf NaN -> s NaN -> ogni prob NaN -> dist_sample cade sul fallback
|
||||
* `g_pbuf[i]>0` (NaN>0 e' falso ovunque) e ritorna 0 PER SEMPRE, in silenzio.
|
||||
* Difesa: mx solo sui finiti; un logit non finito contribuisce prob 0;
|
||||
* se la distribuzione degenera (tutti non finiti / somma non valida) si
|
||||
* ripiega sull'argmax dei finiti e si avvisa UNA volta, mai in silenzio. */
|
||||
int mxi=-1; float mx=0;
|
||||
for(int i=0;i<V;i++) if(isfinite(lo[i]) && (mxi<0 || lo[i]>mx)){ mx=lo[i]; mxi=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]; }
|
||||
if(mxi>=0){
|
||||
for(int i=0;i<V;i++){ g_pbuf[i]=isfinite(lo[i])?expf((lo[i]-mx)*invt):0.f; s+=g_pbuf[i]; }
|
||||
}
|
||||
if(mxi<0 || !isfinite(s) || s<=0.0){ /* distribuzione inutilizzabile */
|
||||
static int warned=0;
|
||||
if(!warned){ warned=1; fprintf(stderr,
|
||||
"[SAMPLE] warning: non-finite logits (NaN/Inf) — falling back to argmax; "
|
||||
"output may be degraded. This usually means a numerical blow-up upstream.\n"); }
|
||||
int a=(mxi>=0)?mxi:0; /* mxi = argmax dei logit FINITI (robusto
|
||||
* anche se lo[0] e' NaN, dove argmax_v fallirebbe) */
|
||||
for(int i=0;i<V;i++) g_pbuf[i]=0.f; g_pbuf[a]=1.f;
|
||||
return; /* delta su un token valido, niente top-p su NaN */
|
||||
}
|
||||
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,62 @@
|
||||
/* Regressione #369: un solo logit NaN o +Inf faceva emettere a dist_build/dist_sample
|
||||
* il token 0 PER SEMPRE, in silenzio (il fallback `g_pbuf[i]>0` e' falso su NaN ovunque).
|
||||
* Ora la distribuzione degenere ripiega sull'argmax dei logit FINITI e avvisa una volta.
|
||||
*
|
||||
* Il test verifica: (a) con logit sani il campionamento resta corretto; (b) con NaN/+Inf
|
||||
* iniettato il token scelto e' l'argmax dei FINITI (mai 0 per default), su ogni posizione
|
||||
* del NaN inclusa lo[0]; (c) nessun NaN sopravvive in g_pbuf. */
|
||||
#define main coli_glm_main_unused
|
||||
#include "../glm.c"
|
||||
#undef main
|
||||
#include <stdio.h>
|
||||
|
||||
static uint32_t rs=0x1234abcd; static uint32_t xr(){rs^=rs<<13;rs^=rs>>17;rs^=rs<<5;return rs;}
|
||||
|
||||
static int pbuf_has_nan(int V){ for(int i=0;i<V;i++) if(!isfinite(g_pbuf[i])) return 1; return 0; }
|
||||
|
||||
int main(void){
|
||||
int fail=0, V=2000;
|
||||
float *lo=malloc(V*sizeof(float));
|
||||
g_temp=0.7f; g_nuc=0.90f; /* il path serve di default */
|
||||
|
||||
/* (a) logit sani: il campionato deve avere prob > 0 e nessun NaN nel buffer */
|
||||
for(int i=0;i<V;i++) lo[i]=(float)((int)(xr()%2000)-1000)/100.f;
|
||||
int known=1337; lo[known]=50.f; /* picco netto */
|
||||
dist_build(lo,V);
|
||||
if(pbuf_has_nan(V)){ printf(" FAIL: NaN in g_pbuf con logit sani\n"); fail=1; }
|
||||
if(g_pbuf[known]<=0.f){ printf(" FAIL: il picco ha prob 0\n"); fail=1; }
|
||||
if(!fail) printf(" logit sani: distribuzione valida, picco vivo ok\n");
|
||||
|
||||
/* (b) NaN/+Inf iniettato in varie posizioni; il finito-argmax deve vincere */
|
||||
float bad[3]; bad[0]=NAN; bad[1]=INFINITY; bad[2]=-INFINITY;
|
||||
const char *bn[3]={"NaN","+Inf","-Inf"};
|
||||
for(int b=0;b<3;b++){
|
||||
for(int pos=0;pos<3;pos++){ /* lo[0], meta', ultimo */
|
||||
for(int i=0;i<V;i++) lo[i]=(float)((int)(xr()%400)-200)/100.f;
|
||||
int amax=777; lo[amax]=9.0f; /* massimo FINITO atteso */
|
||||
int at = pos==0?0 : pos==1?V/2 : V-1;
|
||||
if(at==amax) amax=amax+1; /* non sovrapporre */
|
||||
lo[amax]=9.0f;
|
||||
lo[at]=bad[b]; /* veleno */
|
||||
dist_build(lo,V);
|
||||
if(pbuf_has_nan(V)){ printf(" FAIL: NaN sopravvive (%s @ %d)\n",bn[b],at); fail=1; continue; }
|
||||
/* con -Inf il fallback puo' non scattare (max finito resta), ma il buffer
|
||||
* deve restare valido e sommare ~1: con +Inf/NaN scatta il delta su amax */
|
||||
int picked=-1; float pv=-1;
|
||||
for(int i=0;i<V;i++) if(g_pbuf[i]>pv){pv=g_pbuf[i];picked=i;}
|
||||
if(b<2 && picked!=amax){ /* NaN e +Inf: delta esatto su amax */
|
||||
printf(" FAIL: %s @ %d -> picked %d, atteso argmax finito %d\n",bn[b],at,picked,amax); fail=1;
|
||||
}
|
||||
}
|
||||
}
|
||||
if(!fail) printf(" NaN/+Inf iniettato: argmax dei finiti vince, mai 0/NaN ok\n");
|
||||
|
||||
/* (c) caso estremo: TUTTI non finiti -> non deve crashare, buffer valido */
|
||||
for(int i=0;i<V;i++) lo[i]=NAN;
|
||||
dist_build(lo,V);
|
||||
if(pbuf_has_nan(V)){ printf(" FAIL: tutti-NaN lascia NaN nel buffer\n"); fail=1; }
|
||||
else printf(" tutti non-finiti: nessun crash, buffer valido ok\n");
|
||||
|
||||
printf(fail?"test_sample_nan: FAIL\n":"test_sample_nan: ok\n");
|
||||
return fail;
|
||||
}
|
||||
Reference in New Issue
Block a user