tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-30 00:06:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

use float for softmax in llm.py (#12438)

Browse Source 
fixed numerical issue in `CPU=1 RANGEIFY=1 python3 -m tinygrad.apps.llm`
This commit is contained in:
chenyu
2025-10-03 02:27:56 -04:00
committed by GitHub
parent 4c63f7e786
commit 0f82d92b9d
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/apps/llm.py
+1 -1
View File
@@ -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)