guided: nan_to_num natural-path log_softmax inputs

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.
This commit is contained in:
wassname
2026-05-23 08:42:31 +00:00
parent b777c84e22
commit 0b39d2d3f7
+13 -5
View File
@@ -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: