From 0b39d2d3f730e22d233e0555f72e53edb19d42b6 Mon Sep 17 00:00:00 2001 From: wassname Date: Sat, 23 May 2026 08:42:31 +0000 Subject: [PATCH] guided: nan_to_num natural-path log_softmax inputs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Qwen3.6-27B nf4 + adapter at c=1.0 produced a non-finite raw logit at a single generated step in 1/4 samples (others used forced-prefill path); the natural-path F.log_softmax propagated NaN into mean_pmass_allowed, crashing c_scan. Bound with nan_to_num(±1e4) — leaves argmax-finite rows unchanged. --- src/tinymfv/guided.py | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/src/tinymfv/guided.py b/src/tinymfv/guided.py index a45cefe..63e8c26 100644 --- a/src/tinymfv/guided.py +++ b/src/tinymfv/guided.py @@ -262,14 +262,22 @@ def _rollout_natural_or_forced( assert answer_pos < len(step_scores), ( f"answer_pos={answer_pos} ≥ len(step_scores)={len(step_scores)}" ) - lp_vec = F.log_softmax(step_scores[answer_pos][i].float(), dim=-1) - # Natural NLL: mean NLL over gen_ids[start_pos:answer_pos] - # using step_scores[start_pos:answer_pos]. By construction - # answer_pos > start_pos so this window is non-empty. + # nan_to_num: quantized + adapted forwards occasionally + # emit non-finite raw logits at a single generated step; + # ±1e4 bound keeps log_softmax stable without changing the + # argmax for well-behaved rows. + raw = step_scores[answer_pos][i].float() + lp_vec = F.log_softmax( + torch.nan_to_num(raw, nan=0.0, posinf=1e4, neginf=-1e4), dim=-1 + ) gen_ids_full = phase1_ids[i, prompt_len:] nat_nll_sum = 0.0 for k in range(start_pos, answer_pos): - step_lp = F.log_softmax(step_scores[k][i].float(), dim=-1) + raw_k = step_scores[k][i].float() + step_lp = F.log_softmax( + torch.nan_to_num(raw_k, nan=0.0, posinf=1e4, neginf=-1e4), + dim=-1, + ) nat_nll_sum += float(-step_lp[gen_ids_full[k]].item()) nll_val = nat_nll_sum / max(1, answer_pos - start_pos) elif not emitted_close_i: