benchmark.yml

name: Benchmarks

on:
  push:
    branches:
      - master
      - update_benchmark

jobs:
  testmacbenchmark:
    name: Mac Benchmark
    runs-on: [self-hosted, macOS]
    defaults:
      run:
        shell: bash -o pipefail {0}
    if: github.repository_owner == 'tinygrad'
    env:
      PYTHONPATH: .
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
    - name: Symlink models and datasets
      run: |
        mkdir -p weights
        ln -s ~/tinygrad/extra/disassemblers/applegpu extra/disassemblers/applegpu
        ln -s ~/tinygrad/weights/sd-v1-4.ckpt weights/sd-v1-4.ckpt
        ln -s ~/tinygrad/weights/bpe_simple_vocab_16e6.txt.gz weights/bpe_simple_vocab_16e6.txt.gz
        ln -s ~/tinygrad/weights/LLaMA weights/LLaMA
        ln -s ~/tinygrad/extra/datasets/cifar-10-python.tar.gz extra/datasets/cifar-10-python.tar.gz
    - name: Update process replay reference (master only)
      if: github.ref == 'refs/heads/master'
      run: |
        export TRACE_PATH="$HOME/traces/$GITHUB_SHA" && echo "TRACE_PATH=$TRACE_PATH" >> $GITHUB_ENV
        rm -rf "$HOME/traces" && mkdir -p "$HOME/traces"
        echo "SAVE_TRACE=1" >> $GITHUB_ENV
    - name: Setup process replay
      if: contains(github.event.head_commit.message, '[run_process_replay]')
      run: |
        export TRACE_PATH="$HOME/traces/$GITHUB_SHA" && echo "TRACE_PATH=$TRACE_PATH" >> $GITHUB_ENV
        rm -rf "$TRACE_PATH" && touch $TRACE_PATH
        echo "SAVE_TRACE=1" >> $GITHUB_ENV
    - name: Run LLaMA
      run: |
        JIT=0 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_unjitted.txt
        JIT=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_jitted.txt
    - name: Run LLaMA with BEAM
      run: JIT=1 BEAM=2 CACHELEVEL=0 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_beam.txt
    - name: Run process replay tests
      if: contains(github.event.head_commit.message, '[run_process_replay]')
      run: python3 test/external/replay_benchmarks.py
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (Mac)
        path: |
          onnx_inference_speed.csv
          torch_speed.txt
          llama_unjitted.txt
          llama_jitted.txt
          llama_beam.txt
          llama_four_gpu.txt
          gpt2_unjitted.txt
          gpt2_jitted.txt
          gpt2_half.txt
          gpt2_half_beam.txt
          matmul.txt
          matmul_half.txt
          sd.txt
          beautiful_mnist.txt
          train_cifar.txt
          train_cifar_half.txt
          train_cifar_bf16.txt
          train_cifar_wino.txt