diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml index 0dd5cca0f8..dcf555097b 100644 --- a/.github/workflows/benchmark.yml +++ b/.github/workflows/benchmark.yml @@ -81,6 +81,8 @@ jobs: run: | 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 + - name: Run GPT2 w HALF/BEAM + run: CUDA=1 JIT=1 HALF=1 BEAM=4 CACHELEVEL=0 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half_beam.txt - uses: actions/upload-artifact@v3 with: name: Speed (NVIDIA) @@ -89,6 +91,7 @@ jobs: torch_speed.txt gpt2_unjitted.txt gpt2_jitted.txt + gpt2_half_beam.txt testamdbenchmark: name: AMD Benchmark diff --git a/examples/gpt2.py b/examples/gpt2.py index 7306d5d7a8..108a14134a 100644 --- a/examples/gpt2.py +++ b/examples/gpt2.py @@ -159,9 +159,10 @@ class GPT2: if __name__ == "__main__": Tensor.no_grad = True print(f"using {Device.DEFAULT} backend") + default_prompt = "What is the answer to life, the universe, and everything?" parser = argparse.ArgumentParser(description='Run GPT2 in tinygrad', formatter_class=argparse.ArgumentDefaultsHelpFormatter) - parser.add_argument('--prompt', type=str, default="What is the answer to life, the universe, and everything?", help="Phrase to start with") + parser.add_argument('--prompt', type=str, default=default_prompt, help="Phrase to start with") parser.add_argument('--count', type=int, default=100, help="Max number of tokens to generate") parser.add_argument('--temperature', type=float, default=0.8, help="Temperature in the softmax") parser.add_argument('--model_size', type=str, default="gpt2-medium", help="Size of model to use [gpt2, gpt2-medium, gpt2-large, gpt2-xl]") @@ -191,4 +192,16 @@ if __name__ == "__main__": print('Generating text...') if len(texts) == 1: print(texts[0]) else: - for i,text in enumerate(texts): print(colored(f"Response {i}:", "green"), text) \ No newline at end of file + for i,text in enumerate(texts): print(colored(f"Response {i}:", "green"), text) + + # validate output! + if args.temperature == 0 and args.model_size == "gpt2-medium" and args.count == 10: + expected = { + default_prompt: "What is the answer to life, the universe, and everything?\n\nThe answer is that we are all one", + "Hello.": "Hello. I'm a little late to the party, but", + } + try: + assert texts[0] == expected[args.prompt] + print(colored("output validated", "green")) + except KeyError: + pass