From fb1cc6bf4be9325a6161dbabcb35c393c7b50dee Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Tue, 5 Sep 2023 10:12:16 -0700 Subject: [PATCH] llama jit is default, print tok/sec (#1774) * llama jit is default, print tok/sec * jit not default in CI --- .github/workflows/benchmark.yml | 10 +++++----- examples/llama.py | 22 ++++++++++++---------- test/external/external_multi_gpu.py | 2 +- 3 files changed, 18 insertions(+), 16 deletions(-) diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml index c90ea1b107..bdf7f26f6f 100644 --- a/.github/workflows/benchmark.yml +++ b/.github/workflows/benchmark.yml @@ -35,12 +35,12 @@ jobs: - name: Run LLaMA run: | ln -s ~/tinygrad/weights/LLaMA weights/LLaMA - python3 examples/llama.py --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_unjitted.txt + JIT=0 python3 examples/llama.py --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_unjitted.txt JIT=1 python3 examples/llama.py --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_jitted.txt shell: bash - name: Run GPT2 run: | - python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_unjitted.txt + JIT=0 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_unjitted.txt JIT=1 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt shell: bash - name: Run 10 CIFAR training steps @@ -78,7 +78,7 @@ jobs: shell: bash - name: Run GPT2 run: | - CUDA=1 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_unjitted.txt + CUDA=1 JIT=0 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_unjitted.txt CUDA=1 JIT=1 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt shell: bash - uses: actions/upload-artifact@v3 @@ -116,12 +116,12 @@ jobs: - name: Run LLaMA run: | ln -s ~/tinygrad/weights/LLaMA weights/LLaMA - python3 examples/llama.py --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_unjitted.txt + JIT=0 python3 examples/llama.py --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_unjitted.txt JIT=1 python3 examples/llama.py --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_jitted.txt shell: bash - name: Run GPT2 run: | - python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_unjitted.txt + JIT=0 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_unjitted.txt JIT=1 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt shell: bash - name: Run 10 CIFAR training steps diff --git a/examples/llama.py b/examples/llama.py index f3a3deaf6a..5de309e926 100755 --- a/examples/llama.py +++ b/examples/llama.py @@ -9,7 +9,7 @@ import numpy as np np.set_printoptions(linewidth=200) from typing import Optional, Tuple, Dict -from tinygrad.helpers import Timing, getenv, DEBUG, dtypes +from tinygrad.helpers import Timing, getenv, DEBUG, dtypes, CI from tinygrad.ops import Device from tinygrad.tensor import Tensor from tinygrad.nn import Embedding, Linear @@ -18,6 +18,8 @@ from tinygrad.ops import GlobalCounters from tinygrad.jit import TinyJit from tinygrad.shape.symbolic import Variable, sym_infer +JIT = getenv("JIT", 0 if CI else 1) + # https://github.com/facebookresearch/llama/blob/1076b9c51c77ad06e9d7ba8a4c6df775741732bd/llama/model.py#L47 def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0): freqs = 1.0 / (theta ** (np.arange(0, dim, 2, dtype=np.float32)[:(dim // 2)] / dim)) @@ -113,7 +115,7 @@ class TransformerBlock: def __call__(self, x:Tensor, cache_k:Optional[Tensor], cache_v:Optional[Tensor], start_pos:int, freqs_cis:Tensor, mask:Optional[Tensor], jit_ctx:Optional[Dict[Variable,int]]=None): bsz, seqlen, _ = x.shape - if getenv("JIT") and mask is None: + if JIT and mask is None: assert cache_k is not None and cache_v is not None, "no cache" pos = Variable("pos", 1, 1024) cache_k = cache_k.reshape(cache_k.shape[0], pos, cache_k.shape[2], cache_k.shape[3]) @@ -153,7 +155,7 @@ class Transformer: def __call__(self, tokens:Tensor, start_pos:int, temperature:Optional[float]=None): _bsz, seqlen = tokens.shape - if seqlen == 1 and getenv("JIT"): + if seqlen == 1 and JIT: pos = Variable("pos", 1, 1024) freqs_cis = self.freqs_cis.shrink(((0, self.freqs_cis.shape[0]), (pos, pos+seqlen),(0, self.freqs_cis.shape[2]),(0, self.freqs_cis.shape[3]),(0, self.freqs_cis.shape[4]))) freqs_cis.lazydata.st.var_vals[pos] = start_pos @@ -351,12 +353,12 @@ class LLaMa: # **** main code **** """ -test: +test: python3 examples/llama.py --temperature=0 --count=50 --prompt="Hello." output: Hello. I'm a 20 year old male. I'm a student at the University of Texas at Austin. I'm a sophomore majoring in Computer Science. -test: +test: python3 examples/llama.py --gen='2' --temperature=0 --count=50 --prompt="Hello." output: Hello. I'm a 20 year old girl who is looking for a good lay in Palm Coast. I don't care whether it's at your place or not, as long as it's clean. @@ -565,11 +567,11 @@ After you are done speaking, output [EOS]. You are not Chad. if args.timing: print("") st = GlobalCounters.time_sum_s - with Timing("ran model in ", on_exit=(lambda et: f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on GPU"+ - f", {GlobalCounters.global_ops*1e-9:.2f} GOPS, {GlobalCounters.global_mem*1e-9:.2f} GB"+ - f", {GlobalCounters.global_mem*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s") if DEBUG else None, enabled=args.timing): - probs = llama.model(Tensor([toks[start_pos:]]), start_pos, args.temperature).realize() - with Timing("sync in ", enabled=args.timing): + with Timing("total ", enabled=args.timing, on_exit=lambda x: f", {1e9/x:.2f} tok/sec"): + with Timing("ran model in ", on_exit=(lambda et: f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on GPU"+ + f", {GlobalCounters.global_ops*1e-9:.2f} GOPS, {GlobalCounters.global_mem*1e-9:.2f} GB"+ + f", {GlobalCounters.global_mem*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s") if DEBUG else None, enabled=args.timing): + probs = llama.model(Tensor([toks[start_pos:]]), start_pos, args.temperature).realize() probs_np = probs.numpy() tok = int(np.random.choice(len(probs_np), p=probs_np)) diff --git a/test/external/external_multi_gpu.py b/test/external/external_multi_gpu.py index 8edd67f0a8..b500b52770 100644 --- a/test/external/external_multi_gpu.py +++ b/test/external/external_multi_gpu.py @@ -14,7 +14,7 @@ if __name__ == "__main__": sz = 1024*1024*256 # 1 GB #sz = 1024*64 - with Timing("CPU creation: ", on_exit=lambda x: f", {(sz*4*2)/x:.2f} GB/sec"): + with ("CPU creation: ", on_exit=lambda x: f", {(sz*4*2)/x:.2f} GB/sec"): c0 = Tensor.ones(sz, device="cpu").realize() c1 = (Tensor.ones(sz, device="cpu")/2).realize()