sampling: guard against non-finite logits — was silent token-0 spew (#369)#370
Open
KingIcyCreamProjects wants to merge 2 commits into
Open
sampling: guard against non-finite logits — was silent token-0 spew (#369)#370KingIcyCreamProjects wants to merge 2 commits into
KingIcyCreamProjects wants to merge 2 commits into
Conversation
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>
Per review: the collapse starts one line before the sum — seeding mx=lo[0] means a NaN at index 0 makes mx NaN, every (lo[i]-mx) NaN, and the softmax is doomed at the max-finding, not the normalize. Seed mx from -INFINITY and skip NaNs (x==x), mirroring the argmax_v change; if nothing finite survives, fall back to mx=0 and let the post-sum guard decide. The isfinite(s) guard is now the second line of defense rather than the only one. Clean logits take a byte-identical path (the extra x==x compare is noise next to V expf calls). test_logit_nan gains the NaN-at-index-0 and all-NaN dist_build cases; test_topp's 123-case sweep still passes on this tree, confirming no interaction with the JustVugg#354 heap select. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Fixes #369: on the default sampling path (
TEMP>0,0<NUCLEUS<1), a single NaN/+Inf logit poisoned softmax into an all-NaNg_pbuf, after whichdist_samplesilently emitted token 0 forever — no error, just an unbroken run of one token. The greedy path was equally blind (argmax_vNaN-pinned to index 0).Three small changes, all O(1)/free on the happy path:
argmax_v: skip NaN (x==x) and seed from-INFINITY→ returns the max finite/+Inf entry instead of a NaN-pinned 0. Covers greedy decode and the speculative-verify argmax.dist_buildmax scan (per @woolcoxm's review on sampling: a single NaN/+Inf logit makes dist_build emit token 0 forever, silently (default path) #369): seedmxfrom-INFINITYand NaN-skip the scan — the collapse actually started at the max-finding (mx=lo[0]with a NaN at 0 dooms everylo[i]-mxbefore the sum), not at the normalize. If nothing finite survives, fall back tomx=0and let the post-sum guard decide.dist_buildpost-sum guard: ifsis non-finite or ≤0, collapseg_pbufto a one-hot over the finite argmax and warn once — degrade + diagnose, never silently corrupt. Now the second line of defense, not the only one.Clean logits take a byte-identical path.
Placement vs #354
Rebased onto current
dev(post-#354/#357). As predicted in review, the #354 heap rewrite only touched the nucleus tail — the guard sits right after the normalize, before theif(g_nuc>0 && g_nuc<1.f)block, so the heap never sees an all-NaNg_pbuf.Tests
tests/test_logit_nan.c(wired intoTEST_BINS), 6 assertion blocks: argmax_v NaN-skip / +Inf-pick / all-NaN; dist_build with a NaN mid-vocab → finite one-hot on the max finite logit anddist_sampleemits it (not 0); NaN at index 0 (the mx-seed failure mode from review); all-NaN stays finite; and a clean-logits regression case. Fails on stockdev(first assert: argmax NaN-pinned to 0), passes with this change.Also re-ran
test_topp(123 cases) on this tree: green — no interaction with the #354 partial select.@woolcoxm — rebased, mx-scan point folded in, ready for your review.