diff --git a/single_shot_inference.py b/single_shot_inference.py index 27a06945..0040f402 100644 --- a/single_shot_inference.py +++ b/single_shot_inference.py @@ -1053,6 +1053,15 @@ def main(): logits = lm_head_lin(x_out) if profile: torch.cuda.synchronize() t_lm = time.perf_counter() + # Check thinking start token logit on first step + if step == 0: + ls = logits.float() + for tid, name in [(THINK_START, 'think_start'), (THINK_END, 'think_end'), (USER_TOKEN, 'user'), (ASSISTANT_TOKEN, 'assistant')]: + print(f" {name}({tid}) logit={ls[0, tid].item():.2f}", flush=True) + # Also check "Paris" token + paris_tids = [t for t in range(min(129280, ls.shape[-1])) if 'Paris' in tokenizer.decode([t])] + if paris_tids: + print(f" Paris tokens: {[(t, ls[0,t].item()) for t in paris_tids[:5]]}", flush=True) # Only sync + validate on first 3 steps and every 20th step (reduces pipeline stalls) if step < 3 or (step + 1) % 20 == 0: torch.cuda.synchronize() # catch CUDA errors at source