From 61e255d1977dcd45984f364efb43f43fd6e9aab6 Mon Sep 17 00:00:00 2001 From: chenyu Date: Thu, 28 Dec 2023 23:26:00 -0500 Subject: [PATCH] use max for gpt2 and llama (#2949) not using argmax yet because there's a multinomial outside of function. --- examples/gpt2.py | 8 +++++--- extra/models/llama.py | 15 +++++++++------ 2 files changed, 14 insertions(+), 9 deletions(-) diff --git a/examples/gpt2.py b/examples/gpt2.py index 568b420064..0020373897 100644 --- a/examples/gpt2.py +++ b/examples/gpt2.py @@ -93,9 +93,11 @@ class Transformer: for hi in self.h: h = hi(h, start_pos, mask) - logits = self.lm_head(self.ln_f(h)) - # NOTE: temperature=0 with HALF breaks due to precision, should use argmax instead - ret = (logits[:, -1, :] / (temperature+1e-6)).softmax() + logits = self.lm_head(self.ln_f(h))[:, -1, :].flatten() + if temperature < 1e-6: + ret = (logits == logits.max()) + else: + ret = (logits / temperature).softmax() return ret.half().realize() if HALF else ret.realize() # TODO: fix empty token diff --git a/extra/models/llama.py b/extra/models/llama.py index c4f41b98c0..d97072df9b 100644 --- a/extra/models/llama.py +++ b/extra/models/llama.py @@ -4,7 +4,7 @@ from tinygrad.helpers import getenv # https://github.com/facebookresearch/llama/blob/1076b9c51c77ad06e9d7ba8a4c6df775741732bd/llama/model.py#L47 def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0) -> Tensor: - freqs = 1.0 / (theta ** (Tensor.arange(0, dim, 2, dtype=dtypes.default_float)[:(dim // 2)] / dim)) + freqs = 1.0 / (theta ** (Tensor.arange(0, dim, 2, dtype=dtypes.half)[:(dim // 2)] / dim)) freqs = Tensor.arange(end).unsqueeze(dim=1)*freqs.unsqueeze(dim=0) return Tensor.stack([Tensor.cos(freqs), Tensor.sin(freqs)], dim=-1).reshape(1, end, 1, dim//2, 2) @@ -37,7 +37,7 @@ class RMSNorm: def __call__(self, x:Tensor): x = x.float() - return ((x * (x.pow(2).mean(-1, keepdim=True) + self.eps).rsqrt()) * self.weight) + return (x * (x.pow(2).mean(-1, keepdim=True) + self.eps).rsqrt()) * self.weight class Attention: def __init__(self, dim, n_heads, n_kv_heads, max_context, linear=nn.Linear): @@ -66,8 +66,9 @@ class Attention: self.cache_k = Tensor.zeros(bsz, self.max_context, self.n_kv_heads, self.head_dim, dtype=x.dtype) self.cache_v = Tensor.zeros(bsz, self.max_context, self.n_kv_heads, self.head_dim, dtype=x.dtype) - keys = self.cache_k.shrink((None, (0, start_pos), None, None)).cat(xk, dim=1).contiguous() - values = self.cache_v.shrink((None, (0, start_pos), None, None)).cat(xv, dim=1).contiguous() + # TODO: fix coder, old hack did not work after the uop dtype check + keys = self.cache_k.shrink((None, (0, start_pos), None, None)).cat(xk, dim=1) + values = self.cache_v.shrink((None, (0, start_pos), None, None)).cat(xv, dim=1) # update the cache self.cache_k.assign(keys.pad((None,(0,self.max_context-start_pos-seqlen),None,None)).contiguous()).realize() @@ -116,8 +117,10 @@ class Transformer: h = self.tok_embeddings(tokens) mask = Tensor.full((1, 1, seqlen, start_pos+seqlen), float("-inf"), dtype=h.dtype).triu(start_pos+1).realize() if seqlen > 1 else None for layer in self.layers: h = layer(h, start_pos, freqs_cis, mask) - logits = self.output(self.norm(h)) - return (logits[:, -1, :] / (temperature+1e-6)).softmax().flatten().realize() + logits = self.output(self.norm(h))[:, -1, :].flatten() + if temperature < 1e-6: + return (logits == logits.max()).half().realize() + return (logits / temperature).softmax().half().realize() def __call__(self, tokens:Tensor, start_pos:Variable, temperature:float=0.0): # TODO: better way to handle the first call v.s. the rest?