From 0f82d92b9d772f579454ac707205f537f65cbeb3 Mon Sep 17 00:00:00 2001 From: chenyu Date: Fri, 3 Oct 2025 14:27:56 +0800 Subject: [PATCH] use float for softmax in llm.py (#12438) fixed numerical issue in `CPU=1 RANGEIFY=1 python3 -m tinygrad.apps.llm` --- tinygrad/apps/llm.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tinygrad/apps/llm.py b/tinygrad/apps/llm.py index e331756eea..a718170259 100644 --- a/tinygrad/apps/llm.py +++ b/tinygrad/apps/llm.py @@ -135,7 +135,7 @@ class Transformer: x = self.token_embd(tokens) # (B, T, D) for block in self.blk: x = block(x, start_pos) # TODO: add temperature - return self.output(self.output_norm(x))[:, -1, :].softmax(-1).argmax(-1, keepdim=True) + return self.output(self.output_norm(x))[:, -1, :].softmax(-1, dtype="float").argmax(-1, keepdim=True) def __call__(self, tokens:Tensor, start_pos:int|UOp=0) -> Tensor: return (self.forward_jit if getenv("JIT", 1) and tokens.shape[1] == 1 and isinstance(start_pos, UOp) else self.forward)(tokens, start_pos)