remove AMD_LLVM=0 in mlperf and search ci

tinybox updated to llvm 20
feat: bump to llvm20 (#10784 )
2025-06-11 16:16:30 -07:00 · 2025-06-11 16:04:18 -07:00 · 2025-06-11 16:28:40 -04:00 · 2025-06-11 23:20:41 +03:00 · 2025-06-11 16:00:38 -04:00 · 2025-06-11 14:56:45 -04:00
595 changed files with 37769 additions and 1838534 deletions
@@ -11,5 +11,5 @@ runs:
        git fetch origin $CURRENT_SHA
        export COMMIT_MESSAGE=$(git show -s --format=%B "$CURRENT_SHA")
        export CURRENT_HEAD=$(git rev-parse HEAD)
-        cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
+        cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && IGNORE_OOB=1 PYTHONPATH=. python3 process_replay.py
        git checkout $CURRENT_HEAD  # restore to branch
@@ -41,31 +41,19 @@ runs:
  using: "composite"
  steps:
    - name: Set up Python ${{ inputs.python-version }}
+      id: setup-python
      uses: actions/setup-python@v5
      with:
        python-version: ${{ inputs.python-version }}

    # **** Caching packages ****
-    # TODO: key should include input.deps, but it can't since it can't contain commas

-    - name: Cache Python packages (Linux)
-      if: inputs.key != '' && runner.os == 'Linux'
+    - name: Cache Python packages
+      id: restore-venv
      uses: actions/cache@v4
      with:
-        path: ${{ env.Python3_ROOT_DIR }}/lib/python${{ inputs.python-version }}/site-packages
-        key: python-package-${{ inputs.key }}-${{ hashFiles('**/setup.py') }}
-    - name: Cache Python packages (macOS)
-      if: inputs.key != '' && runner.os == 'macOS'
-      uses: actions/cache@v4
-      with:
-        path: /Users/runner/Library/Python/${{ inputs.python-version }}/lib/python/site-packages
-        key: osx-python-package-${{ inputs.key }}-${{ hashFiles('**/setup.py') }}
-    - name: Cache Python packages (Windows)
-      if: inputs.key != '' && runner.os == 'Windows'
-      uses: actions/cache@v4
-      with:
-        path: ${{ env.Python3_ROOT_DIR }}\Lib\site-packages
-        key: windows-python-package-${{ inputs.key }}-${{ hashFiles('**/setup.py') }}
+        path: ${{ github.workspace }}/.venv
+        key: venv-${{ runner.os }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ hashFiles('**/setup.py') }}-${{ env.PYTHON_CACHE_VERSION }}

    # **** Caching downloads ****

@@ -84,52 +72,118 @@ runs:

    # **** Python deps ****

-    - name: Install dependencies (with extra)
-      if: inputs.deps != ''
-      shell: bash
-      run: pip install ${{ (runner.os == 'macOS' && '--user') || (runner.os != 'macOS' && '') }} -e ".[${{ inputs.deps }}]" ${{ inputs.pydeps }} --extra-index-url https://download.pytorch.org/whl/cpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/Triton-Nightly/pypi/simple/
-    - name: Install dependencies (without extra)
-      if: inputs.deps == ''
-      shell: bash
-      run: pip install ${{ (runner.os == 'macOS' && '--user') || (runner.os != 'macOS' && '') }} -e . ${{ inputs.pydeps }}
-
-    # **** OpenCL ****
-
-    - name: Install OpenCL
-      if: inputs.opencl == 'true'
+    - name: Install dependencies in venv (with extra)
+      if: inputs.deps != '' && steps.restore-venv.outputs.cache-hit != 'true'
      shell: bash
      run: |
-        echo 'Acquire::http::Pipeline-Depth "5";' | sudo tee -a /etc/apt/apt.conf.d/99parallel
-        echo "deb [ allow-insecure=yes ] https://apt.repos.intel.com/oneapi all main" | sudo tee /etc/apt/sources.list.d/oneAPI.list
-        sudo apt update || true
-        sudo apt install --allow-unauthenticated -y --no-install-recommends opencl-headers \
-          intel-oneapi-runtime-openmp=2023.2.1-16 intel-oneapi-runtime-compilers-common=2023.2.1-16 intel-oneapi-runtime-compilers=2023.2.1-16 \
-          intel-oneapi-runtime-dpcpp-sycl-opencl-cpu=2023.2.1-16 intel-oneapi-runtime-tbb-common=2021.10.0-49541 \
-          intel-oneapi-runtime-tbb=2021.10.0-49541 intel-oneapi-runtime-opencl=2023.2.1-16
+        python -m venv .venv
+        if [[ "$RUNNER_OS" == "Windows" ]]; then
+          source .venv/Scripts/activate
+        else
+          . .venv/bin/activate
+        fi
+        python -m pip install -e ".[${{ inputs.deps }}]" ${{ inputs.pydeps }} --extra-index-url https://download.pytorch.org/whl/cpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/Triton-Nightly/pypi/simple/
+    - name: Install dependencies in venv (without extra)
+      if: inputs.deps == '' && steps.restore-venv.outputs.cache-hit != 'true'
+      shell: bash
+      run: |
+        python -m venv .venv
+        if [[ "$RUNNER_OS" == "Windows" ]]; then
+          source .venv/Scripts/activate
+        else
+          . .venv/bin/activate
+        fi
+        python -m pip install -e . ${{ inputs.pydeps }}
+    - name: Set up venv environment
+      shell: bash
+      run: |
+        echo "VIRTUAL_ENV=${{ github.workspace }}/.venv" >> "$GITHUB_ENV"
+        echo "OMP_NUM_THREADS=1" >> "$GITHUB_ENV"
+        # no buffers should be over 300MB in CI
+        echo "MAX_BUFFER_SIZE=300000000" >> "$GITHUB_ENV"
+        if [[ "$RUNNER_OS" == "Windows" ]]; then
+          echo "${{ github.workspace }}/.venv/Scripts" >> "$GITHUB_PATH"
+        else
+          echo "${{ github.workspace }}/.venv/bin" >> "$GITHUB_PATH"
+        fi

-    # **** AMD ****
+    # ******************* apt *******************

-    - name: Install AMD (Linux)
+    - name: Add OpenCL Repo
+      if: inputs.opencl == 'true' && runner.os == 'Linux'
+      shell: bash
+      run: echo "deb [ allow-insecure=yes ] https://apt.repos.intel.com/oneapi all main" | sudo tee /etc/apt/sources.list.d/oneAPI.list
+
+    - name: Add AMD Repo (Linux)
      if: inputs.amd == 'true' && runner.os == 'Linux'
      shell: bash
      run: |
-        echo 'Acquire::http::Pipeline-Depth "5";' | sudo tee -a /etc/apt/apt.conf.d/99parallel
        wget https://repo.radeon.com/rocm/rocm.gpg.key -O - | gpg --dearmor | sudo tee /etc/apt/keyrings/rocm.gpg > /dev/null
-        sudo tee /etc/apt/sources.list.d/rocm.list <<'EOF'
-        deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/6.1.2 jammy main
+        sudo tee /etc/apt/sources.list.d/rocm.list <<EOF
+        deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/6.2 $(lsb_release -cs) main
        EOF
        echo -e 'Package: *\nPin: release o=repo.radeon.com\nPin-Priority: 600' | sudo tee /etc/apt/preferences.d/rocm-pin-600
-        sudo apt update || true
-        sudo apt install --no-install-recommends --allow-unauthenticated -y hsa-rocr comgr hsa-rocr-dev liburing-dev libc6-dev
-        curl -s https://api.github.com/repos/Qazalin/remu/releases/latest | \
-        jq -r '.assets[] | select(.name == "libremu.so").browser_download_url' | \
-        sudo xargs curl -L -o /usr/local/lib/libremu.so
+
+    - name: Add LLVM Repo (Linux)
+      if: inputs.llvm == 'true' && runner.os == 'Linux'
+      shell: bash
+      run: |
+        wget -qO- https://apt.llvm.org/llvm-snapshot.gpg.key | sudo tee /etc/apt/trusted.gpg.d/apt.llvm.org.asc
+        echo "deb http://apt.llvm.org/$(lsb_release -cs)/ llvm-toolchain-$(lsb_release -cs)-20 main" | sudo tee /etc/apt/sources.list.d/llvm.list
+
+    - name: apt-get update + install
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true')
+      shell: bash
+      run: |
+        echo 'Acquire::GzipIndexes "true";' | sudo tee /etc/apt/apt.conf.d/gzip
+        echo 'Acquire::http::Pipeline-Depth "5";' | sudo tee -a /etc/apt/apt.conf.d/99parallel
+        sudo apt -qq update || true
+
+        pkgs=""
+        # **** OpenCL ****
+        if [[ "${{ inputs.opencl }}" == "true" ]]; then
+          pkgs+=" opencl-headers \
+            intel-oneapi-runtime-openmp=2023.2.1-16 intel-oneapi-runtime-compilers-common=2023.2.1-16 intel-oneapi-runtime-compilers=2023.2.1-16 \
+            intel-oneapi-runtime-dpcpp-sycl-opencl-cpu=2023.2.1-16 intel-oneapi-runtime-tbb-common=2021.10.0-49541 \
+            intel-oneapi-runtime-tbb=2021.10.0-49541 intel-oneapi-runtime-opencl=2023.2.1-16"
+        fi
+        # **** AMD ****
+        if [[ "${{ inputs.amd }}" == "true" ]]; then
+          pkgs+=" hsa-rocr comgr hsa-rocr-dev liburing-dev libc6-dev"
+        fi
+        # **** CUDA ****
+        if [[ "${{ inputs.cuda }}" == "true" ]]; then
+          pkgs+=" git g++ cmake ninja-build llvm-15-dev zlib1g-dev libglew-dev \
+            flex bison libfl-dev libboost-thread-dev libboost-filesystem-dev nvidia-cuda-toolkit-gcc libzstd-dev"
+        fi
+        # **** WebGPU (dependencies for software-based vulkan) ****
+        if [[ "${{ inputs.webgpu }}" == "true" ]]; then
+          pkgs+=" libgl1 libglx-mesa0 libgl1-mesa-dri libxcb-xfixes0-dev mesa-vulkan-drivers"
+        fi
+        # **** LLVM ****
+        if [[ "${{ inputs.llvm }}" == "true" ]]; then
+          pkgs+=" libllvm20 clang-20 lld-20"
+        fi
+
+        # ******** do install ********
+        if [[ -n "$pkgs" ]]; then
+          sudo apt-get -y --allow-unauthenticated --no-install-recommends install $pkgs
+        fi
+
+    # **** AMD ****
+
+    - name: Setup AMD (Linux)
+      if: inputs.amd == 'true' && runner.os == 'Linux'
+      shell: bash
+      run: |
+        cargo build --release --manifest-path ./extra/remu/Cargo.toml
+        sudo ln -sf ${{ github.workspace }}/extra/remu/target/release/libremu.so /usr/local/lib/libremu.so
        sudo tee --append /etc/ld.so.conf.d/rocm.conf <<'EOF'
          /opt/rocm/lib
          /opt/rocm/lib64
        EOF
        sudo ldconfig
-    - name: Install AMD comgr+remu (macOS)
+    - name: Setup AMD comgr+remu (macOS)
      if: inputs.amd == 'true' && runner.os == 'macOS'
      shell: bash
      run: |
@@ -137,26 +191,14 @@ runs:
        curl -s -H "Authorization: token $GH_TOKEN" curl -s https://api.github.com/repos/nimlgen/amdcomgr_dylib/releases/latest | \
          jq -r '.assets[] | select(.name == "libamd_comgr.dylib").browser_download_url' | \
          sudo xargs curl -L -o /usr/local/lib/libamd_comgr.dylib
-        curl -s -H "Authorization: token $GH_TOKEN" curl -s https://api.github.com/repos/Qazalin/remu/releases/latest | \
-          jq -r '.assets[] | select(.name == "libremu.dylib").browser_download_url' | \
-          sudo xargs curl -L -o /usr/local/lib/libremu.dylib
+        cargo build --release --manifest-path ./extra/remu/Cargo.toml

    # **** CUDA ****

-    - name: Install cuda packages (Linux)
-      if: inputs.cuda == 'true' && runner.os == 'Linux'
-      shell: bash
-      run: |
-        echo 'Acquire::http::Pipeline-Depth "5";' | sudo tee -a /etc/apt/apt.conf.d/99parallel
-        sudo apt update -y || true
-        sudo apt install -y --no-install-recommends git g++ cmake ninja-build llvm-15-dev zlib1g-dev libglew-dev \
-          flex bison libfl-dev libboost-thread-dev libboost-filesystem-dev nvidia-cuda-toolkit-gcc libzstd-dev
    - name: Install gpuocelot dependencies (MacOS)
      if: inputs.cuda == 'true' && runner.os == 'macOS'
      shell: bash
-      run: |
-        brew update
-        brew install cmake ninja llvm@15 zlib glew flex bison boost zstd ncurses
+      run: brew install --quiet cmake ninja llvm@15 zlib glew flex bison boost zstd ncurses
    - name: Cache gpuocelot
      if: inputs.cuda == 'true'
      id: cache-build
@@ -175,7 +217,7 @@ runs:
        git checkout b16039dc940dc6bc4ea0a98380495769ff35ed99
        mkdir build
        cd build
-        cmake .. -Wno-dev -G Ninja -DOCELOT_BUILD_TOOLS=OFF -DCMAKE_BUILD_ALWAYS=0 -DBUILD_TESTS_CUDA=OFF
+        cmake .. -Wno-dev -G Ninja -DOCELOT_BUILD_TOOLS=OFF -DCMAKE_BUILD_ALWAYS=0 -DBUILD_TESTS_CUDA=OFF -DCMAKE_POLICY_VERSION_MINIMUM=3.5
        ninja
    - name: Install gpuocelot
      if: inputs.cuda == 'true'
@@ -191,13 +233,7 @@ runs:
      shell: bash
      run: |
        sudo curl -L https://github.com/wpmed92/pydawn/releases/download/v0.1.6/libwebgpu_dawn.so -o /usr/local/lib/libwebgpu_dawn.so
-    - name: Install dependencies for software-based vulkan
-      if: inputs.webgpu == 'true' && runner.os == 'Linux'
-      shell: bash
-      run: |
-        sudo apt update -y || true
-        sudo apt install -y libegl1-mesa libgl1-mesa-dri libxcb-xfixes0-dev mesa-vulkan-drivers
-
+        sudo ldconfig
    - name: Install WebGPU dawn (macOS)
      if: inputs.webgpu == 'true' && runner.os == 'macOS'
      shell: bash
@@ -207,18 +243,7 @@ runs:

    # **** LLVM ****

-    - name: Install LLVM (Linux)
-      if: inputs.llvm == 'true' && runner.os == 'Linux'
-      shell: bash
-      run: |
-        echo 'Acquire::http::Pipeline-Depth "5";' | sudo tee -a /etc/apt/apt.conf.d/99parallel
-        wget -qO- https://apt.llvm.org/llvm-snapshot.gpg.key | sudo tee /etc/apt/trusted.gpg.d/apt.llvm.org.asc
-        echo "deb http://apt.llvm.org/$(lsb_release -cs)/ llvm-toolchain-$(lsb_release -cs)-19 main" | sudo tee /etc/apt/sources.list.d/llvm.list
-        sudo apt update -y || true
-        sudo apt install -y --no-install-recommends libllvm19 clang-19 lld-19
-
    - name: Install LLVM (macOS)
      if: inputs.llvm == 'true' && runner.os == 'macOS'
      shell: bash
-      run: |
-        brew install llvm
+      run: brew install llvm@20
@@ -24,11 +24,14 @@ on:
 jobs:
  testmacbenchmark:
    name: Mac Benchmark
+    env:
+      # since sudo is required for usbgpu on macos, move the cache to a new location, as some of the files are owned by root
+      PYTHONPYCACHEPREFIX: /tmp/tiny_python_pycache
    runs-on: [self-hosted, macOS]
    timeout-minutes: 20
    defaults:
      run:
-        shell: bash -o pipefail {0}
+        shell: bash -e -o pipefail {0}
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -49,61 +52,78 @@ jobs:
    - name: reset process replay
      run: python3.11 test/external/process_replay/reset.py
    - name: Run Stable Diffusion
-      run: JIT=1 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
+      run: BENCHMARK_LOG=stable_diffusion JIT=1 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
    - name: Run Stable Diffusion without fp16
-      run: JIT=1 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd_no_fp16.txt
+      run: BENCHMARK_LOG=stable_diffusion_fp32 JIT=1 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd_no_fp16.txt
    - name: Run Stable Diffusion v2
-      run: JIT=1 python3.11 examples/sdv2.py --fp16 --seed 0 --noshow --timing | tee sdv2.txt
+      run: BENCHMARK_LOG=stable_diffusion_v2 JIT=1 python3.11 examples/sdv2.py --fp16 --seed 0 --noshow --timing | tee sdv2.txt
+    # process replay can't capture this, the graph is too large
    - name: Run SDXL
-      run: JIT=1 python3.11 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
+      run: BENCHMARK_LOG=stable_diffusion_xl CAPTURE_PROCESS_REPLAY=0 JIT=1 python3.11 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
    - name: Run model inference benchmark
      run: METAL=1 python3.11 test/external/external_model_benchmark.py
+    - name: Run huggingface_onnx test
+      run: METAL=1 python3.11 extra/huggingface_onnx/run_models.py test --debug FacebookAI/xlm-roberta-large
    - name: Test speed vs torch
      run: BIG=2 MPS=1 python3.11 test/test_speed_v_torch.py | tee torch_speed.txt
    - name: Test tensor cores
-      run: METAL=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded
+      run: METAL=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_emulation TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
    - name: Test AMX tensor cores
      run: |
-        DEBUG=2 CPU=1 AMX=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded
-        DEBUG=2 LLVM=1 AMX=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded
+        DEBUG=2 CPU=1 AMX=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_emulation TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
+        DEBUG=2 LLVM=1 AMX=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_emulation TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
    - name: Run Tensor Core GEMM (float)
-      run: DEBUG=2 python3.11 extra/gemm/simple_matmul.py | tee matmul.txt
+      run: DEBUG=2 SHOULD_USE_TC=1 python3.11 extra/gemm/simple_matmul.py | tee matmul.txt
    - name: Run Tensor Core GEMM (half)
-      run: DEBUG=2 HALF=1 python3.11 extra/gemm/simple_matmul.py | tee matmul_half.txt
+      run: DEBUG=2 SHOULD_USE_TC=1 HALF=1 python3.11 extra/gemm/simple_matmul.py | tee matmul_half.txt
    - name: Run Tensor Core GEMM (bfloat16)
-      run: DEBUG=2 BFLOAT16=1 python3.11 extra/gemm/simple_matmul.py | tee matmul_bfloat16.txt
+      run: DEBUG=2 SHOULD_USE_TC=1 BFLOAT16=1 python3.11 extra/gemm/simple_matmul.py | tee matmul_bfloat16.txt
    - name: Fuzz Padded Tensor Core GEMM
      run: METAL=1 M_START=6 M_STOP=10 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=6 K_STOP=24 K_STEP=1 TC_OPT=2 DEBUG=2 python3.11 ./extra/gemm/fuzz_matmul.py
    - name: Run LLaMA
      run: |
-        JIT=0 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_unjitted.txt
-        JIT=1 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_jitted.txt
+        BENCHMARK_LOG=llama_nojit JIT=0 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_unjitted.txt
+        BENCHMARK_LOG=llama JIT=1 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_jitted.txt
    - name: Run LLaMA with BEAM
-      run: JITBEAM=2 IGNORE_BEAM_CACHE=1 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_beam.txt
+      run: BENCHMARK_LOG=llama_beam JITBEAM=2 IGNORE_BEAM_CACHE=1 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_beam.txt
    - name: Run quantized LLaMA
      run: |
-        python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing --quantize int8 | tee llama_int8.txt
-        python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing --quantize nf4 | tee llama_nf4.txt
+        BENCHMARK_LOG=llama_int8 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing --quantize int8 | tee llama_int8.txt
+        BENCHMARK_LOG=llama_nf4 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing --quantize nf4 | tee llama_nf4.txt
+    - name: Run quantized LLaMA3
+      run: |
+        BENCHMARK_LOG=llama3_int8 python3.11 examples/llama3.py --size 8B --temperature 0 --benchmark --quantize int8 | tee llama3_int8.txt
+        BENCHMARK_LOG=llama3_nf4 python3.11 examples/llama3.py --size 8B --temperature 0 --benchmark --quantize nf4 | tee llama3_nf4.txt
    #- name: Run LLaMA 7B on 4 (virtual) GPUs
    #  run: python3.11 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0  --timing | tee llama_four_gpu.txt
    - name: Run GPT2
      run: |
-        JIT=0 python3.11 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_unjitted.txt
-        JIT=1 python3.11 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
+        BENCHMARK_LOG=gpt2_nojit JIT=0 python3.11 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_unjitted.txt
+        BENCHMARK_LOG=gpt2 JIT=1 python3.11 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
    - name: Run GPT2 w HALF
-      run: HALF=1 python3.11 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
+      run: BENCHMARK_LOG=gpt2_half HALF=1 python3.11 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
    - name: Run GPT2 w HALF/BEAM
-      run: HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3.11 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half_beam.txt
+      run: BENCHMARK_LOG=gpt2_half_beam HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3.11 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half_beam.txt
+    - name: Run OLMoE
+      run: BENCHMARK_LOG=olmoe python3.11 examples/olmoe.py
    - name: Train MNIST
      run: time PYTHONPATH=. TARGET_EVAL_ACC_PCT=96.0 python3.11 examples/beautiful_mnist.py | tee beautiful_mnist.txt
    - name: Run 10 CIFAR training steps
-      run: JIT=1 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar.txt
+      run: BENCHMARK_LOG=cifar_10steps JIT=1 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar.txt
    - name: Run 10 CIFAR training steps w HALF
-      run: JIT=2 STEPS=10 DEFAULT_FLOAT=HALF python3.11 examples/hlb_cifar10.py | tee train_cifar_half.txt
+      run: BENCHMARK_LOG=cifar_10steps_half JIT=2 STEPS=10 DEFAULT_FLOAT=HALF python3.11 examples/hlb_cifar10.py | tee train_cifar_half.txt
    #- name: Run 10 CIFAR training steps w BF16
    #  run: STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3.11 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
    - name: Run 10 CIFAR training steps w winograd
-      run: JIT=1 WINO=1 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar_wino.txt
+      run: BENCHMARK_LOG=cifar_10steps_wino JIT=1 WINO=1 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar_wino.txt
+    - name: UsbGPU boot time
+      run: sudo -E PYTHONPATH=. DEBUG=2 AM_RESET=1 AMD=1 AMD_IFACE=USB time python3.11 test/test_tiny.py TestTiny.test_plus
+    - name: UsbGPU tiny tests
+      run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB python3.11 test/test_tiny.py
+    - name: UsbGPU copy speeds
+      run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
+    - name: UsbGPU openpilot test
+      run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB AMD_LLVM=1 NOLOCALS=0 IMAGE=0 GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (Mac)
@@ -115,6 +135,8 @@ jobs:
          llama_beam.txt
          llama_int8.txt
          llama_nf4.txt
+          llama3_int8.txt
+          llama3_nf4.txt
          llama_four_gpu.txt
          gpt2_unjitted.txt
          gpt2_jitted.txt
@@ -141,7 +163,7 @@ jobs:
    timeout-minutes: 30
    defaults:
      run:
-        shell: bash -o pipefail {0}
+        shell: bash -e -o pipefail {0}
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -174,53 +196,53 @@ jobs:
      run: NV=1 python test/external/external_benchmark_multitensor_allreduce.py
    - name: Test tensor cores
      run: |
-        NV=1 ALLOW_TF32=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded
-        PTX=1 ALLOW_TF32=1 NV=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded
+        NV=1 ALLOW_TF32=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_emulation TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
+        PTX=1 ALLOW_TF32=1 NV=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_emulation TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
    - name: Run Tensor Core GEMM (CUDA)
      run: |
-        CUDA=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul.txt
-        CUDA=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_bfloat16.txt
-        CUDA=1 ALLOW_TF32=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_tf32.txt
+        CUDA=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul.txt
+        CUDA=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_bfloat16.txt
+        CUDA=1 SHOULD_USE_TC=1 ALLOW_TF32=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py | tee matmul_tf32.txt
    - name: Run Tensor Core GEMM (PTX)
-      run: NV=1 PTX=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_ptx.txt
+      run: NV=1 PTX=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_ptx.txt
    - name: Run Tensor Core GEMM (NV)
-      run: NV=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_nv.txt
+      run: NV=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_nv.txt
    - name: Test NV=1
      run: DEBUG=2 NV=1 python -m pytest -rA test/test_tiny.py
    - name: Test CUDA=1
      run: DEBUG=2 CUDA=1 python -m pytest -rA test/test_tiny.py
    - name: Run Stable Diffusion
-      run: NV=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
+      run: BENCHMARK_LOG=stable_diffusion NV=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
    - name: Run SDXL
-      run: NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
+      run: BENCHMARK_LOG=stable_diffusion_xl CAPTURE_PROCESS_REPLAY=0 NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
    - name: Run LLaMA
      run: |
-        NV=1 JIT=0 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_unjitted.txt
-        NV=1 JIT=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_jitted.txt
+        BENCHMARK_LOG=llama_nojit NV=1 JIT=0 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_unjitted.txt
+        BENCHMARK_LOG=llama NV=1 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: NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_beam.txt
+      run: BENCHMARK_LOG=llama_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_beam.txt
    # - name: Run LLaMA 7B on 4 GPUs
    #   run: NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0  --timing | tee llama_four_gpu.txt
    # - name: Run LLaMA 7B on 6 GPUs
    #   run: NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing | tee llama_six_gpu.txt
    - name: Run LLaMA-3 8B BEAM
-      run: NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_beam.txt
+      run: BENCHMARK_LOG=llama3_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_beam.txt
    - name: Run LLaMA-3 8B on 4 GPUs with BEAM
-      run: NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_four_gpu.txt
+      run: BENCHMARK_LOG=llama3_beam_4gpu NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_four_gpu.txt
    # - name: Run LLaMA-3 8B on 6 GPUs
    #   run: NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 6 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_six_gpu.txt
    # - name: Run LLaMA-2 70B
    #   run: NV=1 CAPTURE_PROCESS_REPLAY=0 MAX_CONTEXT=256 python3 examples/llama.py --gen 2 --size 70B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing | tee llama_2_70B.txt
    - name: Run Mixtral 8x7B
-      run: time NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/mixtral.py --temperature 0 --count 10 --timing | tee mixtral.txt
+      run: time BENCHMARK_LOG=mixtral NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/mixtral.py --temperature 0 --count 10 --timing | tee mixtral.txt
    - name: Run GPT2
      run: |
-        NV=1 JIT=0 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_unjitted.txt
-        NV=1 JIT=1 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
+        BENCHMARK_LOG=gpt2_nojit NV=1 JIT=0 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_unjitted.txt
+        BENCHMARK_LOG=gpt2 NV=1 JIT=1 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
    - name: Run GPT2 w HALF
-      run: NV=1 HALF=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
+      run: BENCHMARK_LOG=gpt2_half NV=1 HALF=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
    - name: Run GPT2 w HALF/BEAM
-      run: NV=1 HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half_beam.txt
+      run: BENCHMARK_LOG=gpt2_half_beam NV=1 HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half_beam.txt
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (NVIDIA)
@@ -255,7 +277,7 @@ jobs:
    timeout-minutes: 20
    defaults:
      run:
-        shell: bash -o pipefail {0}
+        shell: bash -e -o pipefail {0}
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -284,26 +306,26 @@ jobs:
    - name: Train MNIST
      run: time PYTHONPATH=. NV=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py | tee beautiful_mnist.txt
    - name: Run 10 CIFAR training steps
-      run: NV=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
+      run: BENCHMARK_LOG=cifar_10steps NV=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
    - name: Run 10 CIFAR training steps w HALF
-      run: NV=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
+      run: BENCHMARK_LOG=cifar_10steps_half NV=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
    - name: Run 10 CIFAR training steps w BF16
-      run: NV=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
+      run: BENCHMARK_LOG=cifar_10steps_bf16 NV=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
    - name: Run 10 CIFAR training steps w winograd
-      run: NV=1 CAPTURE_PROCESS_REPLAY=0 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
+      run: BENCHMARK_LOG=cifar_10steps_half_wino NV=1 CAPTURE_PROCESS_REPLAY=0 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
    - name: Run full CIFAR training w 1 GPU
-      run: time NV=1 DEFAULT_FLOAT=HALF LATEWINO=1 STEPS=1000 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
+      run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF LATEWINO=1 STEPS=1000 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
    - name: Run full CIFAR training steps w 6 GPUS
-      run: time CAPTURE_PROCESS_REPLAY=0 NV=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
+      run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 NV=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
    - name: Run MLPerf resnet eval on training data
-      run: time NV=1 MODEL=resnet python3 examples/mlperf/model_eval.py
+      run: time BENCHMARK_LOG=resnet_eval NV=1 MODEL=resnet python3 examples/mlperf/model_eval.py
    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
-      run: NV=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet_one_gpu.txt
+      run: BENCHMARK_LOG=resnet_10steps NV=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet_one_gpu.txt
    - name: Run 10 MLPerf ResNet50 training steps (6 gpu)
-      run: NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
+      run: BENCHMARK_LOG=resnet_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
    - name: Run 10 MLPerf Bert training steps (6 gpu)
      # TODO: remove BERT_LAYERS once scheduler is fast
-      run: NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
+      run: BENCHMARK_LOG=bert_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=6 BERT_LAYERS=2 FUSE_ARANGE=1 FUSE_ARANGE_UINT=0 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (NVIDIA Training)
@@ -327,13 +349,17 @@ jobs:
    timeout-minutes: 20
    defaults:
      run:
-        shell: bash -o pipefail {0}
+        shell: bash -e -o pipefail {0}
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
-    - name: Insert amdgpu
-      run: sudo modprobe amdgpu
+    - name: Remove amdgpu
+      run: sudo rmmod amdgpu || true
+    - name: Cleanup running AM processes
+      run: python extra/amdpci/am_smi.py --pids --kill
+    #- name: Insert amdgpu
+    #  run: sudo modprobe amdgpu
    - name: Symlink models and datasets
      run: |
        mkdir -p weights
@@ -352,12 +378,12 @@ jobs:
        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
    - name: reset process replay
      run: test/external/process_replay/reset.py
-    - name: setup perflevel
-      run: |
-        examples/mlperf/training_submission_v4.1/tinycorp/benchmarks/bert/implementations/tinybox_red/setup.sh
-        rocm-smi
-    - name: Show off tinybox
-      run: /opt/rocm/bin/rocm-bandwidth-test
+    #- name: setup perflevel
+    #  run: |
+    #    examples/mlperf/training_submission_v4.1/tinycorp/benchmarks/bert/implementations/tinybox_red/setup.sh
+    #    rocm-smi
+    #- name: Show off tinybox
+    #  run: /opt/rocm/bin/rocm-bandwidth-test
    # TODO: unstable on AMD
    #- name: Run model inference benchmark
    #  run: LD_PRELOAD="/opt/rocm/lib/libhsa-runtime64.so" HSA=1 NOCLANG=1 python3 test/external/external_model_benchmark.py
@@ -369,56 +395,59 @@ jobs:
    - name: Test speed vs theoretical
      run: AMD=1 IGNORE_BEAM_CACHE=1 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
    - name: Test tensor cores
-      run: AMD=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded
+      run: |
+        AMD=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_emulation TestLinearizer.test_tensor_cores_padded_amd TestLinearizer.test_tensor_cores_padded_uops
+        AMD=1 AMD_LLVM=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_emulation TestLinearizer.test_tensor_cores_padded_amd TestLinearizer.test_tensor_cores_padded_uops
+        AMD=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
    - name: Run Tensor Core GEMM (AMD)
-      run: AMD=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_amd.txt
+      run: AMD=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py | tee matmul_amd.txt
    - name: Test AMD=1
      run: DEBUG=2 AMD=1 python -m pytest -rA test/test_tiny.py
-    - name: Test HIP=1
-      run: DEBUG=2 HIP=1 python -m pytest -rA test/test_tiny.py
+    #- name: Test HIP=1
+    #  run: DEBUG=2 HIP=1 python -m pytest -rA test/test_tiny.py
    # TODO: AMD compiler bug causes this to fail
    #- name: Fuzz Padded Tensor Core GEMM
    #  run: HSA=1 M_START=12 M_STOP=20 M_STEP=1 N_START=12 N_STOP=20 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 DEBUG=2 python3 ./extra/gemm/fuzz_matmul.py
-    - name: Remove amdgpu
-      run: sleep 5 && sudo rmmod amdgpu # sleep a bit to let the driver unload the prev pid.
+    #- name: Remove amdgpu
+    #  run: sleep 10 && sudo rmmod amdgpu # sleep a bit to let the driver unload the prev pid.
    - name: Test AM cold start time
      run: time AMD=1 AM_RESET=1 python3 test/test_tiny.py TestTiny.test_plus
    - name: Test AM warm start time
      run: time AMD=1 python3 test/test_tiny.py TestTiny.test_plus
    - name: Run Stable Diffusion
-      run: AMD=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
+      run: BENCHMARK_LOG=stable_diffusion AMD=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
    - name: Run SDXL
-      run: AMD=1 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
+      run: BENCHMARK_LOG=stable_diffusion_xl CAPTURE_PROCESS_REPLAY=0 AMD=1 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
    - name: Run LLaMA 7B
      run: |
-        AMD=1 JIT=0 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_unjitted.txt
-        AMD=1 JIT=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_jitted.txt
+        BENCHMARK_LOG=llama_nojit AMD=1 JIT=0 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_unjitted.txt
+        BENCHMARK_LOG=llama AMD=1 JIT=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_jitted.txt
    - name: Run LLaMA 7B with BEAM
-      run: AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_beam.txt
+      run: BENCHMARK_LOG=llama_beam AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_beam.txt
    # - name: Run LLaMA 7B on 4 GPUs
    #   run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0  --timing | tee llama_four_gpu.txt
    # - name: Run LLaMA 7B on 6 GPUs
    #   run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing | tee llama_six_gpu.txt
    - name: Run LLaMA-3 8B BEAM
-      run: AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_beam.txt
+      run: BENCHMARK_LOG=llama3_beam AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_beam.txt
    - name: Run LLaMA-3 8B on 4 GPUs with BEAM
-      run: AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_four_gpu.txt
+      run: BENCHMARK_LOG=llama3_beam_4gpu AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_four_gpu.txt
    # - name: Run LLaMA-3 8B on 6 GPUs
    #   run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 6 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_six_gpu.txt
-    - name: Restore amdgpu
-      run: sudo modprobe amdgpu
+    #- name: Restore amdgpu
+    #  run: sudo modprobe amdgpu
    # - name: Run LLaMA-2 70B
    #   run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 2 --size 70B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing | tee llama_2_70B.txt
    - name: Run Mixtral 8x7B
-      run: time AMD=1 python3 examples/mixtral.py --temperature 0 --count 10 --timing | tee mixtral.txt
+      run: time BENCHMARK_LOG=mixtral AMD=1 python3 examples/mixtral.py --temperature 0 --count 10 --timing | tee mixtral.txt
    - name: Run GPT2
      run: |
-        AMD=1 JIT=0 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_unjitted.txt
-        AMD=1 JIT=1 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
+        BENCHMARK_LOG=gpt2_nojit AMD=1 JIT=0 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_unjitted.txt
+        BENCHMARK_LOG=gpt2 AMD=1 JIT=1 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
    - name: Run GPT2 w HALF
-      run: AMD=1 HALF=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
+      run: BENCHMARK_LOG=gpt2_half AMD=1 HALF=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
    - name: Run GPT2 w HALF/BEAM
-      run: AMD=1 HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half_beam.txt
+      run: BENCHMARK_LOG=gpt2_half_beam AMD=1 HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half_beam.txt
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (AMD)
@@ -450,13 +479,15 @@ jobs:
    timeout-minutes: 30
    defaults:
      run:
-        shell: bash -o pipefail {0}
+        shell: bash -e -o pipefail {0}
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
    - name: Remove amdgpu
      run: sudo rmmod amdgpu || true
+    - name: Cleanup running AM processes
+      run: python extra/amdpci/am_smi.py --pids --kill
    - name: Symlink models and datasets
      run: |
        mkdir -p weights
@@ -477,26 +508,19 @@ jobs:
    - name: Train MNIST
      run: time PYTHONPATH=. AMD=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py | tee beautiful_mnist.txt
    - name: Run 10 CIFAR training steps
-      run: AMD=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
+      run: BENCHMARK_LOG=cifar_10steps AMD=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
    - name: Run 10 CIFAR training steps w HALF
-      run: AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
+      run: BENCHMARK_LOG=cifar_10steps_half AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
    - name: Run 10 CIFAR training steps w BF16
-      run: AMD=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
+      run: BENCHMARK_LOG=cifar_10steps_bf16 AMD=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
    - name: Run 10 CIFAR training steps w winograd
-      run: AMD=1 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
+      run: BENCHMARK_LOG=cifar_10steps_half_wino AMD=1 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
    - name: Run full CIFAR training w 1 GPU
-      run: time AMD=1 DEFAULT_FLOAT=HALF LATEWINO=1 STEPS=1000 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
+      run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF LATEWINO=1 STEPS=1000 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
    - name: Run full CIFAR training steps w 6 GPUS
-      run: time AMD=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
-    - name: Run MLPerf resnet eval
-      run: time AMD=1 MODEL=resnet python3 examples/mlperf/model_eval.py
-    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
-      run: AMD=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet_one_gpu.txt
-    - name: Run 10 MLPerf ResNet50 training steps (6 gpu)
-      run: AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
-    - name: Run 10 MLPerf Bert training steps (6 gpu)
-      # TODO: remove BERT_LAYERS once scheduler is fast
-      run: AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
+      run: time BENCHMARK_LOG=cifar_6gpu AMD=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
+    - name: Run full CIFAR training steps w 6 GPUS (REMOTE)
+      run: time BENCHMARK_LOG=cifar_6gpu_remote REMOTE=1 REMOTEDEV=AMD DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu_remote.txt
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (AMD Training)
@@ -508,6 +532,55 @@ jobs:
          train_cifar_wino.txt
          train_cifar_one_gpu.txt
          train_cifar_six_gpu.txt
+          train_cifar_six_gpu_remote.txt
+    - name: Run process replay tests
+      run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
+
+  testmlperfamdbenchmark:
+    name: tinybox red MLPerf Benchmark
+    runs-on: [self-hosted, Linux, tinybox]
+    timeout-minutes: 30
+    defaults:
+      run:
+        shell: bash -e -o pipefail {0}
+    if: github.repository_owner == 'tinygrad'
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v4
+    - name: Remove amdgpu
+      run: sudo rmmod amdgpu || true
+    - name: Cleanup running AM processes
+      run: python extra/amdpci/am_smi.py --pids --kill
+    - name: Symlink models and datasets
+      run: |
+        mkdir -p weights
+        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
+        ln -s /raid/weights/mixtral-8x7b-32kseqlen weights/mixtral-8x7b-32kseqlen
+        ln -s /raid/weights/LLaMA-2 weights/LLaMA-2
+        mkdir -p extra/datasets
+        ln -s /raid/datasets/imagenet extra/datasets/imagenet
+    - name: setup staging db
+      if: github.ref == 'refs/heads/update_benchmark_staging'
+      run: |
+        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
+        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
+    - name: reset process replay
+      run: test/external/process_replay/reset.py
+    - name: Run MLPerf resnet eval
+      run: time BENCHMARK_LOG=resnet_eval AMD=1 MODEL=resnet python3 examples/mlperf/model_eval.py
+    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
+      run: BENCHMARK_LOG=resnet_10steps AMD=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet_one_gpu.txt
+    - name: Run 10 MLPerf ResNet50 training steps (6 gpu)
+      run: BENCHMARK_LOG=resnet_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
+    - name: Run 10 MLPerf Bert training steps (6 gpu)
+      # TODO: remove BERT_LAYERS once scheduler is fast
+      run: BENCHMARK_LOG=bert_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=6 BERT_LAYERS=2 FUSE_ARANGE=1 FUSE_ARANGE_UINT=0 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
+    - uses: actions/upload-artifact@v4
+      with:
+        name: Speed (AMD MLPerf)
+        path: |
          train_resnet.txt
          train_resnet_one_gpu.txt
          train_bert.txt
@@ -520,7 +593,7 @@ jobs:
    timeout-minutes: 20
    defaults:
      run:
-        shell: bash -o pipefail {0}
+        shell: bash -e -o pipefail {0}
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -535,19 +608,27 @@ jobs:
    - name: validate openpilot 0.9.7
      run: PYTHONPATH=. FLOAT16=0 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx | tee openpilot_image_0_9_7.txt
    - name: benchmark openpilot 0.9.4
-      run: PYTHONPATH=. QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx | tee openpilot_0_9_4.txt
+      run: BENCHMARK_LOG=openpilot_0_9_4 PYTHONPATH=. QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx | tee openpilot_0_9_4.txt
    - name: benchmark openpilot 0.9.7
-      run: PYTHONPATH=. QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx | tee openpilot_0_9_7.txt
+      run: BENCHMARK_LOG=openpilot_0_9_7 PYTHONPATH=. QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx | tee openpilot_0_9_7.txt
    - name: benchmark openpilot w IMAGE=2 0.9.4
-      run: PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx | tee openpilot_image_0_9_4.txt
+      run: BENCHMARK_LOG=openpilot_0_9_4_image PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx | tee openpilot_image_0_9_4.txt
    - name: benchmark openpilot w IMAGE=2 0.9.7
-      run: PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx | tee openpilot_image_0_9_7.txt
+      run: BENCHMARK_LOG=openpilot_0_9_7_image PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx | tee openpilot_image_0_9_7.txt
    - name: openpilot compile3 0.9.7
      run: PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx
    - name: openpilot compile3 0.9.7+ tomb raider
      run: PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/e8bea2c78ffa92685ece511e9b554122aaf1a79d/selfdrive/modeld/models/supercombo.onnx
    - name: openpilot dmonitoring compile3 0.9.7
      run: PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/dmonitoring_model.onnx
+    - name: benchmark MobileNetV2 on DSP
+      run: |
+        # generate quantized weights
+        ln -s /data/home/tiny/tinygrad/extra/datasets/imagenet extra/datasets/imagenet
+        ln -s /data/home/tiny/tinygrad/testsig-*.so .
+        PYTHONPATH=. CC=clang-19 CPU=1 QUANT=1 CNT=0 python3 examples/test_onnx_imagenet.py https://github.com/xamcat/mobcat-samples/raw/refs/heads/master/onnx_runtime/InferencingSample/InferencingSample/mobilenetv2-7.onnx /tmp/model.quant.onnx
+        # benchmark on DSP with NOOPT=1, the devectorizer has issues
+        PYTHONPATH=. CC=clang-19 DSP=1 DONT_REALIZE_EXPAND=1 NOOPT=1 CNT=2 DEBUG=2 python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
    - name: Run process replay tests
      run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
    - uses: actions/upload-artifact@v4
@@ -0,0 +1,33 @@
+name: Benchmark with kernel search
+
+on:
+  push:
+    branches:
+      - update_benchmark_search
+  workflow_dispatch:
+
+jobs:
+  run_script_job:
+    runs-on: [self-hosted, Linux, tinybox]
+    if: github.repository_owner == 'tinygrad'
+
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v4
+    - name: Remove amdgpu
+      run: sudo rmmod amdgpu || true
+    - name: Cleanup running AM processes
+      run: python extra/amdpci/am_smi.py --pids --kill
+    - name: Run SDXL with new search
+      # TODO: GCVM_L2_PROTECTION_FAULT_STATUS with llvm19
+      run: |
+        BENCHMARK_LOG=search_sdxl PYTHONPATH=. AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 DISABLE_COMPILER_CACHE=1 python examples/sdxl.py --noshow --timing --seed 0
+    - name: Run SDXL with cached search
+      run: |
+        BENCHMARK_LOG=search_sdxl_cached PYTHONPATH=. AMD=1 JITBEAM=2 python examples/sdxl.py --noshow --timing --seed 0
+    - name: Run winograd cifar with new search
+      run: |
+        BENCHMARK_LOG=search_wino_cifar WINO=1 DEFAULT_FLOAT=HALF FUSE_ARANGE=1 JITBEAM=4 IGNORE_BEAM_CACHE=1 DISABLE_COMPILER_CACHE=1 BS=1024 STEPS=500 python examples/hlb_cifar10.py
+    - name: Run winograd cifar with cached search
+      run: |
+        BENCHMARK_LOG=search_wino_cifar_cached WINO=1 DEFAULT_FLOAT=HALF FUSE_ARANGE=1 JITBEAM=4 BS=1024 STEPS=500 python examples/hlb_cifar10.py
@@ -0,0 +1,29 @@
+name: Run MLPerf Training
+
+on:
+  schedule:
+    - cron: '5 8 * * *'  # Runs at 08:05 UTC (12:05 AM Pacific Time)
+  push:
+    branches:
+      - update_mlperf
+  workflow_dispatch:
+
+jobs:
+  run_script_job:
+    runs-on: [self-hosted, Linux, tinybox]
+    if: github.repository_owner == 'tinygrad'
+
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v4
+    - name: Cleanup running AM processes
+      run: python extra/amdpci/am_smi.py --pids --kill
+    - name: Symlink datasets
+      run: |
+        mkdir -p extra/datasets
+        ln -s /raid/datasets/imagenet extra/datasets/imagenet
+    - name: Run resnet
+      run: |
+        rm "~/.cache/tinygrad/cache_mlperf.db" || true
+        BENCHMARK_LOG=mlpert_train_resnet LOGMLPERF=0 CACHEDB="~/.cache/tinygrad/cache_mlperf.db" examples/mlperf/training_submission_v5.1/tinycorp/benchmarks/resnet/implementations/tinybox_red/run_and_time.sh
+        rm "~/.cache/tinygrad/cache_mlperf.db"
@@ -1,7 +1,8 @@
 name: Unit Tests
 env:
  # increment this when downloads substantially change to avoid the internet
-  DOWNLOAD_CACHE_VERSION: '9'
+  DOWNLOAD_CACHE_VERSION: '10'
+  PYTHON_CACHE_VERSION: '2'
  CAPTURE_PROCESS_REPLAY: 1
  GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}

@@ -29,9 +30,9 @@ jobs:
    - name: External Benchmark Schedule
      run: PYTHONPATH="." python3 test/external/external_benchmark_schedule.py
    - name: Speed Test
-      run: LLVM=1 LLVMOPT=1 python3 test/test_speed_v_torch.py
+      run: LLVM=1 python3 test/test_speed_v_torch.py
    - name: Speed Test (BEAM=2)
-      run: BEAM=2 LLVM=1 LLVMOPT=1 python3 test/test_speed_v_torch.py
+      run: BEAM=2 LLVM=1 python3 test/test_speed_v_torch.py

  docs:
    name: Docs
@@ -53,11 +54,6 @@ jobs:
        source venv/bin/activate
        pip install $GITHUB_WORKSPACE
        python -c "from tinygrad.tensor import Tensor; print(Tensor([1,2,3,4,5]))"
-        # Test using VIZ=1 after package installation
-        VIZ=1 python -c "from tinygrad.tensor import Tensor; Tensor([1,2,3,4,5]).realize()" & VIZ_PID=$!
-        echo "started VIZ server at: $(ps -p "$VIZ_PID" -o pid=)"
-        i=0; while ((i++ < 10)); do curl -sSf localhost:8000 > /dev/null && break || { echo "VIZ verification attempt $i/10"; sleep 1; }; done; ((i > 10)) && echo "Could not verify VIZ server" && exit 1
-        kill $VIZ_PID
        pip install mypy
        mypy -c "from tinygrad.tensor import Tensor; print(Tensor([1,2,3,4,5]))"
    - name: Run beautiful_mnist with tinygrad only
@@ -87,8 +83,8 @@ jobs:

  autogen:
    name: Autogen
-    runs-on: ubuntu-22.04
-    timeout-minutes: 10
+    runs-on: ubuntu-24.04
+    timeout-minutes: 15
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
@@ -100,6 +96,8 @@ jobs:
        cuda: 'true'
        webgpu: 'true'
        llvm: 'true'
+    - name: Install autogen support packages
+      run: sudo apt-get install -y --no-install-recommends llvm-14-dev libclang-14-dev
    - name: Verify OpenCL autogen
      run: |
        cp tinygrad/runtime/autogen/opencl.py /tmp/opencl.py.bak
@@ -116,14 +114,17 @@ jobs:
    - name: Verify AMD autogen
      run: |
        cp tinygrad/runtime/autogen/hsa.py /tmp/hsa.py.bak
+        cp tinygrad/runtime/autogen/kfd.py /tmp/kfd.py.bak
        cp tinygrad/runtime/autogen/comgr.py /tmp/comgr.py.bak
        cp tinygrad/runtime/autogen/amd_gpu.py /tmp/amd_gpu.py.bak
        cp tinygrad/runtime/autogen/sqtt.py /tmp/sqtt.py.bak
        ./autogen_stubs.sh hsa
+        ./autogen_stubs.sh kfd
        ./autogen_stubs.sh comgr
        ./autogen_stubs.sh amd
        ./autogen_stubs.sh sqtt
        diff /tmp/hsa.py.bak tinygrad/runtime/autogen/hsa.py
+        diff /tmp/kfd.py.bak tinygrad/runtime/autogen/kfd.py
        diff /tmp/comgr.py.bak tinygrad/runtime/autogen/comgr.py
        diff /tmp/amd_gpu.py.bak tinygrad/runtime/autogen/amd_gpu.py
        diff /tmp/sqtt.py.bak tinygrad/runtime/autogen/sqtt.py
@@ -149,7 +150,9 @@ jobs:
  torchbackend:
    name: Torch Backend Tests
    runs-on: ubuntu-latest
-    timeout-minutes: 10
+    timeout-minutes: 15
+    env:
+      IGNORE_OOB: 0
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
@@ -176,8 +179,8 @@ jobs:
      run: PYTHONPATH=. python3 extra/torch_backend/test.py
    - name: Test one op in torch tests
      run: PYTHONPATH=. DEBUG=2 python3 extra/torch_backend/torch_tests.py TestTinyBackendPRIVATEUSE1.test_unary_log_tiny_float32
-    - name: Test Ops with TINY_BACKEND (expect failure)
-      run: PYTHONPATH=. LLVM=1 LLVMOPT=0 TINY_BACKEND=1 python3 -m pytest -n auto test/test_ops.py --durations=20 || true
+    - name: Test Ops with TINY_BACKEND
+      run: PYTHONPATH=. LLVM=1 LLVMOPT=0 TINY_BACKEND=1 python3 -m pytest -n auto test/test_ops.py --durations=20
    - name: Test in-place operations on views
      run: PYTHONPATH=. TORCH_DEBUG=1 python3 extra/torch_backend/test_inplace.py
    - name: Test multi-gpu
@@ -186,7 +189,9 @@ jobs:
  torchbackendmore:
    name: Torch Backend Tests More
    runs-on: ubuntu-latest
-    timeout-minutes: 10
+    timeout-minutes: 15
+    env:
+      IGNORE_OOB: 0
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
@@ -201,7 +206,7 @@ jobs:
        sudo apt update || true
        sudo apt install -y --no-install-recommends ninja-build
    - name: Test beautiful_mnist in torch with TINY_BACKEND
-      run: PYTHONPATH=. LLVM=1 TARGET_EVAL_ACC_PCT=96.0 TINY_BACKEND=1 python3 examples/other_mnist/beautiful_mnist_torch.py
+      run: SPLIT_REDUCEOP=0 FUSE_ARANGE=1 PYTHONPATH=. LLVM=1 TARGET_EVAL_ACC_PCT=96.0 TINY_BACKEND=1 python3 examples/other_mnist/beautiful_mnist_torch.py
    - name: Test some torch tests (expect failure)
      run: PYTHONPATH=. python3 -m pytest extra/torch_backend/torch_tests.py -v --tb=no || true

@@ -209,6 +214,8 @@ jobs:
    name: Tensor Core tests
    runs-on: ubuntu-latest
    timeout-minutes: 10
+    env:
+      IGNORE_OOB: 0
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
@@ -224,7 +231,8 @@ jobs:
    - name: Test emulated METAL tensor cores
      run: |
        DEBUG=2 EMULATE_METAL=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_big_gemm
-        PYTHONPATH=. DEBUG=2 EMULATE_METAL=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded
+        PYTHONPATH=. DEBUG=2 EMULATE_METAL=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores
+        PYTHONPATH=. DEBUG=2 EMULATE_METAL=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
    - name: Test emulated AMX tensor cores
      run: PYTHONPATH=. DEBUG=2 AMX=1 EMULATE_AMX=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm
    - name: Test emulated AMD tensor cores
@@ -233,17 +241,28 @@ jobs:
        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=1 python3 ./extra/gemm/simple_matmul.py
        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=1 python3 ./extra/gemm/simple_matmul.py
-        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded
+        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores
+        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores_padded_amd TestLinearizer.test_tensor_cores_padded_uops
    - name: Test emulated AMD MFMA tensor cores
      run: |
        PYTHONPATH=. DEBUG=2 EMULATE_AMD_MFMA=1 FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
-        PYTHONPATH=. DEBUG=2 EMULATE_AMD_MFMA=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded
+        PYTHONPATH=. DEBUG=2 EMULATE_AMD_MFMA=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores
+        PYTHONPATH=. DEBUG=2 EMULATE_AMD_MFMA=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
+    - name: Test emulated AMD RDNA4 tensor cores
+      run: |
+        PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
+        PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
+        PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=1 python3 ./extra/gemm/simple_matmul.py
+        PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=1 python3 ./extra/gemm/simple_matmul.py
+        PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores
+        PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
    - name: Test emulated CUDA tensor cores
      run: |
        DEBUG=2 EMULATE_CUDA=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm_fp16
        DEBUG=2 EMULATE_CUDA=1 ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm
        DEBUG=2 EMULATE_CUDA_SM75=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm_fp16
-        PYTHONPATH="." DEBUG=2 EMULATE_CUDA=1 ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded
+        PYTHONPATH="." DEBUG=2 EMULATE_CUDA=1 ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores
+        PYTHONPATH="." DEBUG=2 EMULATE_CUDA=1 ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
    - name: Test emulated INTEL OpenCL tensor cores
      run: DEBUG=2 EMULATE_INTEL=1 FORWARD_ONLY=1 PYTHON=1 HALF=1 N=64 python3 ./extra/gemm/simple_matmul.py
    - name: Full test tensor cores
@@ -255,12 +274,12 @@ jobs:
        PYTHONPATH=. DEBUG=2 AMX=1 EMULATE_AMX=1 FORWARD_ONLY=1 PYTHON=1 python3 ./test/test_linearizer.py TestLinearizer.test_tensor_cores
    - name: Test tensor cores (TC=3)
      run: |
-        TC=3 DEBUG=3 EMULATE_METAL=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm
-        TC=3 PYTHONPATH=. DEBUG=3 EMULATE_AMD=1 PYTHON=1 N=16 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
-        TC=3 PYTHONPATH=. DEBUG=3 EMULATE_AMD_MFMA=1 PYTHON=1 N=16 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
-        TC=3 DEBUG=3 EMULATE_CUDA=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm_fp16
-        TC=3 PYTHONPATH=. DEBUG=3 EMULATE_INTEL=1 PYTHON=1 N=16 HALF=1 python3 ./extra/gemm/simple_matmul.py
-        TC=3 PYTHONPATH=. DEBUG=3 AMX=1 EMULATE_AMX=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm
+        PYTHONPATH=. DEBUG=2 PYTHON=1 EMULATE_METAL=1 python3 ./test/test_linearizer.py TestLinearizer.test_tensor_cores_emulation
+        PYTHONPATH=. DEBUG=2 PYTHON=1 EMULATE_AMD=1 python3 ./test/test_linearizer.py TestLinearizer.test_tensor_cores_emulation
+        PYTHONPATH=. DEBUG=2 PYTHON=1 EMULATE_AMD_MFMA=1 python3 ./test/test_linearizer.py TestLinearizer.test_tensor_cores_emulation
+        PYTHONPATH=. DEBUG=2 PYTHON=1 EMULATE_CUDA=1 python3 ./test/test_linearizer.py TestLinearizer.test_tensor_cores_emulation
+        PYTHONPATH=. DEBUG=2 PYTHON=1 EMULATE_INTEL=1 python3 ./test/test_linearizer.py TestLinearizer.test_tensor_cores_emulation
+        PYTHONPATH=. DEBUG=2 PYTHON=1 EMULATE_AMX=1 AMX=1 python3 ./test/test_linearizer.py TestLinearizer.test_tensor_cores_emulation
    - name: Test device flop counts
      run: |
        PYTHONPATH=. DEBUG=2 EMULATE_METAL=1 PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
@@ -273,6 +292,8 @@ jobs:
    name: Python Backend
    runs-on: ubuntu-latest
    timeout-minutes: 10
+    env:
+      IGNORE_OOB: 0
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
@@ -291,6 +312,8 @@ jobs:
      run: PYTHONPATH=. PYTHON=1 python3 test/test_symbolic_ops.py
    - name: test_linearizer_failures with Python emulator
      run: PYTHONPATH=. PYTHON=1 python3 -m pytest -rA test/test_linearizer_failures.py::TestLinearizerFailures::test_failure_1
+    - name: test_renderer_failures with Python emulator
+      run: PYTHONPATH=. PYTHON=1 python3 -m pytest -rA test/test_renderer_failures.py::TestRendererFailures

  linter:
    name: Linters
@@ -313,6 +336,7 @@ jobs:
      run: |
        pip3 install --upgrade --force-reinstall ruff==0.11.0
        python3 -m ruff check .
+        python3 -m ruff check examples/mlperf/model_train.py --ignore E501
    - name: Lint tinygrad with pylint
      run: python -m pylint tinygrad/
    - name: Run mypy
@@ -330,13 +354,25 @@ jobs:
      uses: ./.github/actions/setup-tinygrad
      with:
        key: unittest-12
+        pydeps: "pillow"
        deps: testing_unit
    - name: Test README
      run: awk '/```python/{flag=1;next}/```/{flag=0}flag' README.md > README.py &&  PYTHONPATH=. python README.py
    - name: Run unit tests
-      run: PYTHONPATH="." python -m pytest -n=auto test/unit/
-    - name: Repo line count < 12000 lines
-      run: MAX_LINE_COUNT=12000 python sz.py
+      run: PYTHONPATH="." python -m pytest -n=auto test/unit/ --durations=20
+    - name: Run targetted tests on NULL backend
+      run: PYTHONPATH="." NULL=1 python3 test/test_multitensor.py TestMultiTensor.test_data_parallel_resnet_train_step
+    - name: Run SDXL on NULL backend
+      run: MAX_BUFFER_SIZE=0 PYTHONPATH="." NULL=1 DEBUG=1 python3 examples/sdxl.py --seed 0 --noshow --timing --fakeweights
+    # TODO: support fake weights
+    #- name: Run LLaMA 7B on 4 fake devices
+    #  run: NULL=1 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 3 --temperature 0 --timing
+    - name: Run GC tests
+      run: PYTHONPATH="." python test/external/external_uop_gc.py
+    - name: Run process replay tests
+      uses: ./.github/actions/process-replay
+    - name: Repo line count < 14000 lines
+      run: MAX_LINE_COUNT=14000 python sz.py

  fuzzing:
    name: Fuzzing
@@ -348,19 +384,25 @@ jobs:
    - name: Setup Environment
      uses: ./.github/actions/setup-tinygrad
      with:
-        key: fuzzing-minimal
-        deps: testing_minimal
+        key: fuzzing-unit
+        deps: testing_unit
    - name: Fuzz Test symbolic
      run: python test/external/fuzz_symbolic.py
+    - name: Fuzz Test fast idiv
+      run: python test/external/fuzz_fast_idiv.py
    - name: Fuzz Test shapetracker
      run: |
        PYTHONPATH="." python test/external/fuzz_shapetracker.py
        PYTHONPATH="." python test/external/fuzz_shapetracker_math.py
+    - name: Fuzz Test shape ops
+      run: python test/external/fuzz_shape_ops.py

  testgpuimage:
    name: 'GPU IMAGE Tests'
    runs-on: ubuntu-22.04
    timeout-minutes: 10
+    env:
+      IGNORE_OOB: 0
    steps:
      - name: Checkout Code
        uses: actions/checkout@v4
@@ -378,6 +420,8 @@ jobs:
        run: |
          PYTHONPATH="." GPU=1 IMAGE=2 python -m pytest -n=auto test/test_ops.py --durations=20
          PYTHONPATH="." GPU=1 IMAGE=2 python3 test/models/test_end2end.py TestEnd2End.test_linear_mnist
+      - name: Run fused optimizer tests
+        run: PYTHONPATH="." GPU=1 FUSE_OPTIM=1 python -m pytest -n=auto test/models/test_mnist.py
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

@@ -385,6 +429,8 @@ jobs:
    name: 'openpilot Compile Tests'
    runs-on: ubuntu-22.04
    timeout-minutes: 10
+    env:
+      IGNORE_OOB: 0
    steps:
      - name: Checkout Code
        uses: actions/checkout@v4
@@ -394,20 +440,25 @@ jobs:
          key: openpilot-compile
          deps: testing
          opencl: 'true'
+          llvm: 'true'
      - name: Test openpilot model kernel count and gate usage
        run: |
-          PYTHONPATH="." ALLOWED_KERNEL_COUNT=209 ALLOWED_READ_IMAGE=2138 ALLOWED_GATED_READ_IMAGE=29 FLOAT16=0 GPU=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx
+          PYTHONPATH="." ALLOWED_KERNEL_COUNT=209 ALLOWED_READ_IMAGE=2137 ALLOWED_GATED_READ_IMAGE=29 FLOAT16=0 GPU=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx
      - name: Test openpilot alt model correctness (float32)
        run: PYTHONPATH="." FLOAT16=0 DEBUGCL=1 GPU=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/3799fe46b3a629e491d4b8498b8ae83e4c88c304/selfdrive/modeld/models/supercombo.onnx
      - name: Test openpilot fastvits model correctness (float32)
        run: PYTHONPATH="." FLOAT16=0 DEBUGCL=1 GPU=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
+      - name: Test openpilot LLVM compile
+        run: PYTHONPATH="." LLVM=1 LLVMOPT=1 JIT=2 BEAM=0 IMAGE=0 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

-  testopencl:
-    name: 'ONNX+Optimization Tests'
+  testonnxcpu:
+    name: 'ONNX (CPU) Tests'
    runs-on: ubuntu-22.04
    timeout-minutes: 20
+    env:
+      IGNORE_OOB: 0

    steps:
      - name: Checkout Code
@@ -415,13 +466,10 @@ jobs:
      - name: Setup Environment
        uses: ./.github/actions/setup-tinygrad
        with:
-          key: onnxoptl
-          deps: testing,testing_tf
+          key: onnxoptc
+          deps: testing
          python-version: '3.11'
-          opencl: 'true'
          llvm: 'true'
-      - name: Test ONNX (GPU)
-        run: GPU=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
      - name: Test ONNX (CPU)
        run: CPU=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
      - name: Test ONNX (LLVM)
@@ -430,11 +478,36 @@ jobs:
        run: CPU=1 PYTHONPATH=. python3 test/external/external_test_onnx_ops.py
      - name: Test Quantize ONNX
        run: CPU=1 PYTHONPATH=. python3 test/test_quantize_onnx.py
-      - name: Run CLOUD=1 Test
+      - name: Run REMOTE=1 Test
        run: |
-          CLOUDDEV=CPU CLOUD=1 python3 test/test_tiny.py
-          CLOUDDEV=GPU CLOUD=1 python3 test/test_tiny.py
-          CLOUDDEV=GPU IMAGE=2 CLOUD=1 python3 test/test_tiny.py
+          REMOTEDEV=CPU REMOTE=1 python3 -m pytest test/test_tiny.py test/test_jit.py test/test_multitensor.py
+      - name: Run process replay tests
+        uses: ./.github/actions/process-replay
+
+  testopencl:
+    name: 'ONNX (GPU)+Optimization Tests'
+    runs-on: ubuntu-22.04
+    timeout-minutes: 20
+    env:
+      IGNORE_OOB: 0
+
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v4
+      - name: Setup Environment
+        uses: ./.github/actions/setup-tinygrad
+        with:
+          key: onnxoptl
+          deps: testing
+          pydeps: "tensorflow==2.15.1 tensorflow_addons"
+          python-version: '3.11'
+          opencl: 'true'
+      - name: Test ONNX (GPU)
+        run: GPU=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
+      - name: Run REMOTE=1 Test
+        run: |
+          REMOTEDEV=GPU REMOTE=1 python3 -m pytest test/test_tiny.py test/test_image_dtype.py test/test_jit.py
+          REMOTEDEV=GPU IMAGE=2 REMOTE=1 python3 -m pytest test/test_tiny.py test/test_image_dtype.py
      - name: Test Optimization Helpers
        run: PYTHONPATH="." DEBUG=1 python3 extra/optimization/test_helpers.py
      - name: Test Action Space
@@ -451,7 +524,9 @@ jobs:
  testmodels:
    name: Models (llvm+cpu+gpu)
    runs-on: ubuntu-22.04
-    timeout-minutes: 10
+    timeout-minutes: 15
+    env:
+      IGNORE_OOB: 0
    steps:
      - name: Checkout Code
        uses: actions/checkout@v4
@@ -475,6 +550,8 @@ jobs:
    name: Linux (DSP)
    runs-on: ubuntu-24.04
    timeout-minutes: 15
+    env:
+      IGNORE_OOB: 0
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
@@ -496,10 +573,12 @@ jobs:
        tags: qemu-hexagon:latest
        cache-from: type=gha
        cache-to: type=gha,mode=min
+    - name: Set MOCKDSP env
+      run: printf "MOCKDSP=1" >> $GITHUB_ENV
    - name: Run test_tiny on DSP
      run: DEBUG=2 DSP=1 python test/test_tiny.py
    - name: Test transcendentals
-      run: CC=clang-19 PYTHONPATH="." DEBUG=2 DSP=1 python test/test_transcendental.py TestTranscendentalVectorized
+      run: CC=clang-20 PYTHONPATH="." DEBUG=2 DSP=1 python test/test_transcendental.py TestTranscendentalVectorized
    - name: Test quantize onnx
      run: PYTHONPATH="." DEBUG=2 DSP=1 python3 test/test_quantize_onnx.py
    - name: Test LLVM=1 DEVECTORIZE=0
@@ -535,11 +614,55 @@ jobs:
    - name: Run process replay tests
      uses: ./.github/actions/process-replay

-  tests:
+  testamd:
    strategy:
      fail-fast: false
      matrix:
-        backend: [llvm, cpu, gpu, ptx, amd, nv] #, triton]
+        backend: [amd, amdllvm]
+
+    name: Linux (${{ matrix.backend }})
+    runs-on: ubuntu-22.04
+    timeout-minutes: 20
+    env:
+      IGNORE_OOB: 0
+      AMD: 1
+      MOCKGPU: 1
+      FORWARD_ONLY: 1
+      AMD_LLVM: ${{ matrix.backend == 'amdllvm' && '1' || matrix.backend != 'amdllvm' && '0' }}
+      PYTHONPATH: ${{ github.workspace }}
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v4
+      - name: Setup Environment
+        uses: ./.github/actions/setup-tinygrad
+        with:
+          key: ${{ matrix.backend }}-minimal
+          deps: testing_minimal
+          amd: 'true'
+          llvm: ${{ matrix.backend == 'amdllvm' && 'true' }}
+      - name: Check Device.DEFAULT and print some source
+        run: |
+          python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['AMD'], Device.DEFAULT"
+          DEBUG=5 FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
+      - name: Run LLVM test
+        if: matrix.backend=='amdllvm'
+        run: python test/test_amd_llvm.py
+      - name: Run pytest (amd)
+        run: python -m pytest -n=auto test/test_ops.py test/test_dtype.py test/test_dtype_alu.py test/test_linearizer.py test/test_randomness.py test/test_jit.py test/test_graph.py test/test_multitensor.py test/test_hcq.py test/external/external_test_am.py --durations=20
+      - name: Run TRANSCENDENTAL math
+        run: TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20
+      - name: Run TestOps.test_add with SQTT
+        run: |
+          PROFILE=1 SQTT=1 DEBUG=5 python3 test/test_ops.py TestOps.test_add
+          extra/sqtt/rgptool.py create "/tmp/profile.pkl.$USER" -o /tmp/gpu0.rgp
+      - name: Run process replay tests
+        uses: ./.github/actions/process-replay
+
+  testnvidia:
+    strategy:
+      fail-fast: false
+      matrix:
+        backend: [ptx, nv]

    name: Linux (${{ matrix.backend }})
    runs-on: ubuntu-22.04
@@ -552,29 +675,50 @@ jobs:
        uses: ./.github/actions/setup-tinygrad
        with:
          key: ${{ matrix.backend }}-minimal
-          deps: testing_minimal${{matrix.backend=='ptx'&&',cuda'||matrix.backend=='triton'&&',triton'||''}}
-          opencl: ${{ matrix.backend == 'gpu' && 'true' }}
-          amd: ${{ matrix.backend == 'amd' && 'true' }}
-          llvm: ${{ (matrix.backend == 'amd' || matrix.backend == 'llvm') && 'true' }}
-          cuda: ${{ (matrix.backend == 'ptx' || matrix.backend == 'triton' || matrix.backend == 'nv') && 'true' }}
+          deps: testing_minimal
+          cuda: 'true'
      - name: Set env
-        run: printf "${{ matrix.backend == 'llvm' && 'LLVM=1' || matrix.backend == 'cpu' && 'CPU=1' || matrix.backend == 'gpu' && 'GPU=1' || matrix.backend == 'PTX' && 'FORWARD_ONLY=1\nJIT=1\nOPT=2\nCUDA=1\nPTX=1\nMOCKGPU=1' || matrix.backend == 'triton' && 'FORWARD_ONLY=1\nJIT=1\nOPT=2\nNV=1\nMOCKGPU=1\nTRITON=1\nTRITON_PTXAS_PATH=/usr/bin/ptxas' || matrix.backend == 'amd' && 'AMD=1\nMOCKGPU=1\nFORWARD_ONLY=1' || matrix.backend == 'nv' && 'NV=1\nMOCKGPU=1\nFORWARD_ONLY=1' }}" >> $GITHUB_ENV
+        run: printf "${{ matrix.backend == 'PTX' && 'FORWARD_ONLY=1\nJIT=1\nOPT=2\nCUDA=1\nPTX=1\nMOCKGPU=1' || matrix.backend == 'nv' && 'NV=1\nMOCKGPU=1\nFORWARD_ONLY=1' }}" >> $GITHUB_ENV
      - name: Check Device.DEFAULT and print some source
        run: |
-          PYTHONPATH=${{ github.workspace }} python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['LLVM','CPU','CUDA','GPU','AMD','NV'], Device.DEFAULT"
+          PYTHONPATH=${{ github.workspace }} python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['CUDA','NV'], Device.DEFAULT"
          DEBUG=5 PYTHONPATH=${{ github.workspace }} FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
-      - name: Run pytest (not cuda or amd)
-        if: matrix.backend!='ptx' && matrix.backend!='triton' && matrix.backend != 'amd' && matrix.backend != 'nv'
-        run: python -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --durations=20
      - name: Run pytest (cuda)
-        if: matrix.backend=='ptx'||matrix.backend=='triton'||matrix.backend=='nv'
-        run: python -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --ignore test/test_gc.py --durations=20
-      - name: Run pytest (amd)
-        if: matrix.backend=='amd'
-        run: python -m pytest -n=auto test/test_ops.py test/test_dtype.py test/test_dtype_alu.py test/test_linearizer.py test/test_randomness.py test/imported/test_indexing.py test/test_hcq.py test/external/external_test_am.py --durations=20
-      - name: Run pytest (amd with llvm backend)
-        if: matrix.backend=='amd'
-        run: AMD_LLVM=1 python -m pytest -n=auto test/test_ops.py test/test_dtype.py test/test_dtype_alu.py test/test_linearizer.py test/test_randomness.py test/imported/test_indexing.py test/test_hcq.py test/external/external_test_am.py test/test_amd_llvm.py --durations=20
+        # skip multitensor because it's slow
+        run: python -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --ignore test/test_gc.py --ignore test/test_multitensor.py --durations=20
+      - name: Run process replay tests
+        uses: ./.github/actions/process-replay
+
+  tests:
+    strategy:
+      fail-fast: false
+      matrix:
+        backend: [llvm, cpu, gpu]
+
+    name: Linux (${{ matrix.backend }})
+    runs-on: ubuntu-22.04
+    timeout-minutes: 20
+    env:
+      IGNORE_OOB: 0
+
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v4
+      - name: Setup Environment
+        uses: ./.github/actions/setup-tinygrad
+        with:
+          key: ${{ matrix.backend }}-minimal
+          deps: testing_minimal
+          opencl: ${{ matrix.backend == 'gpu' && 'true' }}
+          llvm: ${{ matrix.backend == 'llvm' && 'true' }}
+      - name: Set env
+        run: printf "${{ matrix.backend == 'llvm' && 'LLVM=1' || matrix.backend == 'cpu' && 'CPU=1' || matrix.backend == 'gpu' && 'GPU=1' }}" >> $GITHUB_ENV
+      - name: Check Device.DEFAULT and print some source
+        run: |
+          PYTHONPATH=${{ github.workspace }} python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['LLVM','CPU','GPU'], Device.DEFAULT"
+          DEBUG=5 PYTHONPATH=${{ github.workspace }} FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
+      - name: Run pytest (not cuda)
+        run: python -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --durations=20
      - name: Run TRANSCENDENTAL math
        run: TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20
      - name: Run process replay tests
@@ -586,6 +730,8 @@ jobs:
    name: MacOS (unit)
    runs-on: macos-14
    timeout-minutes: 20
+    env:
+      IGNORE_OOB: 0

    steps:
    - name: Checkout Code
@@ -600,11 +746,11 @@ jobs:
        cuda: 'true'
        llvm: 'true'
    - name: Run real world test
-      run: JIT=2 METAL=1 python -m pytest -n=auto test/models/test_real_world.py --durations=20
+      run: METAL=1 python -m pytest -n=auto test/models/test_real_world.py --durations=20
    - name: Test models (Metal)
-      run: JIT=2 METAL=1 python -m pytest -n=auto test/models --durations=20
+      run: METAL=1 python -m pytest -n=auto test/models -v --durations=20
    - name: Run ONNX
-      run: JIT=2 METAL=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
+      run: METAL=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
    - name: Test tensor core ops (fake)
      run: TC=2 METAL=1 DEBUG=3 python test/test_ops.py TestOps.test_gemm
    - name: Test tensor core ops (real)
@@ -615,8 +761,6 @@ jobs:
      run: PYTHONPATH="." METAL=1 IGNORE_BEAM_CACHE=1 python3 -m pytest extra/optimization/test_beam_search.py
    - name: Fuzz Test linearizer
      run: PYTHONPATH="." METAL=1 DEPTH=4 FUZZ_N=50 FUZZ_MAX_SIZE=1000000 python test/external/fuzz_linearizer.py
-    # - name: Fuzz Test models schedule
-    # run: FUZZ_SCHEDULE=1 FUZZ_SCHEDULE_MAX_PATHS=5 python -m pytest test/models/test_train.py test/models/test_end2end.py
    - name: Run TRANSCENDENTAL math
      run: TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20
    - name: Run pytest (amd)
@@ -656,7 +800,10 @@ jobs:
      uses: ./.github/actions/setup-tinygrad
      with:
        key: osx-webgpu
+        deps: testing_minimal
        webgpu: 'true'
+    - name: Test infinity math in WGSL
+      run: WEBGPU=1 python -m pytest -n=auto test/test_renderer_failures.py::TestWGSLFailures::test_multiply_infinity --durations=20
    - name: Build WEBGPU Efficientnet
      run: WEBGPU=1 WEBGPU_BACKEND="WGPUBackendType_Metal" python3 -m examples.compile_efficientnet
    - name: Clean npm cache
@@ -665,6 +812,67 @@ jobs:
      run: npm install puppeteer
    - name: Run WEBGPU Efficientnet
      run: node test/web/test_webgpu.js
+    - name: Run VIZ tests as external package
+      run: |
+        mkdir $GITHUB_WORKSPACE/test_dir
+        cd $GITHUB_WORKSPACE/test_dir
+        python -m venv venv
+        source venv/bin/activate
+        pip install $GITHUB_WORKSPACE
+        cp $GITHUB_WORKSPACE/test/web/test_viz.js .
+        node test_viz.js
+
+  osxremote:
+    name: MacOS (remote metal)
+    runs-on: macos-15
+    timeout-minutes: 10
+    env:
+      REMOTE: 1
+      REMOTEDEV: METAL
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v4
+      - name: Setup Environment
+        uses: ./.github/actions/setup-tinygrad
+        with:
+          key: macos-remote
+          deps: testing_minimal
+      - name: Check Device.DEFAULT and print some source
+        run: |
+          python -c "from tinygrad import Device; assert Device.DEFAULT == 'REMOTE', Device.DEFAULT"
+          python -c "from tinygrad import Device; assert Device.default.properties.real_device == 'METAL', Device.default.properties.real_device"
+          DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus
+      - name: Run REMOTE=1 Test
+        run: |
+          python3 -m pytest test/test_tiny.py test/test_jit.py test/test_subbuffer.py test/test_graph.py test/test_multitensor.py test/test_tensor_variable.py
+
+  amdremote:
+    name: Linux (remote amd)
+    runs-on: ubuntu-22.04
+    timeout-minutes: 20
+    env:
+      REMOTE: 1
+      REMOTEDEV: 'AMD'
+      PYTHONPATH: ${{ github.workspace }}
+      MOCKGPU: 1
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v4
+      - name: Setup Environment
+        uses: ./.github/actions/setup-tinygrad
+        with:
+          key: linux-remote-amd
+          deps: testing_minimal
+          amd: 'true'
+          llvm: 'true'
+      - name: Check Device.DEFAULT and print some source
+        run: |
+          python -c "from tinygrad import Device; assert Device.DEFAULT == 'REMOTE', Device.DEFAULT"
+          python -c "from tinygrad import Device; assert Device.default.properties.real_device == 'AMD', Device.default.properties.real_device"
+          DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus
+      - name: Run REMOTE=1 Test
+        run: |
+          python3 -m pytest test/test_tiny.py test/test_jit.py test/test_subbuffer.py test/test_graph.py test/test_multitensor.py test/test_tensor_variable.py

  osxtests:
    strategy:
@@ -674,6 +882,8 @@ jobs:
    name: MacOS (${{ matrix.backend }})
    runs-on: macos-15
    timeout-minutes: 20
+    env:
+      IGNORE_OOB: 0
    steps:
      - name: Checkout Code
        uses: actions/checkout@v4
@@ -684,7 +894,7 @@ jobs:
          deps: testing_minimal
          llvm: ${{ matrix.backend == 'llvm' && 'true' }}
      - name: Set env
-        run: printf "${{ matrix.backend == 'llvm' && 'LLVM=1' || matrix.backend == 'cpu' && 'CPU=1' || matrix.backend == 'metal' && 'METAL=1\nJIT=2'}}" >> $GITHUB_ENV
+        run: printf "${{ matrix.backend == 'llvm' && 'LLVM=1' || matrix.backend == 'cpu' && 'CPU=1' || matrix.backend == 'metal' && 'METAL=1'}}" >> $GITHUB_ENV
      - name: Check Device.DEFAULT and print some source
        run: |
          python -c "from tinygrad import Device; assert Device.DEFAULT == '${{ matrix.backend }}'.upper(), Device.DEFAULT"
@@ -707,7 +917,9 @@ jobs:

    name: Windows (${{ matrix.backend }})
    runs-on: windows-latest
-    timeout-minutes: 10
+    timeout-minutes: 15
+    env:
+      IGNORE_OOB: 0
    steps:
      - name: Checkout Code
        uses: actions/checkout@v4
@@ -60,3 +60,5 @@ comgr_*
 site/
 profile_stats
 *.log
+target
+.mypy_cache
@@ -28,7 +28,7 @@ repos:
        pass_filenames: false
      - id: tests
        name: subset of tests
-        entry: env PYTHONPATH="." python3 -m pytest -n=4 test/unit/ test/test_ops.py test/test_dtype.py test/test_schedule.py test/test_assign.py test/test_symbolic_shapetracker.py
+        entry: env PYTHONPATH="." python3 -m pytest -n=4 test/test_ops.py test/test_dtype.py test/test_schedule.py test/test_assign.py
        language: system
        always_run: true
        pass_filenames: false
@@ -7,7 +7,7 @@ extension-pkg-whitelist=scipy,cereal.messaging.messaging_pyx,PyQt5,av

 # Add files or directories to the blacklist. They should be base names, not
 # paths.
-ignore=CVS,autogen,msm_kgsl.py,runtime
+ignore=CVS,autogen,msm_kgsl.py,runtime,.venv

 # Add files or directories matching the regex patterns to the blacklist. The
 # regex matches against base names, not paths.
@@ -0,0 +1,17 @@
+# tinygrad agents
+
+Hello agent. You are one of the most talented programmers of your generation.
+
+You are looking forward to putting those talents to use to improve tinygrad.
+
+## philosophy
+
+tinygrad is a **tensor** library focused on beauty and minimalism, while still matching the functionality of PyTorch and JAX.
+
+Every line must earn its keep. Prefer readability over cleverness. We believe that if carefully designed, 10 lines can have the impact of 1000.
+
+Never mix functionality changes with whitespace changes. All functionality changes must be tested.
+
+## style
+
+Use **2-space indentation**, and keep lines to a maximum of **150 characters**. Match the existing style.
@@ -21,11 +21,11 @@ tinygrad: For something between [PyTorch](https://github.com/pytorch/pytorch) an

 ---

-This may not be the best deep learning framework, but it is a deep learning framework.
+Despite tinygrad's size, it is a fully featured deep learning framework.

-Due to its extreme simplicity, it aims to be the easiest framework to add new accelerators to, with support for both inference and training. If XLA is CISC, tinygrad is RISC.
+Due to its extreme simplicity, it is the easiest framework to add new accelerators to, with support for both inference and training. If XLA is CISC, tinygrad is RISC.

-tinygrad is still alpha software, but we [raised some money](https://geohot.github.io/blog/jekyll/update/2023/05/24/the-tiny-corp-raised-5M.html) to make it good. Someday, we will tape out chips.
+tinygrad is now beta software, we [raised some money](https://geohot.github.io/blog/jekyll/update/2023/05/24/the-tiny-corp-raised-5M.html) to make it good. Someday, we will tape out chips.

 ## Features

@@ -9,7 +9,7 @@ if [[ ! $(clang2py -V) ]]; then
  pip install clang==14.0.6
  git clone https://github.com/nimlgen/ctypeslib.git
  cd ctypeslib
-  pip install --user .
+  pip install .
  clang2py -V
  popd
 fi
@@ -35,7 +35,7 @@ def _try_dlopen_$name():
  for candidate in PATHS_TO_TRY:
    try: return ctypes.CDLL(candidate)
    except OSError: pass
-  raise RuntimeError("library $name not found")
+  return None
 EOF
 }

@@ -78,16 +78,17 @@ generate_kfd() {
  clang2py /usr/include/linux/kfd_ioctl.h -o $BASE/kfd.py -k cdefstum

  fixup $BASE/kfd.py
-  sed -i "s\import ctypes\import ctypes, os\g" $BASE/kfd.py
-  sed -i "s\import fcntl, functools\import functools" $BASE/kfd.py
-  sed -i "s\import ctypes,os\a from tinygrad.runtime.support import HWInterface\g" $BASE/kfd.py
-  sed -i "s\def _do_ioctl(__idir, __base, __nr, __user_struct, __fd, **kwargs):\def _do_ioctl(__idir, __base, __nr, __user_struct, __fd:HWInterface, **kwargs):\g" $BASE/kfd.py
-  sed -i "s\fcntl.ioctl(__fd, (__idir<<30)\__fd.ioctl((__idir<<30)\g" $BASE/kfd.py
+  sed -i "s/import ctypes/import ctypes, os/g" $BASE/kfd.py
+  sed -i "s/import fcntl, functools/import functools/g" $BASE/kfd.py
+  sed -i "/import functools/a from tinygrad.runtime.support.hcq import FileIOInterface" $BASE/kfd.py
+  sed -i "s/def _do_ioctl(__idir, __base, __nr, __user_struct, __fd, \*\*kwargs):/def _do_ioctl(__idir, __base, __nr, __user_struct, __fd:FileIOInterface, \*\*kwargs):/g" $BASE/kfd.py
+  sed -i "s/fcntl.ioctl(__fd, (__idir<<30)/__fd.ioctl((__idir<<30)/g" $BASE/kfd.py
+  sed -i "s/!!/not not /g" $BASE/kfd.py
  python3 -c "import tinygrad.runtime.autogen.kfd"
 }

 generate_cuda() {
-  clang2py /usr/include/cuda.h -o $BASE/cuda.py -l /usr/lib/x86_64-linux-gnu/libcuda.so
+  clang2py /usr/include/cuda.h --clang-args="-D__CUDA_API_VERSION_INTERNAL" -o $BASE/cuda.py -l /usr/lib/x86_64-linux-gnu/libcuda.so
  sed -i "s\import ctypes\import ctypes, ctypes.util\g" $BASE/cuda.py
  sed -i "s\ctypes.CDLL('/usr/lib/x86_64-linux-gnu/libcuda.so')\ctypes.CDLL(ctypes.util.find_library('cuda'))\g" $BASE/cuda.py
  fixup $BASE/cuda.py
@@ -104,10 +105,10 @@ generate_nvrtc() {
 }

 generate_nv() {
-  NVKERN_COMMIT_HASH=d6b75a34094b0f56c2ccadf14e5d0bd515ed1ab6
+  NVKERN_COMMIT_HASH=81fe4fb417c8ac3b9bdcc1d56827d116743892a5
  NVKERN_SRC=/tmp/open-gpu-kernel-modules-$NVKERN_COMMIT_HASH
  if [ ! -d "$NVKERN_SRC" ]; then
-    git clone https://github.com/tinygrad/open-gpu-kernel-modules $NVKERN_SRC
+    git clone https://github.com/NVIDIA/open-gpu-kernel-modules $NVKERN_SRC
    pushd .
    cd $NVKERN_SRC
    git reset --hard $NVKERN_COMMIT_HASH
@@ -116,15 +117,19 @@ generate_nv() {

  clang2py -k cdefstum \
    extra/nv_gpu_driver/clc6c0qmd.h \
+    extra/nv_gpu_driver/clcec0qmd.h \
    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/cl0080.h \
    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/cl2080_notification.h \
    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/clc56f.h \
-    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/clc56f.h \
-    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/clc56f.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/clc86f.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/clc96f.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/clc761.h \
    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/cl83de.h \
    $NVKERN_SRC/src/nvidia/generated/g_allclasses.h \
    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/clc6c0.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/clcdc0.h \
    $NVKERN_SRC/kernel-open/nvidia-uvm/clc6b5.h \
+    $NVKERN_SRC/kernel-open/nvidia-uvm/clc9b5.h \
    $NVKERN_SRC/kernel-open/nvidia-uvm/uvm_ioctl.h \
    $NVKERN_SRC/kernel-open/nvidia-uvm/uvm_linux_ioctl.h \
    $NVKERN_SRC/kernel-open/nvidia-uvm/hwref/ampere/ga100/dev_fault.h \
@@ -149,6 +154,8 @@ generate_nv() {
  sed -i "s\import ctypes\import ctypes, os\g" $BASE/nv_gpu.py
  sed -i 's/#\?\s\([A-Za-z0-9_]\+\) = MW ( \([0-9]\+\) : \([0-9]\+\) )/\1 = (\2 , \3)/' $BASE/nv_gpu.py # NVC6C0_QMDV03_00 processing
  sed -i 's/#\sdef NVC6C0_QMD\([A-Za-z0-9_()]\+\):/def NVC6C0_QMD\1:/' $BASE/nv_gpu.py
+  sed -i 's/#\sdef NVCEC0_QMD\([A-Za-z0-9_()]\+\):/def NVCEC0_QMD\1:/' $BASE/nv_gpu.py
+  sed -E -i -n '/^def (NVCEC0_QMDV05_00_RELEASE)(_ENABLE)\(i\):/{p;s//\1'"0"'\2=\1\2(0)\n\1'"1"'\2=\1\2(1)/;H;b};p;${x;s/^\n//;p}' "$BASE/nv_gpu.py"
  sed -i 's/#\s*return MW(\([0-9i()*+]\+\):\([0-9i()*+]\+\))/    return (\1 , \2)/' $BASE/nv_gpu.py
  sed -i 's/#\?\s*\(.*\)\s*=\s*\(NV\)\?BIT\(32\)\?\s*(\s*\([0-9]\+\)\s*)/\1 = (1 << \4)/' $BASE/nv_gpu.py # name = BIT(x) -> name = (1 << x)
  sed -i "s/UVM_\([A-Za-z0-9_]\+\) = \['i', '(', '\([0-9]\+\)', ')'\]/UVM_\1 = \2/" $BASE/nv_gpu.py # UVM_name = ['i', '(', '<num>', ')'] -> UVM_name = <num>
@@ -220,7 +227,7 @@ generate_libc() {

  sed -i "s\import ctypes\import ctypes, ctypes.util, os\g" $BASE/libc.py
  sed -i "s\FIXME_STUB\libc\g" $BASE/libc.py
-  sed -i "s\FunctionFactoryStub()\None if (libc_path := ctypes.util.find_library('c')) is None else ctypes.CDLL(libc_path)\g" $BASE/libc.py
+  sed -i "s\FunctionFactoryStub()\None if (libc_path := ctypes.util.find_library('c')) is None else ctypes.CDLL(libc_path, use_errno=True)\g" $BASE/libc.py

  fixup $BASE/libc.py
 }
@@ -287,11 +294,19 @@ generate_vfio() {
  fixup $BASE/vfio.py
  sed -i "s\import ctypes\import ctypes, os\g" $BASE/vfio.py
  sed -i "s\import fcntl, functools\import functools" $BASE/vfio.py
-  sed -i "s\import ctypes,os\a from tinygrad.runtime.support import HWInterface\g" $BASE/vfio.py
+  sed -i "s\import ctypes,os\a from tinygrad.runtime.support import FileIOInterface\g" $BASE/vfio.py
  sed -i "s\fcntl.ioctl(__fd, (__idir<<30)\return __fd.ioctl((__idir<<30)\g" $BASE/vfio.py
 }

 generate_am() {
+  AMKERN_COMMIT_HASH=ceb12c04e2b5b53ec0779362831f5ee40c4921e4
+  AMKERN_SRC=/tmp/ROCK-Kernel-Driver-$AMKERN_COMMIT_HASH
+  if [ ! -d "$AMKERN_SRC" ]; then
+    git clone https://github.com/ROCm/ROCK-Kernel-Driver $AMKERN_SRC --depth 1
+  fi
+  AMKERN_AMD=$AMKERN_SRC/drivers/gpu/drm/amd/
+  AMKERN_INC=$AMKERN_AMD/include/
+
  clang2py -k cdefstum \
    extra/amdpci/headers/v11_structs.h \
    extra/amdpci/headers/v12_structs.h \
@@ -302,7 +317,7 @@ generate_am() {
    extra/amdpci/headers/amdgpu_psp.h \
    extra/amdpci/headers/amdgpu_irq.h \
    extra/amdpci/headers/amdgpu_doorbell.h \
-    extra/amdpci/headers/soc15_ih_clientid.h \
+    $AMKERN_INC/soc15_ih_clientid.h \
    --clang-args="-include stdint.h" \
    -o $BASE/am/am.py
  fixup $BASE/am/am.py
@@ -310,190 +325,73 @@ generate_am() {
  sed -i "s\AMDGPU_PTE_MTYPE_VG10(2)\AMDGPU_PTE_MTYPE_VG10(0, 2)\g" $BASE/am/am.py # incorrect parsing (TODO: remove when clang2py is gone).

  clang2py -k cdefstum \
-    extra/hip_gpu_driver/kfd_pm4_headers_ai.h \
-    extra/hip_gpu_driver/soc15d.h \
+    $AMKERN_AMD/amdkfd/kfd_pm4_headers_ai.h \
+    $AMKERN_AMD/amdgpu/soc15d.h \
    -o $BASE/am/pm4_soc15.py
  fixup $BASE/am/pm4_soc15.py

  clang2py -k cdefstum \
-    extra/hip_gpu_driver/kfd_pm4_headers_ai.h \
-    extra/hip_gpu_driver/nvd.h \
+    $AMKERN_AMD/amdkfd/kfd_pm4_headers_ai.h \
+    $AMKERN_AMD/amdgpu/nvd.h \
    -o $BASE/am/pm4_nv.py
  fixup $BASE/am/pm4_nv.py

  clang2py -k cdefstum \
-    extra/amdpci/headers/vega10_enum.h \
+    $AMKERN_INC/vega10_enum.h \
    -o $BASE/am/vega10.py
  fixup $BASE/am/vega10.py

  clang2py -k cdefstum \
-    extra/amdpci/headers/navi10_enum.h \
+    $AMKERN_INC/navi10_enum.h \
    -o $BASE/am/navi10.py
  fixup $BASE/am/navi10.py

  clang2py -k cdefstum \
-    extra/amdpci/headers/soc21_enum.h \
+    $AMKERN_INC/soc21_enum.h \
    -o $BASE/am/soc21.py
  fixup $BASE/am/soc21.py

  clang2py -k cdefstum \
-    extra/amdpci/headers/soc24_enum.h \
+    $AMKERN_INC/soc24_enum.h \
    -o $BASE/am/soc24.py
  fixup $BASE/am/soc24.py

-  clang2py -k cdefstum \
-    extra/amdpci/headers/mp_13_0_0_offset.h \
-    extra/amdpci/headers/mp_13_0_0_sh_mask.h \
-    -o $BASE/am/mp_13_0_0.py
-  fixup $BASE/am/mp_13_0_0.py
-
-  # 14_0_3 reuses 14_0_2
-  clang2py -k cdefstum \
-    extra/amdpci/headers/mp_14_0_2_offset.h \
-    extra/amdpci/headers/mp_14_0_2_sh_mask.h \
-    -o $BASE/am/mp_14_0_3.py
-  fixup $BASE/am/mp_14_0_3.py
-
-  clang2py -k cdefstum \
-    extra/amdpci/headers/mp_11_0_offset.h \
-    extra/amdpci/headers/mp_11_0_sh_mask.h \
-    -o $BASE/am/mp_11_0.py
-  fixup $BASE/am/mp_11_0.py
-
-  clang2py -k cdefstum \
-    extra/amdpci/headers/gc_9_4_3_offset.h \
-    extra/amdpci/headers/gc_9_4_3_sh_mask.h \
-    extra/amdpci/overlay/gc_9_4_3.h \
-    -o $BASE/am/gc_9_4_3.py
-  fixup $BASE/am/gc_9_4_3.py
-
-  clang2py -k cdefstum \
-    extra/amdpci/headers/gc_10_3_0_offset.h \
-    extra/amdpci/headers/gc_10_3_0_sh_mask.h \
-    -o $BASE/am/gc_10_3_0.py
-  fixup $BASE/am/gc_10_3_0.py
-
-  clang2py -k cdefstum \
-    extra/amdpci/headers/gc_11_0_0_offset.h \
-    extra/amdpci/headers/gc_11_0_0_sh_mask.h \
-    -o $BASE/am/gc_11_0_0.py
-  fixup $BASE/am/gc_11_0_0.py
-
-  clang2py -k cdefstum \
-    extra/amdpci/headers/gc_12_0_0_offset.h \
-    extra/amdpci/headers/gc_12_0_0_sh_mask.h \
-    -o $BASE/am/gc_12_0_0.py
-  fixup $BASE/am/gc_12_0_0.py
-
  clang2py -k cdefstum \
    extra/hip_gpu_driver/sdma_registers.h \
-    extra/hip_gpu_driver/vega10_sdma_pkt_open.h \
+    $AMKERN_AMD/amdgpu/vega10_sdma_pkt_open.h \
    --clang-args="-I/opt/rocm/include -x c++" \
    -o $BASE/am/sdma_4_0_0.py
  fixup $BASE/am/sdma_4_0_0.py

  clang2py -k cdefstum \
    extra/hip_gpu_driver/sdma_registers.h \
-    extra/hip_gpu_driver/navi10_sdma_pkt_open.h \
+    $AMKERN_AMD/amdgpu/navi10_sdma_pkt_open.h \
    --clang-args="-I/opt/rocm/include -x c++" \
    -o $BASE/am/sdma_5_0_0.py
  fixup $BASE/am/sdma_5_0_0.py

  clang2py -k cdefstum \
    extra/hip_gpu_driver/sdma_registers.h \
-    extra/hip_gpu_driver/sdma_v6_0_0_pkt_open.h \
+    $AMKERN_AMD/amdgpu/sdma_v6_0_0_pkt_open.h \
    --clang-args="-I/opt/rocm/include -x c++" \
    -o $BASE/am/sdma_6_0_0.py
  fixup $BASE/am/sdma_6_0_0.py

  clang2py -k cdefstum \
-    extra/amdpci/headers/mmhub_3_0_0_offset.h \
-    extra/amdpci/headers/mmhub_3_0_0_sh_mask.h \
-    -o $BASE/am/mmhub_3_0_0.py
-  fixup $BASE/am/mmhub_3_0_0.py
-
-  clang2py -k cdefstum \
-    extra/amdpci/headers/mmhub_3_0_2_offset.h \
-    extra/amdpci/headers/mmhub_3_0_2_sh_mask.h \
-    -o $BASE/am/mmhub_3_0_2.py
-  fixup $BASE/am/mmhub_3_0_2.py
-
-  clang2py -k cdefstum \
-    extra/amdpci/headers/nbio_2_3_offset.h \
-    extra/amdpci/headers/nbio_2_3_sh_mask.h \
-    -o $BASE/am/nbio_2_3_0.py
-  fixup $BASE/am/nbio_2_3_0.py
-
-  clang2py -k cdefstum \
-    extra/amdpci/headers/mmhub_4_1_0_offset.h \
-    extra/amdpci/headers/mmhub_4_1_0_sh_mask.h \
-    -o $BASE/am/mmhub_4_1_0.py
-  fixup $BASE/am/mmhub_4_1_0.py
-
-  clang2py -k cdefstum \
-    extra/amdpci/headers/nbio_4_3_0_offset.h \
-    extra/amdpci/headers/nbio_4_3_0_sh_mask.h \
-    -o $BASE/am/nbio_4_3_0.py
-  fixup $BASE/am/nbio_4_3_0.py
-
-  clang2py -k cdefstum \
-    extra/amdpci/headers/nbif_6_3_1_offset.h \
-    extra/amdpci/headers/nbif_6_3_1_sh_mask.h \
-    -o $BASE/am/nbif_6_3_1.py
-  fixup $BASE/am/nbif_6_3_1.py
-
-  clang2py -k cdefstum \
-    extra/amdpci/headers/nbio_7_9_0_offset.h \
-    extra/amdpci/headers/nbio_7_9_0_sh_mask.h \
-    -o $BASE/am/nbio_7_9_0.py
-  fixup $BASE/am/nbio_7_9_0.py
-
-  clang2py -k cdefstum \
-    extra/amdpci/headers/nbio_7_11_0_offset.h \
-    extra/amdpci/headers/nbio_7_11_0_sh_mask.h \
-    -o $BASE/am/nbio_7_11_0.py
-  fixup $BASE/am/nbio_7_11_0.py
-
-  clang2py -k cdefstum \
-    extra/amdpci/headers/osssys_6_0_0_offset.h \
-    extra/amdpci/headers/osssys_6_0_0_sh_mask.h \
-    -o $BASE/am/osssys_6_0_0.py
-  fixup $BASE/am/osssys_6_0_0.py
-
-  clang2py -k cdefstum \
-    extra/amdpci/headers/osssys_7_0_0_offset.h \
-    extra/amdpci/headers/osssys_7_0_0_sh_mask.h \
-    -o $BASE/am/osssys_7_0_0.py
-  fixup $BASE/am/osssys_7_0_0.py
-
-  clang2py -k cdefstum \
-    extra/amdpci/headers/smu_v13_0_0_ppsmc.h \
-    extra/amdpci/headers/smu13_driver_if_v13_0_0.h \
+    $AMKERN_AMD/pm/swsmu/inc/pmfw_if/smu_v13_0_0_ppsmc.h \
+    $AMKERN_AMD/pm/swsmu/inc/pmfw_if/smu13_driver_if_v13_0_0.h \
    extra/amdpci/headers/amdgpu_smu.h \
    -o $BASE/am/smu_v13_0_0.py
  fixup $BASE/am/smu_v13_0_0.py

  clang2py -k cdefstum \
-    extra/amdpci/headers/smu_v14_0_0_pmfw.h \
-    extra/amdpci/headers/smu_v14_0_2_ppsmc.h \
-    extra/amdpci/headers/smu14_driver_if_v14_0_0.h \
-    extra/amdpci/headers/smu14_driver_if_v14_0.h \
+    $AMKERN_AMD/pm/swsmu/inc/pmfw_if/smu_v14_0_0_pmfw.h \
+    $AMKERN_AMD/pm/swsmu/inc/pmfw_if/smu_v14_0_2_ppsmc.h \
+    $AMKERN_AMD/pm/swsmu/inc/pmfw_if/smu14_driver_if_v14_0.h \
    extra/amdpci/headers/amdgpu_smu.h \
    --clang-args="-include stdint.h" \
    -o $BASE/am/smu_v14_0_3.py
  fixup $BASE/am/smu_v14_0_3.py
-
-  clang2py -k cdefstum \
-    extra/amdpci/headers/hdp_6_0_0_offset.h \
-    extra/amdpci/headers/hdp_6_0_0_sh_mask.h \
-    -o $BASE/am/hdp_6_0_0.py
-  fixup $BASE/am/hdp_6_0_0.py
-
-  clang2py -k cdefstum \
-    extra/amdpci/headers/hdp_7_0_0_offset.h \
-    extra/amdpci/headers/hdp_7_0_0_sh_mask.h \
-    -o $BASE/am/hdp_7_0_0.py
-  fixup $BASE/am/hdp_7_0_0.py
 }

 generate_sqtt() {
@@ -515,6 +413,19 @@ generate_webgpu() {
  python3 -c "import tinygrad.runtime.autogen.webgpu"
 }

+generate_libusb() {
+  clang2py -k cdefstum \
+    /usr/include/libusb-1.0/libusb.h \
+    -o $BASE/libusb.py
+
+  fixup $BASE/libusb.py
+  sed -i "s\import ctypes\import ctypes, os\g" $BASE/libusb.py
+  sed -i "s/FIXME_STUB/libusb/g" "$BASE/libusb.py"
+  sed -i "s/libusb_le16_to_cpu = libusb_cpu_to_le16//g" "$BASE/libusb.py"
+  sed -i "s/FunctionFactoryStub()/None if (lib_path:=os.getenv('LIBUSB_PATH', ctypes.util.find_library('usb-1.0'))) is None else ctypes.CDLL(lib_path)/g" "$BASE/libusb.py"
+  python3 -c "import tinygrad.runtime.autogen.libusb"
+}
+
 if [ "$1" == "opencl" ]; then generate_opencl
 elif [ "$1" == "hip" ]; then generate_hip
 elif [ "$1" == "comgr" ]; then generate_comgr
@@ -535,6 +446,7 @@ elif [ "$1" == "adreno" ]; then generate_adreno
 elif [ "$1" == "pci" ]; then generate_pci
 elif [ "$1" == "vfio" ]; then generate_vfio
 elif [ "$1" == "webgpu" ]; then generate_webgpu
+elif [ "$1" == "libusb" ]; then generate_libusb
 elif [ "$1" == "all" ]; then generate_opencl; generate_hip; generate_comgr; generate_cuda; generate_nvrtc; generate_hsa; generate_kfd; generate_nv; generate_amd; generate_io_uring; generate_libc; generate_am; generate_webgpu
 else echo "usage: $0 <type>"
 fi
@@ -1,7 +1,7 @@
 # tinygrad is a tensor library, and as a tensor library it has multiple parts
 # 1. a "runtime". this allows buffer management, compilation, and running programs
 # 2. a "Device" that uses the runtime but specifies compute in an abstract way for all
-# 3. a "LazyBuffer" that fuses the compute into kernels, using memory only when needed
+# 3. a "UOp" that fuses the compute into kernels, using memory only when needed
 # 4. a "Tensor" that provides an easy to use frontend with autograd ".backward()"


@@ -39,7 +39,7 @@ DEVICE = "CPU"   # NOTE: you can change this!
 import struct
 from tinygrad.dtype import dtypes
 from tinygrad.device import Buffer, Device
-from tinygrad.ops import UOp, Ops
+from tinygrad.uop.ops import UOp, Ops
 from tinygrad.shape.shapetracker import ShapeTracker

 # allocate some buffers + load in values
@@ -51,11 +51,11 @@ b = Buffer(DEVICE, 1, dtypes.int32).allocate().copyin(memoryview(bytearray(struc
 # describe the computation
 buf_1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int32.ptr(), (), 1)
 buf_2 = UOp(Ops.DEFINE_GLOBAL, dtypes.int32.ptr(), (), 2)
-ld_1 = UOp(Ops.LOAD, dtypes.int32, (buf_1, ShapeTracker.from_shape((1,)).to_uop()))
-ld_2 = UOp(Ops.LOAD, dtypes.int32, (buf_2, ShapeTracker.from_shape((1,)).to_uop()))
+ld_1 = UOp(Ops.LOAD, dtypes.int32, (buf_1.view(ShapeTracker.from_shape((1,))),))
+ld_2 = UOp(Ops.LOAD, dtypes.int32, (buf_2.view(ShapeTracker.from_shape((1,))),))
 alu = ld_1 + ld_2
 output_buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int32.ptr(), (), 0)
-st_0 = UOp(Ops.STORE, dtypes.void, (output_buf, ShapeTracker.from_shape((1,)).to_uop(), alu))
+st_0 = UOp(Ops.STORE, dtypes.void, (output_buf.view(ShapeTracker.from_shape((1,))), alu))
 s = UOp(Ops.SINK, dtypes.void, (st_0,))

 # convert the computation to a "linearized" format (print the format)
@@ -74,30 +74,47 @@ fxn.exec([out, a, b])
 assert out.as_buffer().cast('I')[0] == 5


-print("******** third, the LazyBuffer ***********")
+print("******** third, the UOp ***********")

 from tinygrad.engine.realize import run_schedule
 from tinygrad.engine.schedule import create_schedule_with_vars
+from tinygrad.engine.grouper import get_kernelize_map

 # allocate some values + load in values
-a = UOp.metaop(Ops.EMPTY, (1,), dtypes.int32, DEVICE)
-b = UOp.metaop(Ops.EMPTY, (1,), dtypes.int32, DEVICE)
+a = UOp.new_buffer(DEVICE, 1, dtypes.int32)
+b = UOp.new_buffer(DEVICE, 1, dtypes.int32)
 a.buffer.allocate().copyin(memoryview(bytearray(struct.pack("I", 2))))
 b.buffer.allocate().copyin(memoryview(bytearray(struct.pack("I", 3))))

 # describe the computation
-out = a.alu(Ops.ADD, b)
+out = a + b
+s = UOp(Ops.SINK, dtypes.void, (out,))

-# schedule the computation as a list of kernels
-sched, _, becomes_map = create_schedule_with_vars(out.sink())
-for si in sched: print(si.ast.op)  # NOTE: the first two convert it to CPU
-# NOTE: UOps are no longer mutable, the scheduler gives you a map to lookup which BUFFER the result was written to
-out = becomes_map[out]
+# group the computation into kernels
+becomes_map = get_kernelize_map(s)
+
+# the compute maps to an assign
+assign = becomes_map[a+b]
+
+# the first source is the output buffer (data)
+assert assign.src[0].op is Ops.BUFFER
+# the second source is the kernel (compute)
+assert assign.src[1].op is Ops.KERNEL
+
+# schedule the kernel graph in a linear list
+s = UOp(Ops.SINK, dtypes.void, (assign,))
+sched, _ = create_schedule_with_vars(s)
+assert len(sched) == 1

 # DEBUGGING: print the compute ast
 print(sched[-1].ast)
 # NOTE: sched[-1].ast is the same as st_0 above

+# the output will be stored in a new buffer
+out = assign.buf_uop
+assert out.op is Ops.BUFFER and not out.buffer.is_allocated()
+print(out)
+
 # run that schedule
 run_schedule(sched)

@@ -36,7 +36,7 @@ optim.schedule_step()   # this will step the optimizer without running realize
 # 3. Create a schedule.

 # The weight Tensors have been assigned to, but not yet realized. Everything is still lazy at this point
-# l1.lazydata and l2.lazydata define a computation graph
+# l1.uop and l2.uop define a computation graph

 from tinygrad.engine.schedule import ScheduleItem
 schedule: List[ScheduleItem] = Tensor.schedule(l1, l2)
@@ -1,6 +1,6 @@
 # AM Driver

-AM driver is a userspace driver targeting AMD's 7900XTX. You only need tinygrad to send compute tasks to your GPU!
+AM driver is a userspace driver targeting AMD's RDNA3/RDNA4. You only need tinygrad to send compute tasks to your GPU!

 ## How to run?
 Make sure that amdgpu module is unloaded and just run tinygrad with `AMD=1`!
@@ -0,0 +1,109 @@
+# Kernel Creation
+
+Tinygrad lazily builds up a graph of Tensor operations. The Tensor graph includes a mix of:
+
+- Buffer and Assignment Ops: `BUFFER`, `BUFFER_VIEW`, `COPY`, `ASSIGN`
+- Movement Ops: `RESHAPE`, `EXPAND`, `PERMUTE`, `PAD`, `SHRINK`, `FLIP`
+- Compute Ops: `ADD`, `MUL`, `REDUCE_AXIS`, ...
+
+`Tensor.kernelize` creates the kernels and buffers needed to realize the output Tensor(s).
+
+## Kernelize flow
+
+Let's see how a multiply add Tensor graph becomes a fused elementwise kernel.
+
+```py
+# initialize 3 input buffers on the device
+a = Tensor([1]).realize()
+b = Tensor([2]).realize()
+c = Tensor([3]).realize()
+
+# create the Tensor graph
+mul = a*b
+out = mul+c
+
+print(mul) # <Tensor <UOp METAL (1,) int (<Ops.MUL: 48>, None)> on METAL with grad None>
+print(out) # <Tensor <UOp METAL (1,) int (<Ops.ADD: 52>, None)> on METAL with grad None>
+
+out.kernelize()
+
+print(mul) # <Tensor <UOp METAL (1,) int (<Ops.MUL: 48>, None)> on METAL with grad None>
+print(out) # <Tensor <UOp METAL (1,) int (<Ops.ASSIGN: 66>, None)> on METAL with grad None>
+```
+
+The multiply Tensor stays the same because it is fused. The output Tensor's UOp becomes a new ASSIGN UOp:
+
+```py
+print(out.uop)
+```
+
+The first source is the output BUFFER:
+
+```
+UOp(Ops.BUFFER, dtypes.int, arg=1, src=(
+  UOp(Ops.DEVICE, dtypes.void, arg='METAL', src=()),
+  UOp(Ops.UNIQUE, dtypes.void, arg=6, src=()),))
+```
+
+And the second source is the KERNEL and its 4 buffer edges (output_buffer, a, b, c):
+
+```
+UOp(Ops.KERNEL, dtypes.void, arg=<Kernel 12 SINK(<Ops.STORE: 45>,) (__add__, __mul__)>, src=(
+  UOp(Ops.BUFFER, dtypes.int, arg=1, src=(
+    x1:=UOp(Ops.DEVICE, dtypes.void, arg='METAL', src=()),
+    UOp(Ops.UNIQUE, dtypes.void, arg=6, src=()),)),
+  UOp(Ops.BUFFER, dtypes.int, arg=1, src=(
+     x1,
+    UOp(Ops.UNIQUE, dtypes.void, arg=1, src=()),)),
+  UOp(Ops.BUFFER, dtypes.int, arg=1, src=(
+     x1,
+    UOp(Ops.UNIQUE, dtypes.void, arg=3, src=()),)),
+  UOp(Ops.BUFFER, dtypes.int, arg=1, src=(
+     x1,
+    UOp(Ops.UNIQUE, dtypes.void, arg=5, src=()),)),))
+```
+
+KERNEL describes the compute AST, metadata and memory dependencies.
+
+BUFFER holds a reference to the device memory where the output will be stored.
+
+Once a Tensor is kernelized, all children will LOAD its BUFFER, instead of fusing it:
+
+```py
+child = out+2
+child.kernelize()
+print(child.uop.src[1].arg.ast)
+```
+
+```
+UOp(Ops.SINK, dtypes.void, arg=None, src=(
+  UOp(Ops.STORE, dtypes.void, arg=None, src=(
+    UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(1), arg=0, src=()),
+    x2:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1,), strides=(0,), offset=0, mask=None, contiguous=True),)), src=()),
+    UOp(Ops.ADD, dtypes.int, arg=None, src=(
+      UOp(Ops.LOAD, dtypes.int, arg=None, src=(
+        UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(1), arg=1, src=()),
+         x2,)),
+      UOp(Ops.CONST, dtypes.int, arg=2, src=(
+         x2,)),)),)),))
+```
+
+`Tensor.realize` will execute the kernels and write outputs to memory:
+
+```py
+Tensor.realize(out)
+print(out)        # <Tensor <UOp METAL (1,) int (<Ops.BUFFER: 23>, <buf real:True device:METAL size:1 dtype:dtypes.int offset:0>)> on METAL with grad None>
+print(out.item()) # 5
+```
+
+<hr />
+
+**Summary**
+
+- The large Tensor graph is built from a mix of data, compute and movement Ops.
+
+- `Tensor.kernelize` splits the Tensor graph into data (BUFFER), compute (KERNEL) and links dependencies with ASSIGN.
+
+- `Tensor.realize` executes KERNELs on device and replaces the Tensor graph with just a BUFFER.
+
+- Kernelize can be called multiple times on a Tensor. This allows for incrementally building the kernel fusion layout of a large Tensor graph, without having to call `realize` or `schedule`.
@@ -68,4 +68,4 @@ We have a simple framework in tinygrad for adding these ALU blocks and achieving

 ### Indexing

-Indexing determines the address of the memory we need to load. GPUs often have less integer math resources than floating point math, so this can sometimes be the bottleneck. We have a symbolic math engine in our rewrite rules to simplifiy indexing before it's emitted to the kernel. Newer NVIDIA GPUs have a "Tensor Memory Accelerator" to assist with fast indexing, however, this is not supported in tinygrad yet.
+Indexing determines the address of the memory we need to load. GPUs often have less integer math resources than floating point math, so this can sometimes be the bottleneck. We have a symbolic math engine in our rewrite rules to simplify indexing before it's emitted to the kernel. Newer NVIDIA GPUs have a "Tensor Memory Accelerator" to assist with fast indexing, however, this is not supported in tinygrad yet.
@@ -1,10 +1,10 @@
-::: tinygrad.ops.UOp
+::: tinygrad.uop.ops.UOp
    options:
        members: false
        members_order: source
        show_labels: false

-::: tinygrad.ops.Ops
+::: tinygrad.uop.ops.Ops
    options:
        members: true
        members_order: source
@@ -30,13 +30,12 @@ These control the behavior of core tinygrad even when used as a library.

 Variable | Possible Value(s) | Description
 ---|---|---
-DEBUG               | [1-6]      | enable debugging output, with 4 you get operations, timings, speed, generated code and more
+DEBUG               | [1-7]      | enable debugging output (operations, timings, speed, generated code and more)
 GPU                 | [1]        | enable the GPU (OpenCL) backend
 CUDA                | [1]        | enable CUDA backend
 AMD                 | [1]        | enable AMD backend
 NV                  | [1]        | enable NV backend
 METAL               | [1]        | enable Metal backend (for Mac M1 and after)
-METAL_XCODE         | [1]        | enable Metal using macOS Xcode SDK
 CPU                 | [1]        | enable CPU (Clang) backend
 LLVM                | [1]        | enable LLVM backend
 BEAM                | [#]        | number of beams in kernel beam search
@@ -49,4 +48,17 @@ VISIBLE_DEVICES     | [list[int]]| restricts the NV/AMD devices that are availab
 JIT                 | [0-2]      | 0=disabled, 1=[jit enabled](quickstart.md#jit) (default), 2=jit enabled, but graphs are disabled
 VIZ                 | [1]        | 0=disabled, 1=[viz enabled](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/viz)
 ALLOW_TF32          | [1]        | enable TensorFloat-32 tensor cores on Ampere or newer GPUs.
-WEBGPU_BACKEND      | [WGPUBackendType_Metal, ...]          | Force select a backend for WebGPU (Metal, DirectX, OpenGL, Vulkan...)
+WEBGPU_BACKEND      | [WGPUBackendType_Metal, ...]          | Force select a backend for WebGPU (Metal, DirectX, OpenGL, Vulkan...)
+CUDA_PATH           | str        | Use `CUDA_PATH/include` for CUDA headers for CUDA and NV backends. If not set, TinyGrad will use `/usr/local/cuda/include`, `/usr/include` and `/opt/cuda/include`.
+
+## Debug breakdown
+
+Variable | Value | Description
+---|---|---
+DEBUG               | >= 1       | Enables debugging and lists devices being used
+DEBUG               | >= 2       | Provides performance metrics for operations, including timing, memory usage, bandwidth for each kernel execution
+DEBUG               | >= 3       | Outputs buffers used for each kernel (shape, dtype and strides) and the applied optimizations at a kernel level
+DEBUG               | >= 4       | Outputs the generated kernel code
+DEBUG               | >= 5       | Displays the intermediate representation of the computation UOps (AST)
+DEBUG               | >= 6       | Displays the intermediate representation of the computation UOps in a linearized manner, detailing the operation sequence
+DEBUG               | >= 7       | Outputs the assembly code generated for the target hardware
@@ -42,7 +42,7 @@ There's nothing special about a "Module" class in tinygrad, it's just a normal c

 ### tinygrad is functional

-In tinygrad, you can do [`x.conv2d(w, b)`](tensor/ops.md/#tinygrad.Tensor.conv2d) or [`x.sparse_categorical_cross_entropy(y)`](tensor/ops.md/#tinygrad.Tensor.sparse_categorical_crossentropy). We do also have a [`Conv2D`](nn.md/#tinygrad.nn.Conv2d) class like PyTorch if you want a place to keep the state, but all stateless operations don't have classes.
+In tinygrad, you can do [`x.conv2d(w, b)`](tensor/ops.md/#tinygrad.Tensor.conv2d) or [`x.sparse_categorical_crossentropy(y)`](tensor/ops.md/#tinygrad.Tensor.sparse_categorical_crossentropy). We do also have a [`Conv2D`](nn.md/#tinygrad.nn.Conv2d) class like PyTorch if you want a place to keep the state, but all stateless operations don't have classes.

 ### tinygrad is lazy

@@ -0,0 +1,293 @@
+#!/usr/bin/env python3
+
+# this file is a "ramp" for people new to tinygrad to think about how to approach it
+# it is runnable and editable.
+# whenever you see stuff like DEBUG=2 or CPU=1 discussed, these are environment variables
+# in a unix shell like bash `DEBUG=2 CPU=1 python docs/ramp.py`
+
+# this pip installs tinygrad master for the system
+# the -e allows you to edit the tinygrad folder and update system tinygrad
+# tinygrad is pure Python, so you are encouraged to do this
+# git pull in the tinygrad directory will also get you the latest
+"""
+git clone https://github.com/tinygrad/tinygrad.git
+cd tinygrad
+python3 -m pip install -e .
+"""
+
+# %% ********
+print("******* PART 1 *******")
+
+# we start with a Device.
+# a Device is where Tensors are stored and compute is run
+# tinygrad autodetects the best device on your system and makes it the DEFAULT
+from tinygrad import Device
+print(Device.DEFAULT)  # on Mac, you can see this prints METAL
+
+# now, lets create a Tensor
+from tinygrad import Tensor, dtypes
+t = Tensor([1,2,3,4])
+
+# you can see this Tensor is on the DEFAULT device with int dtype and shape (4,)
+assert t.device == Device.DEFAULT
+assert t.dtype == dtypes.int
+assert t.shape == (4,)
+
+# unlike in torch, if we print it, it doesn't print the contents
+# this is because tinygrad is lazy
+# this Tensor has not been computed yet
+print(t)
+# <Tensor <UOp METAL (4,) int (<Ops.COPY: 7>, None)> on METAL with grad None>
+
+# the ".uop" property on Tensor contains the specification of how to compute it
+print(t.uop)
+"""
+UOp(Ops.COPY, dtypes.int, arg=None, src=(
+  UOp(Ops.BUFFER, dtypes.int, arg=4, src=(
+    UOp(Ops.UNIQUE, dtypes.void, arg=0, src=()),
+    UOp(Ops.DEVICE, dtypes.void, arg='PYTHON', src=()),)),
+  UOp(Ops.DEVICE, dtypes.void, arg='METAL', src=()),))
+"""
+# as you can see, it's specifying a copy from PYTHON device
+# which is where the [1,2,3,4] array lives
+
+# UOps are the specification language in tinygrad
+# they are immutable and form a DAG
+# they have a "Ops", a "dtype", a tuple of srcs (parents), and an arg
+
+t.realize()
+# if we want to "realize" a tensor, we can with the "realize" method
+# now when we look at the uop, it's changed
+print(t.uop)
+"""
+UOp(Ops.BUFFER, dtypes.int, arg=4, src=(
+  UOp(Ops.UNIQUE, dtypes.void, arg=1, src=()),
+  UOp(Ops.DEVICE, dtypes.void, arg='METAL', src=()),))
+"""
+# the copy was actually run, and now the "uop" of the Tensor is just a BUFFER
+# if you run this script with DEBUG=2 in the environment, you can see the copy happen
+# *** METAL      1 copy       16,   METAL <- PYTHON ...
+
+# now let's do some compute
+# we look at the uop to see the specification of the compute
+t_times_2 = t * 2
+print(t_times_2.uop)
+"""
+UOp(Ops.MUL, dtypes.int, arg=None, src=(
+  UOp(Ops.BUFFER, dtypes.int, arg=4, src=(
+    UOp(Ops.UNIQUE, dtypes.void, arg=1, src=()),
+    x2:=UOp(Ops.DEVICE, dtypes.void, arg='METAL', src=()),)),
+  UOp(Ops.EXPAND, dtypes.int, arg=(4,), src=(
+    UOp(Ops.RESHAPE, dtypes.int, arg=(1,), src=(
+      UOp(Ops.CONST, dtypes.int, arg=2, src=(
+        UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(), strides=(), offset=0, mask=None, contiguous=True),)), src=(
+           x2,)),)),)),)),))
+"""
+# the BUFFER from above is being multiplied by a CONST 2
+# it's RESHAPEd and EXPANDed to broadcast the CONST to the BUFFER
+
+# we can check the result with
+assert t_times_2.tolist() == [2, 4, 6, 8]
+
+# UOps are both immutable and globally unique
+# if i multiply the Tensor by 4 twice, these result Tensors will have the same uop specification
+t_times_4_try_1 = t * 4
+t_times_4_try_2 = t * 4
+assert t_times_4_try_1.uop is t_times_4_try_2.uop
+# the specification isn't just the same, it's the exact same Python object
+assert t_times_4_try_1 is not t_times_4_try_2
+# the Tensor is a different Python object
+
+# if we realize `t_times_4_try_1` ...
+t_times_4_try_1.realize()
+print(t_times_4_try_2.uop)
+"""
+UOp(Ops.BUFFER, dtypes.int, arg=4, src=(
+  UOp(Ops.UNIQUE, dtypes.void, arg=4, src=()),
+  UOp(Ops.DEVICE, dtypes.void, arg='METAL', src=()),))
+"""
+# ... `t_times_4_try_2` also becomes the same BUFFER
+assert t_times_4_try_1.uop is t_times_4_try_2.uop
+# so this print doesn't require any computation, just a copy back to the CPU so we can print it
+print("** only the copy start")
+print(t_times_4_try_2.tolist())  # [4, 8, 12, 16]
+print("** only the copy end")
+# you can confirm this with DEBUG=2, seeing what's printed in between the "**" prints
+
+# tinygrad has an auto differentiation engine that operates according to these same principles
+# the derivative of "log(x)" is "1/x", and you can see this on line 20 of gradient.py
+t_float = Tensor([3.0])
+t_log = t_float.log()
+t_log_grad, = t_log.sum().gradient(t_float)
+# due to how log is implemented, this gradient contains a lot of UOps
+print(t_log_grad.uop)
+# ...not shown here...
+# but if you run with DEBUG=4 (CPU=1 used here for simpler code), you can see the generated code
+"""
+void E_(float* restrict data0, float* restrict data1) {
+  float val0 = *(data1+0);
+  *(data0+0) = (0.6931471805599453f*(1/(val0*0.6931471805599453f)));
+}
+"""
+# the derivative is close to 1/3
+assert (t_log_grad.item() - 1/3) < 1e-6
+
+# %% ********
+print("******* PART 2 *******")
+
+# we redefine the same t here so this cell can run on it's own
+from tinygrad import Tensor
+t = Tensor([1,2,3,4])
+
+# what's above gives you enough of an understanding to go use tinygrad as a library
+# however, a lot of the beauty of tinygrad is in how easy it is to interact with the internals
+# NOTE: the APIs here are subject to change
+
+t_plus_3_plus_4 = t + 3 + 4
+print(t_plus_3_plus_4.uop)
+"""
+UOp(Ops.ADD, dtypes.int, arg=None, src=(
+  UOp(Ops.ADD, dtypes.int, arg=None, src=(
+    UOp(Ops.BUFFER, dtypes.int, arg=4, src=(
+      UOp(Ops.UNIQUE, dtypes.void, arg=1, src=()),
+      x3:=UOp(Ops.DEVICE, dtypes.void, arg='CPU', src=()),)),
+    UOp(Ops.EXPAND, dtypes.int, arg=(4,), src=(
+      UOp(Ops.RESHAPE, dtypes.int, arg=(1,), src=(
+        UOp(Ops.CONST, dtypes.int, arg=3, src=(
+          x7:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(), strides=(), offset=0, mask=None, contiguous=True),)), src=(
+             x3,)),)),)),)),)),
+  UOp(Ops.EXPAND, dtypes.int, arg=(4,), src=(
+    UOp(Ops.RESHAPE, dtypes.int, arg=(1,), src=(
+      UOp(Ops.CONST, dtypes.int, arg=4, src=(
+         x7,)),)),)),))
+"""
+# you can see it's adding both 3 and 4
+
+# but by the time we are actually running the code, it's adding 7
+# `kernelize` will simplify and group the operations in the graph into kernels
+t_plus_3_plus_4.kernelize()
+print(t_plus_3_plus_4.uop)
+"""
+UOp(Ops.ASSIGN, dtypes.int, arg=None, src=(
+  x0:=UOp(Ops.BUFFER, dtypes.int, arg=4, src=(
+    UOp(Ops.UNIQUE, dtypes.void, arg=7, src=()),
+    x2:=UOp(Ops.DEVICE, dtypes.void, arg='CPU', src=()),)),
+  UOp(Ops.KERNEL, dtypes.void, arg=<Kernel 12 SINK(<Ops.STORE: 48>,) (__add__,)>, src=(
+     x0,
+    UOp(Ops.BUFFER, dtypes.int, arg=4, src=(
+      UOp(Ops.UNIQUE, dtypes.void, arg=1, src=()),
+       x2,)),)),))
+"""
+# ASSIGN has two srcs, src[0] is the BUFFER that's assigned to, and src[1] is the thing to assign
+# src[1] is the GPU Kernel that's going to be run
+# we can get the ast of the Kernel as follows
+kernel_ast = t_plus_3_plus_4.uop.src[1].arg.ast
+
+# almost everything in tinygrad functions as a rewrite of the UOps
+# the codegen rewrites the ast to a simplified form ready for "rendering"
+from tinygrad.codegen import full_rewrite_to_sink
+rewritten_ast = full_rewrite_to_sink(kernel_ast)
+print(rewritten_ast)
+"""
+UOp(Ops.SINK, dtypes.void, arg=None, src=(
+  UOp(Ops.STORE, dtypes.void, arg=None, src=(
+    UOp(Ops.INDEX, dtypes.int.ptr(4), arg=None, src=(
+      UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(4), arg=0, src=()),
+      x3:=UOp(Ops.SPECIAL, dtypes.int, arg=('gidx0', 4), src=()),)),
+    UOp(Ops.ADD, dtypes.int, arg=None, src=(
+      UOp(Ops.LOAD, dtypes.int, arg=None, src=(
+        UOp(Ops.INDEX, dtypes.int.ptr(4), arg=None, src=(
+          UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(4), arg=1, src=()),
+           x3,)),)),
+      UOp(Ops.CONST, dtypes.int, arg=7, src=()),)),)),))
+"""
+# you can see at this point we are adding 7, not 3 and 4
+
+# with DEBUG=4, we can see the code.
+# since optimizations are on, it UPCASTed the operation, explicitly writing out all 4 +7s
+t_plus_3_plus_4.realize()
+"""
+void E_4n2(int* restrict data0, int* restrict data1) {
+  int val0 = *(data1+0);
+  int val1 = *(data1+1);
+  int val2 = *(data1+2);
+  int val3 = *(data1+3);
+  *(data0+0) = (val0+7);
+  *(data0+1) = (val1+7);
+  *(data0+2) = (val2+7);
+  *(data0+3) = (val3+7);
+}
+"""
+# the function name E_4n2 is "E" for elementwise op (as opposed to "r" for reduce op)
+# "4" for the size, and "n2" for name deduping (it's the 3rd function with the same E and 4 in this session)
+# when you print the name with DEBUG=2, you'll see the 4 is yellow, meaning that it's upcasted
+# if you run with NOOPT=1 ...
+"""
+void E_4n2(int* restrict data0, int* restrict data1) {
+  for (int ridx0 = 0; ridx0 < 4; ridx0++) {
+    int val0 = *(data1+ridx0);
+    *(data0+ridx0) = (val0+7);
+  }
+}
+"""
+# ... you get this unoptimized code with a loop and the 4 is blue (for global). the color code is in kernel.py
+
+# %% ********
+print("******* PART 3 *******")
+
+# now, we go even lower and understand UOps better and how the graph rewrite engine works.
+# it's much simpler than what's in LLVM or MLIR
+
+from tinygrad import dtypes
+from tinygrad.uop import UOp, Ops
+
+# first, we'll construct some const UOps
+a = UOp(Ops.CONST, dtypes.int, arg=2)
+b = UOp(Ops.CONST, dtypes.int, arg=2)
+
+# if you have been paying attention, you should know these are the same Python object
+assert a is b
+
+# UOps support normal Python math operations, so a_plus_b expresses the spec for 2 + 2
+a_plus_b = a + b
+print(a_plus_b)
+"""
+UOp(Ops.ADD, dtypes.int, arg=None, src=(
+  x0:=UOp(Ops.CONST, dtypes.int, arg=2, src=()),
+   x0,))
+"""
+
+# we could actually render this 2+2 into a language like c and run it
+# or, we can use tinygrad's graph rewrite engine to "constant fold"
+
+from tinygrad.uop.ops import graph_rewrite, UPat, PatternMatcher
+
+# a `PatternMatcher` is a list of tuples. for each element in the list:
+# [0] is the pattern to match, and [1] is the function to run.
+# this function can return either a UOp to replace the pattern with, or None to not replace
+simple_pm = PatternMatcher([
+  (UPat(Ops.ADD, src=(UPat(Ops.CONST, name="c1"), UPat(Ops.CONST, name="c2"))),
+   lambda c1,c2: UOp(Ops.CONST, dtype=c1.dtype, arg=c1.arg+c2.arg)),
+])
+# this pattern matches the addition of two CONST and rewrites it into a single CONST UOp
+
+# to actually apply the pattern to a_plus_b, we use graph_rewrite
+a_plus_b_simplified = graph_rewrite(a_plus_b, simple_pm)
+print(a_plus_b_simplified)
+"""
+UOp(Ops.CONST, dtypes.int, arg=4, src=())
+"""
+# 2+2 is in fact, 4
+
+# we can also use syntactic sugar to write the pattern nicer
+simpler_pm = PatternMatcher([
+  (UPat.cvar("c1")+UPat.cvar("c2"), lambda c1,c2: c1.const_like(c1.arg+c2.arg))
+])
+assert graph_rewrite(a_plus_b, simple_pm) is graph_rewrite(a_plus_b, simpler_pm)
+# note again the use of is, UOps are immutable and globally unique
+
+# %% ********
+
+# that brings you to an understanding of the most core concepts in tinygrad
+# you can run this with VIZ=1 to use the web based graph rewrite explorer
+# hopefully now you understand it. the nodes in the graph are just UOps
@@ -5,7 +5,7 @@ tinygrad supports various runtimes, enabling your code to scale across a wide ra
 | Runtime | Description | Requirements |
 |---------|-------------|--------------|
 | [NV](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_nv.py) | Provides acceleration for NVIDIA GPUs | Ampere/Ada series GPUs |
-| [AMD](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_amd.py) | Provides acceleration for AMD GPUs | RDNA2/RDNA3 series GPUs |
+| [AMD](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_amd.py) | Provides acceleration for AMD GPUs | RDNA2/RDNA3/RDNA4 series GPUs. You can select one of the interfaces for communication by setting `AMD_IFACE=(KFD|PCI)`. See [AMD interfaces](#amd-interfaces) for more details. |
 | [QCOM](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_qcom.py) | Provides acceleration for QCOM GPUs | 6xx series GPUs |
 | [METAL](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_metal.py) | Utilizes Metal for acceleration on Apple devices | M1+ Macs; Metal 3.0+ for `bfloat` support |
 | [CUDA](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cuda.py) | Utilizes CUDA for acceleration on NVIDIA GPUs | NVIDIA GPU with CUDA support |
@@ -64,3 +64,11 @@ rawbuf_ptr = to_mv(cl_buf_desc_ptr, 0x100).cast('Q')[20] # offset 0xA0 is a raw
 # create tiny tensor
 tiny = Tensor.from_blob(rawbuf_ptr, (h*w*4,), dtype=dtypes.imagef((h,w)), device='QCOM')
 ```
+
+## AMD Interfaces
+AMD backend supports several interfaces for communicating with devices:
+
+* `KFD`: uses the amdgpu driver
+* `PCI`: uses the [AM driver](developer/am.md)
+
+You can force an interface by setting `AMD_IFACE` to one of these values. In the case of `AMD_IFACE=PCI`, this may unbind your GPU from the amdgpu driver.
@@ -22,7 +22,9 @@
 ::: tinygrad.Tensor.rand
 ::: tinygrad.Tensor.rand_like
 ::: tinygrad.Tensor.randn
+::: tinygrad.Tensor.randn_like
 ::: tinygrad.Tensor.randint
+::: tinygrad.Tensor.randperm
 ::: tinygrad.Tensor.normal
 ::: tinygrad.Tensor.uniform
 ::: tinygrad.Tensor.scaled_uniform
@@ -18,6 +18,7 @@
 ::: tinygrad.Tensor.repeat_interleave
 ::: tinygrad.Tensor.split
 ::: tinygrad.Tensor.chunk
+::: tinygrad.Tensor.unfold
 ::: tinygrad.Tensor.meshgrid
 ::: tinygrad.Tensor.squeeze
 ::: tinygrad.Tensor.unsqueeze
@@ -37,8 +37,10 @@
 ::: tinygrad.Tensor.scatter
 ::: tinygrad.Tensor.scatter_reduce
 ::: tinygrad.Tensor.masked_select
+::: tinygrad.Tensor.masked_fill
 ::: tinygrad.Tensor.sort
 ::: tinygrad.Tensor.topk
+::: tinygrad.Tensor.multinomial

 ## Neural Network (functional)

@@ -78,10 +78,7 @@ if __name__ == "__main__":

  @TinyJit
  def get_action(obs:Tensor) -> Tensor:
-    # TODO: with no_grad
-    Tensor.no_grad = True
    ret = model(obs)[0].exp().multinomial().realize()
-    Tensor.no_grad = False
    return ret

  st, steps = time.perf_counter(), 0
@@ -3,14 +3,19 @@ start_tm = time.perf_counter()
 import math
 from typing import Tuple, cast
 import numpy as np
-from tinygrad import Tensor, nn, GlobalCounters, TinyJit, dtypes
+from tinygrad import Tensor, nn, GlobalCounters, TinyJit, dtypes, Device
 from tinygrad.helpers import partition, trange, getenv, Context
 from extra.lr_scheduler import OneCycleLR

+GPUS = [f'{Device.DEFAULT}:{i}' for i in range(getenv("GPUS", 1))]
+
+# override tinygrad defaults
 dtypes.default_float = dtypes.half
+Context(FUSE_ARANGE=1, FUSE_OPTIM=1).__enter__()

 # from https://github.com/tysam-code/hlb-CIFAR10/blob/main/main.py
 batchsize = getenv("BS", 1024)
+assert batchsize % len(GPUS) == 0, f"{batchsize=} is not a multiple of {len(GPUS)=}"
 bias_scaler = 64
 hyp = {
  'opt': {
@@ -67,7 +72,7 @@ class ConvGroup:
    cast(Tensor, self.norm2.weight).requires_grad = False
  def __call__(self, x:Tensor) -> Tensor:
    x =    self.norm1(self.conv1(x).max_pool2d().float()).cast(dtypes.default_float).quick_gelu()
-    return self.norm2(self.conv2(x).float()).cast(dtypes.default_float).quick_gelu()
+    return self.norm2(self.conv2(x).float()).cast(dtypes.default_float).quick_gelu() + x

 class SpeedyConvNet:
  def __init__(self):
@@ -78,23 +83,25 @@ class SpeedyConvNet:
    self.linear = nn.Linear(depths['block3'], depths['num_classes'], bias=False)
  def __call__(self, x:Tensor) -> Tensor:
    x = self.whiten(x).quick_gelu()
+    # ************* HACKS *************
+    x = x.pad((1,0,0,1)) # TODO: this pad should not be here! copied from hlb_cifar10 for speed
+    # ************* HACKS *************
    x = x.sequential([self.conv_group_1, self.conv_group_2, self.conv_group_3])
    return self.linear(x.max(axis=(2,3))) * hyp['opt']['scaling_factor']

 if __name__ == "__main__":
  # *** dataset ***
  X_train, Y_train, X_test, Y_test = nn.datasets.cifar()
-  # TODO: without this line indexing doesn't fuse!
-  X_train, Y_train, X_test, Y_test = [x.contiguous() for x in [X_train, Y_train, X_test, Y_test]]
  cifar10_std, cifar10_mean = X_train.float().std_mean(axis=(0, 2, 3))
-  def preprocess(X:Tensor, Y:Tensor) -> Tuple[Tensor, Tensor]:
-    return ((X - cifar10_mean.view(1, -1, 1, 1)) / cifar10_std.view(1, -1, 1, 1)).cast(dtypes.default_float), Y.one_hot(depths['num_classes'])
+  def preprocess(X:Tensor) -> Tensor: return ((X - cifar10_mean.view(1, -1, 1, 1)) / cifar10_std.view(1, -1, 1, 1)).cast(dtypes.default_float)

  # *** model ***
  model = SpeedyConvNet()
  state_dict = nn.state.get_state_dict(model)
-
-  #for k,v in nn.state.torch_load("/tmp/cifar_net.pt").items(): print(k)
+  if len(GPUS) > 1:
+    cifar10_std.to_(GPUS)
+    cifar10_mean.to_(GPUS)
+    for x in state_dict.values(): x.to_(GPUS)

  params_bias, params_non_bias = partition(state_dict.items(), lambda x: 'bias' in x[0])
  opt_bias     = nn.optim.SGD([x[1] for x in params_bias],     lr=0.01, momentum=.85, nesterov=True, weight_decay=hyp['opt']['bias_decay'])
@@ -111,40 +118,37 @@ if __name__ == "__main__":
  lr_sched_bias     = OneCycleLR(opt_bias,     max_lr=hyp['opt']['bias_lr'],     pct_start=pct_start, div_factor=initial_div_factor, final_div_factor=1./(initial_div_factor*final_lr_ratio), total_steps=total_train_steps)
  lr_sched_non_bias = OneCycleLR(opt_non_bias, max_lr=hyp['opt']['non_bias_lr'], pct_start=pct_start, div_factor=initial_div_factor, final_div_factor=1./(initial_div_factor*final_lr_ratio), total_steps=total_train_steps)

-  def loss_fn(out, Y):
-    return out.cross_entropy(Y, reduction='none', label_smoothing=0.2).mul(hyp['opt']['loss_scale_scaler']*loss_batchsize_scaler).sum().div(hyp['opt']['loss_scale_scaler'])
+  def loss_fn(out:Tensor, Y:Tensor) -> Tensor:
+    ret = out.sparse_categorical_crossentropy(Y, reduction='none', label_smoothing=0.2)
+    return ret.mul(hyp['opt']['loss_scale_scaler']*loss_batchsize_scaler).sum().div(hyp['opt']['loss_scale_scaler'])

  @TinyJit
  @Tensor.train()
  def train_step(idxs:Tensor) -> Tensor:
-    with Context(SPLIT_REDUCEOP=0, FUSE_ARANGE=1):
-      X = X_train[idxs]
-      Y = Y_train[idxs].realize(X)
-    X, Y = preprocess(X, Y)
-    out = model(X)
+    X, Y = X_train[idxs], Y_train[idxs]
+    if len(GPUS) > 1:
+      X.shard_(GPUS, axis=0)
+      Y.shard_(GPUS, axis=0)
+    out = model(preprocess(X))
    loss = loss_fn(out, Y)
    opt.zero_grad()
    loss.backward()
-    opt.step()
-    lr_sched_bias.step()
-    lr_sched_non_bias.step()
-    return loss / (batchsize*loss_batchsize_scaler)
+    return (loss / (batchsize*loss_batchsize_scaler)).realize(*opt.schedule_step(),
+                                                              *lr_sched_bias.schedule_step(), *lr_sched_non_bias.schedule_step())

  eval_batchsize = 2500
  @TinyJit
-  @Tensor.test()
  def val_step() -> Tuple[Tensor, Tensor]:
-    # TODO with Tensor.no_grad()
-    Tensor.no_grad = True
    loss, acc = [], []
    for i in range(0, X_test.size(0), eval_batchsize):
-      X, Y = preprocess(X_test[i:i+eval_batchsize], Y_test[i:i+eval_batchsize])
-      out = model(X)
+      X, Y = X_test[i:i+eval_batchsize], Y_test[i:i+eval_batchsize]
+      if len(GPUS) > 1:
+        X.shard_(GPUS, axis=0)
+        Y.shard_(GPUS, axis=0)
+      out = model(preprocess(X))
      loss.append(loss_fn(out, Y))
-      acc.append((out.argmax(-1).one_hot(depths['num_classes']) * Y).sum() / eval_batchsize)
-    ret = Tensor.stack(*loss).mean() / (batchsize*loss_batchsize_scaler), Tensor.stack(*acc).mean()
-    Tensor.no_grad = False
-    return ret
+      acc.append((out.argmax(-1) == Y).sum() / eval_batchsize)
+    return Tensor.stack(*loss).mean() / (batchsize*loss_batchsize_scaler), Tensor.stack(*acc).mean()

  np.random.seed(1337)
  for epoch in range(math.ceil(hyp['misc']['train_epochs'])):
@@ -1,4 +1,4 @@
-# model based off https://towardsdatascience.com/going-beyond-99-mnist-handwritten-digits-recognition-cfff96337392
+# model based off https://medium.com/data-science/going-beyond-99-mnist-handwritten-digits-recognition-cfff96337392
 from typing import List, Callable
 from tinygrad import Tensor, TinyJit, nn, GlobalCounters
 from tinygrad.helpers import getenv, colored, trange
@@ -34,7 +34,6 @@ if __name__ == "__main__":
    return loss

  @TinyJit
-  @Tensor.test()
  def get_test_acc() -> Tensor: return (model(X_test).argmax(axis=1) == Y_test).mean()*100

  test_acc = float('nan')
@@ -1,12 +1,12 @@
-import sys, onnx, time, pickle
-from tinygrad import TinyJit, Device, GlobalCounters, fetch, getenv
-from tinygrad.frontend.onnx import OnnxRunner
+import sys, time, pickle
+from tinygrad import TinyJit, GlobalCounters, fetch, getenv
+from tinygrad.frontend.onnx import OnnxRunner, onnx_load
 from extra.onnx_helpers import get_example_inputs, validate

 def load_onnx_model(onnx_file):
-  onnx_model = onnx.load(onnx_file)
+  onnx_model = onnx_load(onnx_file)
  run_onnx = OnnxRunner(onnx_model)
-  run_onnx_jit = TinyJit(lambda **kwargs: next(iter(run_onnx({k:v.to(Device.DEFAULT) for k,v in kwargs.items()}).values())), prune=True)
+  run_onnx_jit = TinyJit(lambda **kwargs: next(iter(run_onnx({k:v.to(None) for k,v in kwargs.items()}).values())), prune=True, optimize=True)
  return run_onnx_jit, run_onnx.graph_inputs

 if __name__ == "__main__":
@@ -34,8 +34,3 @@ if __name__ == "__main__":
  if getenv("ORT"):
    validate(onnx_file, new_inputs, rtol=1e-3, atol=1e-3)
    print("model validated")
-
-  if (fn:=getenv("SAVE_PKL", "")) != "":
-    with open(fn, "wb") as f:
-      pickle.dump(run_onnx_jit, f)
-    print(f"pkl saved to {fn}")
@@ -23,8 +23,6 @@ def create_fixed_tokenizer(output_file):
 # echo -en "write 2+2\nwrite hello world\ny\n" | TEMP=0 python3 examples/coder.py

 if __name__ == "__main__":
-  Tensor.no_grad = True
-
  # https://huggingface.co/teknium/OpenHermes-2.5-Mistral-7B/blob/main/config.json
  with Timing("create model: "):
    model = Transformer(4096, 14336, n_heads=32, n_layers=32, norm_eps=1e-5, vocab_size=32002, n_kv_heads=8, max_context=4096, jit=getenv("JIT", 1))
@@ -34,8 +32,8 @@ if __name__ == "__main__":
    part2 = nn.state.torch_load(fetch("https://huggingface.co/teknium/OpenHermes-2.5-Mistral-7B/resolve/main/pytorch_model-00002-of-00002.bin?download=true"))

  with Timing("weights -> model: "):
-    nn.state.load_state_dict(model, fix_bf16(convert_from_huggingface(part1, model, 32, 8)), strict=False)
-    nn.state.load_state_dict(model, fix_bf16(convert_from_huggingface(part2, model, 32, 8)), strict=False)
+    nn.state.load_state_dict(model, fix_bf16(convert_from_huggingface(part1, 32, 32, 8)), strict=False)
+    nn.state.load_state_dict(model, fix_bf16(convert_from_huggingface(part2, 32, 32, 8)), strict=False)

  if not os.path.isfile("/tmp/tokenizer.model"): create_fixed_tokenizer("/tmp/tokenizer.model")
  spp = SentencePieceProcessor(model_file="/tmp/tokenizer.model")
@@ -159,7 +159,6 @@ def init_vits(
  text_mapper = TextMapper(apply_cleaners=True, symbols=symbols)

  # Load the model.
-  Tensor.no_grad = True
  if seed is not None:
    Tensor.manual_seed(seed)
    np.random.seed(seed)
@@ -221,7 +220,6 @@ def mp_output_stream(q: mp.Queue, counter: mp.Value, num_channels: int, sample_r
 if __name__ == "__main__":
  import nltk
  nltk.download("punkt")
-  Tensor.no_grad = True
  # Parse CLI arguments
  parser = argparse.ArgumentParser("Have a tiny conversation with tinygrad")

@@ -3,10 +3,11 @@ import os, argparse, contextlib
 from typing import Optional, Union
 with contextlib.suppress(ImportError): import tiktoken
 from tinygrad import Tensor, TinyJit, Device, GlobalCounters, Variable, dtypes
-from tinygrad.ops import UOp
+from tinygrad.uop.ops import UOp
 from tinygrad.helpers import Timing, DEBUG, JIT, getenv, fetch, colored, trange
 from tinygrad.nn import Embedding, Linear, LayerNorm
 from tinygrad.nn.state import gguf_load, torch_load, load_state_dict, get_state_dict
+from extra.bench_log import BenchEvent, WallTimeEvent

 MAX_CONTEXT = getenv("MAX_CONTEXT", 128)
 HALF = getenv("HALF")
@@ -84,7 +85,10 @@ class Transformer:
      seqlen = tokens.shape[1]
      tok_emb = self.wte(tokens)

-    pos_emb = self.wpe(self.allpos.shrink((None, (start_pos, start_pos+seqlen))))
+    # not symbolic when consuming the prompt
+    selected_pos = (0, seqlen) if start_pos.val == 0 else (start_pos, start_pos+1)
+    pos_emb = self.wpe(self.allpos.shrink((None, selected_pos)))
+
    h = tok_emb + pos_emb

    if HALF: h = h.half()
@@ -134,11 +138,12 @@ class GPT2:
    # lm head and wte are tied
    weights['lm_head.weight'] = weights['wte.weight']

-    load_state_dict(model, weights)
+    with WallTimeEvent(BenchEvent.LOAD_WEIGHTS):
+      load_state_dict(model, weights)

-    if HALF:
-      for l in get_state_dict(model).values():
-        l.replace(l.half().realize())
+      if HALF:
+        for l in get_state_dict(model).values():
+          l.replace(l.half().realize())

    return GPT2(model, tokenizer)

@@ -167,7 +172,8 @@ class GPT2:
      return key
    state_dict = { _remap_gguf_key(k): v for k, v in state_dict.items() }
    model = Transformer(**gpt2_params)
-    load_state_dict(model, state_dict)
+    with WallTimeEvent(BenchEvent.LOAD_WEIGHTS):
+      load_state_dict(model, state_dict)
    return GPT2(model, tiktoken.get_encoding("gpt2"))

  def __init__(self, model, tokenizer):
@@ -185,11 +191,12 @@ class GPT2:
      with Timing("ran model in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on GPU" if DEBUG>=2 else "")+
                  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>=2 else "")) if DEBUG else None, enabled=timing):
-        if batch_size == 1 and len(toks[0][start_pos:]) == 1:
-          tokens = Variable("tokens", 0, VOCAB_SIZE).bind(toks[0][start_pos])
-        else:
-          tokens = Tensor([x[start_pos:] for x in toks])
-        tok = self.model(tokens, Variable("start_pos", 1 if start_pos else 0, MAX_CONTEXT).bind(start_pos), temperature).tolist()
+        with WallTimeEvent(BenchEvent.STEP):
+          if batch_size == 1 and len(toks[0][start_pos:]) == 1:
+            tokens = Variable("tokens", 0, VOCAB_SIZE-1).bind(toks[0][start_pos])
+          else:
+            tokens = Tensor([x[start_pos:] for x in toks])
+          tok = self.model(tokens, Variable("start_pos", 1 if start_pos else 0, MAX_CONTEXT-1).bind(start_pos), temperature).tolist()
      start_pos = len(toks[0])
      for i,t in enumerate(tok): toks[i].append(t)
    return [self.tokenizer.decode(x) for x in toks]
@@ -197,7 +204,6 @@ class GPT2:
 # **** main code ****

 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?"

@@ -3,7 +3,8 @@ from extra.mcts_search import mcts_search
 from examples.mlperf.helpers import get_mlperf_bert_model
 from tinygrad import Tensor, Device, dtypes, nn
 from tinygrad.codegen.kernel import Kernel
-from tinygrad.ops import Ops, sym_infer
+from tinygrad.codegen.heuristic import hand_coded_optimizations
+from tinygrad.uop.ops import Ops, sym_infer
 from tinygrad.device import Compiled
 from tinygrad.engine.search import beam_search, bufs_from_lin
 from tinygrad.helpers import DEBUG, ansilen, getenv, colored, TRACEMETA
@@ -83,7 +84,7 @@ if __name__ == "__main__":

    # always try hand coded opt
    lin = Kernel(si.ast, opts=device.renderer)
-    lin.hand_coded_optimizations()
+    lin.apply_opts(hand_coded_optimizations(lin))
    lins.append((lin, "HC"))

    # maybe try tensor cores
@@ -11,6 +11,7 @@ from tinygrad import nn, dtypes, Tensor, Device, GlobalCounters, TinyJit
 from tinygrad.nn.state import get_state_dict, get_parameters
 from tinygrad.nn import optim
 from tinygrad.helpers import Context, BEAM, WINO, getenv, colored, prod
+from extra.bench_log import BenchEvent, WallTimeEvent

 cifar_mean = [0.4913997551666284, 0.48215855929893703, 0.4465309133731618]
 cifar_std = [0.24703225141799082, 0.24348516474564, 0.26158783926049628]
@@ -117,7 +118,7 @@ class SpeedyResNet:
 # hyper-parameters were exactly the same as the original repo
 bias_scaler = 58
 hyp = {
-  'seed' : 209,
+  'seed' : 200,
  'opt': {
    'bias_lr':            1.76 * bias_scaler/512,
    'non_bias_lr':        1.76 / 512,
@@ -266,13 +267,10 @@ def train_cifar():

    @TinyJit
    def update(self, net, decay):
-      # TODO with Tensor.no_grad()
-      Tensor.no_grad = True
      for net_ema_param, (param_name, net_param) in zip(get_state_dict(self.net_ema).values(), get_state_dict(net).items()):
        # batchnorm currently is not being tracked
        if not ("num_batches_tracked" in param_name) and not ("running" in param_name):
          net_ema_param.assign(net_ema_param.detach()*decay + net_param.detach()*(1.-decay)).realize()
-      Tensor.no_grad = False

  set_seed(getenv('SEED', hyp['seed']))

@@ -395,20 +393,23 @@ def train_cifar():
      if STEPS == 0 or i == STEPS: break

      GlobalCounters.reset()
-      X, Y = next(batcher)
-      if len(GPUS) > 1:
-        X.shard_(GPUS, axis=0)
-        Y.shard_(GPUS, axis=0)

-      with Context(BEAM=getenv("LATEBEAM", BEAM.value), WINO=getenv("LATEWINO", WINO.value)):
-        loss = train_step_jitted(model, optim.OptimizerGroup(opt_bias, opt_non_bias), [lr_sched_bias, lr_sched_non_bias], X, Y)
-        et = time.monotonic()
-        loss_cpu = loss.numpy()
-      # EMA for network weights
-      if getenv("EMA") and i > hyp['ema']['steps'] and (i+1) % hyp['ema']['every_n_steps'] == 0:
-        if model_ema is None:
-          model_ema = modelEMA(W, model)
-        model_ema.update(model, Tensor([projected_ema_decay_val*(i/STEPS)**hyp['ema']['decay_pow']]))
+      with WallTimeEvent(BenchEvent.STEP):
+        X, Y = next(batcher)
+        if len(GPUS) > 1:
+          X.shard_(GPUS, axis=0)
+          Y.shard_(GPUS, axis=0)
+
+        with Context(BEAM=getenv("LATEBEAM", BEAM.value), WINO=getenv("LATEWINO", WINO.value)):
+          loss = train_step_jitted(model, optim.OptimizerGroup(opt_bias, opt_non_bias), [lr_sched_bias, lr_sched_non_bias], X, Y)
+          et = time.monotonic()
+          loss_cpu = loss.numpy()
+        # EMA for network weights
+        if getenv("EMA") and i > hyp['ema']['steps'] and (i+1) % hyp['ema']['every_n_steps'] == 0:
+          if model_ema is None:
+            model_ema = modelEMA(W, model)
+          model_ema.update(model, Tensor([projected_ema_decay_val*(i/STEPS)**hyp['ema']['decay_pow']]))
+
      cl = time.monotonic()
      device_str = loss.device if isinstance(loss.device, str) else f"{loss.device[0]} * {len(loss.device)}"
      #  53  221.74 ms run,    2.22 ms python,  219.52 ms CL,  803.39 loss, 0.000807 LR, 4.66 GB used,   3042.49 GFLOPS,    674.65 GOPS
@@ -424,4 +425,5 @@ def train_cifar():
      raise ValueError(colored(f"{eval_acc_pct=} < {target}", "red"))

 if __name__ == "__main__":
-  train_cifar()
+  with WallTimeEvent(BenchEvent.FULL):
+    train_cifar()
@@ -13,6 +13,7 @@ from extra.models.llama import Transformer, convert_from_huggingface, fix_bf16
 from sentencepiece import SentencePieceProcessor
 import tiktoken, sys
 from tiktoken.load import load_tiktoken_bpe
+from extra.bench_log import BenchEvent, WallTimeEvent

 MAX_CONTEXT = getenv("MAX_CONTEXT", 4096)

@@ -206,40 +207,42 @@ class LLaMa:

    model = Transformer(**params["args"], linear=linear, max_context=MAX_CONTEXT, jit=bool(JIT))

-    if model_path.is_dir():
-      weights = concat_weights([load(filename) for filename in [f"{model_path}/consolidated.{i:02d}.pth" for i in range(params["files"])]], device[0] if isinstance(device, tuple) else device)
-    else:
-      weights = load(str(model_path))
-    if "model.embed_tokens.weight" in weights:
-      weights = convert_from_huggingface(weights, model, params["args"]["n_heads"], params["args"].get("n_kv_heads", params["args"]["n_heads"]))
+    with WallTimeEvent(BenchEvent.LOAD_WEIGHTS):
+      if model_path.is_dir():
+        weights = concat_weights([load(filename) for filename in [f"{model_path}/consolidated.{i:02d}.pth" for i in range(params["files"])]], device[0] if isinstance(device, tuple) else device)
+      else:
+        weights = load(str(model_path))
+      if "model.embed_tokens.weight" in weights:
+        weights = convert_from_huggingface(weights, params["args"]["n_layers"], params["args"]["n_heads"], params["args"].get("n_kv_heads", params["args"]["n_heads"]))

-    weights = fix_bf16(weights)
+      weights = fix_bf16(weights)

-    with Context(BEAM=0):
-      # quantize
-      if quantize is not None:
-        weights = linear.quantize(weights, device)
-        for _,v in weights.items(): v.realize()
+      # prevent tracking model weights
+      # this is a part of a larger problem with BUFFER UOps and gc in TRACK_MATCH_STATS=2
+      with Context(BEAM=0, TRACK_MATCH_STATS=0):
+        # quantize
+        if quantize is not None:
+          weights = linear.quantize(weights, device)
+          for _,v in weights.items(): v.realize()

-      # shard
-      if isinstance(device, tuple):
-        for k,v in nn.state.get_state_dict(model).items():
-          if 'scale' in k: v.shard_(device, axis=None)  # from quantized
-          elif '.attention.' in k:
-            if getenv("SHARD_KVCACHE") and ('.wq.' in k or '.wk.' in k or '.wv.' in k): v.shard_(device, axis=0)
-            else: v.shard_(device, axis=-1)
-          elif '.feed_forward.w1.' in k: v.shard_(device, axis=0)
-          elif '.feed_forward.w3.' in k: v.shard_(device, axis=0)
-          elif '.feed_forward.' in k: v.shard_(device, axis=-1)
-          elif 'tok_embeddings.weight' in k: v.shard_(device, axis=0)
-          elif 'output.weight' in k: v.shard_(device, axis=-1)
-          #elif k.endswith('.weight'): v.shard_(device, axis=-1)
-          #elif 'norm.' in k: v.shard_(device, axis=-1)
-          else: v.shard_(device, axis=None)
-          #print(k, v.shape, v.lazydata.axis)
+        # shard
+        if isinstance(device, tuple):
+          for k,v in nn.state.get_state_dict(model).items():
+            if 'scale' in k: v.shard_(device, axis=None)  # from quantized
+            elif '.attention.' in k:
+              if getenv("SHARD_KVCACHE") and ('.wq.' in k or '.wk.' in k or '.wv.' in k): v.shard_(device, axis=0)
+              else: v.shard_(device, axis=-1)
+            elif '.feed_forward.w1.' in k: v.shard_(device, axis=0)
+            elif '.feed_forward.w3.' in k: v.shard_(device, axis=0)
+            elif '.feed_forward.' in k: v.shard_(device, axis=-1)
+            elif 'tok_embeddings.weight' in k: v.shard_(device, axis=0)
+            elif 'output.weight' in k: v.shard_(device, axis=-1)
+            #elif k.endswith('.weight'): v.shard_(device, axis=-1)
+            #elif 'norm.' in k: v.shard_(device, axis=-1)
+            else: v.shard_(device, axis=None)

-      # replace weights in model
-      load_state_dict(model, weights, strict=False, consume=True)
+        # replace weights in model
+        load_state_dict(model, weights, strict=False, consume=True)

    return LLaMa(model, tokenizer)

@@ -327,7 +330,6 @@ int main()
 \end{code}
 """
 if __name__ == "__main__":
-  Tensor.no_grad = True
  print(f"using {Device.DEFAULT} backend")

  parser = argparse.ArgumentParser(description="Run LLaMA in tinygrad", formatter_class=argparse.ArgumentDefaultsHelpFormatter)
@@ -443,7 +445,7 @@ After you are done speaking, output [EOS]. You are not Chad.
  print(f"using LLaMA{LLAMA_SUFFIX}-{args.size} model")
  device = tuple(f"{Device.DEFAULT}:{i}" for i in range(args.shard)) if args.shard > 1 else Device.DEFAULT
  llama = LLaMa.build(MODEL_PATH, TOKENIZER_PATH, model_gen=args.gen, model_size=args.size, quantize=args.quantize, device=device)
-  param_bytes = sum(x.lazydata.size * x.dtype.itemsize for x in get_parameters(llama.model))
+  param_bytes = sum(x.uop.size * x.dtype.itemsize for x in get_parameters(llama.model))

  outputted = pre_prompt if chatbot else args.prompt
  start_pos, toks = 0, [llama.tokenizer.bos_id()] + llama.tokenizer.encode(outputted)
@@ -475,11 +477,12 @@ After you are done speaking, output [EOS]. You are not Chad.
      next_tok = Tensor([toks[start_pos:]], device=device) if tok_tensor is None or (len(toks)-start_pos) > 1 else tok_tensor.reshape(1, 1)
      with Profiling(enabled=args.profile):
        with Timing("total ", enabled=args.timing, on_exit=lambda x: f", {1e9/x:.2f} tok/s, {GlobalCounters.global_mem/x:.2f} GB/s, param {param_bytes/x:.2f} GB/s"):
-          with Timing("enqueue in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on GPU" if DEBUG>=2 else "")+
-                      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, param {param_bytes*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s" if DEBUG>=2 else "")) if DEBUG else None, enabled=args.timing):
-            tok_tensor = llama.model(next_tok, start_pos, args.temperature)
-          tok = tok_tensor.item()
+          with WallTimeEvent(BenchEvent.STEP):
+            with Timing("enqueue in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on GPU" if DEBUG>=2 else "")+
+                        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, param {param_bytes*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s" if DEBUG>=2 else "")) if DEBUG else None, enabled=args.timing):
+              tok_tensor = llama.model(next_tok, start_pos, args.temperature)
+            tok = tok_tensor.item()

      # use the kv cache
      start_pos = len(toks)
@@ -7,6 +7,7 @@ from extra.models.llama import Transformer, convert_from_huggingface, convert_fr
 from tinygrad.nn.state import safe_load, torch_load, load_state_dict, get_parameters, gguf_load
 from tinygrad import Tensor, dtypes, nn, Context, Device, GlobalCounters
 from tinygrad.helpers import Profiling, Timing, DEBUG, colored, fetch, tqdm
+from extra.bench_log import BenchEvent, WallTimeEvent

 class Tokenizer:
  pat_str = r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"
@@ -126,7 +127,8 @@ def NF4Linear(block_size):
      return x.linear(unscaled.reshape(self.out_features, self.in_features).T)

    @staticmethod
-    def quantize(state_dict: dict[str, Tensor], device, scale_dtype=dtypes.float16) -> dict[str, Tensor]:
+    def quantize(state_dict: dict[str, Tensor], device, scale_dtype=dtypes.float16, quantize_embeds=False) -> dict[str, Tensor]:
+      assert not quantize_embeds  # TODO: support this?
      new_state_dict = {}
      for k, v in state_dict.items():
        if "feed_forward" in k or "attention.w" in k:
@@ -165,40 +167,42 @@ def build_transformer(model_path: Path, model_size="8B", quantize=None, scale_dt
  model = Transformer(**MODEL_PARAMS[model_size]["args"], linear=linear, embedding=embedding, max_context=max_context, jit=True)

  if not load_weights: return model
+
  # load weights
-  if model_path.is_dir():
-    if (model_path / "model.safetensors.index.json").exists(): weights = load(str(model_path / "model.safetensors.index.json"))
-    elif (model_path / "model.safetensors").exists(): weights = load(str(model_path / "model.safetensors"))
-    else: weights = concat_weights([load(str(model_path / f"consolidated.{i:02d}.pth")) for i in range(MODEL_PARAMS[model_size]["files"])], device[0] if isinstance(device, tuple) else device)
-  else:
-    weights = load(str(model_path))
-  if "model.embed_tokens.weight" in weights:
-    weights = convert_from_huggingface(weights, model, MODEL_PARAMS[model_size]["args"]["n_heads"], MODEL_PARAMS[model_size]["args"]["n_kv_heads"])
-  elif "token_embd.weight" in weights:
-    weights = convert_from_gguf(weights, model)
-  weights = fix_bf16(weights)
+  with WallTimeEvent(BenchEvent.LOAD_WEIGHTS):
+    if model_path.is_dir():
+      if (model_path / "model.safetensors.index.json").exists(): weights = load(str(model_path / "model.safetensors.index.json"))
+      elif (model_path / "model.safetensors").exists(): weights = load(str(model_path / "model.safetensors"))
+      else: weights = concat_weights([load(str(model_path / f"consolidated.{i:02d}.pth")) for i in range(MODEL_PARAMS[model_size]["files"])], device[0] if isinstance(device, tuple) else device)
+    else:
+      weights = load(str(model_path))
+    if "model.embed_tokens.weight" in weights:
+      weights = convert_from_huggingface(weights, MODEL_PARAMS[model_size]["args"]["n_layers"], MODEL_PARAMS[model_size]["args"]["n_heads"], MODEL_PARAMS[model_size]["args"]["n_kv_heads"])
+    elif "token_embd.weight" in weights:
+      weights = convert_from_gguf(weights, MODEL_PARAMS[model_size]["args"]["n_layers"])
+    weights = fix_bf16(weights)

-  with Context(BEAM=0):
-    # quantize
-    if quantize == "float16": weights = {k:v.cast(quantize).contiguous() for k,v in weights.items()}
-    elif quantize is not None:
-      weights = linear.quantize(weights, device, scale_dtype, quantize_embeds)
-      for _,v in weights.items(): v.realize()
+    with Context(BEAM=0):
+      # quantize
+      if quantize == "float16": weights = {k:v.cast(quantize).contiguous() for k,v in weights.items()}
+      elif quantize is not None:
+        weights = linear.quantize(weights, device, scale_dtype, quantize_embeds)
+        for _,v in weights.items(): v.realize()

-    # shard
-    if isinstance(device, tuple):
-      for k,v in nn.state.get_state_dict(model).items():
-        if 'scale' in k: v.shard_(device, axis=None)  # from quantized
-        elif '.attention.' in k: v.shard_(device, axis=-1)
-        elif '.feed_forward.w1.' in k: v.shard_(device, axis=0)
-        elif '.feed_forward.w3.' in k: v.shard_(device, axis=0)
-        elif '.feed_forward.' in k: v.shard_(device, axis=-1)
-        elif 'tok_embeddings.weight' in k: v.shard_(device, axis=0)
-        elif 'output.weight' in k: v.shard_(device, axis=0)
-        else: v.shard_(device, axis=None)
+      # shard
+      if isinstance(device, tuple):
+        for k,v in nn.state.get_state_dict(model).items():
+          if 'scale' in k: v.shard_(device, axis=None)  # from quantized
+          elif '.attention.' in k: v.shard_(device, axis=-1)
+          elif '.feed_forward.w1.' in k: v.shard_(device, axis=0)
+          elif '.feed_forward.w3.' in k: v.shard_(device, axis=0)
+          elif '.feed_forward.' in k: v.shard_(device, axis=-1)
+          elif 'tok_embeddings.weight' in k: v.shard_(device, axis=0)
+          elif 'output.weight' in k: v.shard_(device, axis=0)
+          else: v.shard_(device, axis=None)

-    # replace weights in model
-    load_state_dict(model, weights, strict=False, consume=True)
+      # replace weights in model
+      load_state_dict(model, weights, strict=False, consume=True)
  return model

 # default settings
@@ -229,8 +233,6 @@ def prefill(model, toks, start_pos=0):
  return start_pos

 if __name__ == "__main__":
-  Tensor.no_grad = True
-
  parser = argparse.ArgumentParser()
  parser.add_argument("--download_model", action="store_true", help="Download a model")
  parser.add_argument("--model", type=Path, help="Model path")
@@ -282,7 +284,7 @@ if __name__ == "__main__":

  device = tuple(f"{Device.DEFAULT}:{i}" for i in range(args.shard)) if args.shard > 1 else Device.DEFAULT
  model = build_transformer(args.model, model_size=args.size, quantize=args.quantize, device=device)
-  param_bytes = sum(x.lazydata.size * x.dtype.itemsize for x in get_parameters(model))
+  param_bytes = sum(x.uop.size * x.dtype.itemsize for x in get_parameters(model))

  if not args.no_api and not args.benchmark:
    from bottle import Bottle, request, response, HTTPResponse, abort, static_file
@@ -434,11 +436,12 @@ if __name__ == "__main__":
      st = GlobalCounters.time_sum_s
      with Profiling(enabled=args.profile):
        with Timing("total ", on_exit=lambda x: f", {1e9/x:.2f} tok/s, {GlobalCounters.global_mem/x:.2f} GB/s, param {param_bytes/x:.2f} GB/s"):
-          with Timing("enqueue in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on GPU" if DEBUG>=2 else "")+
-                      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, param {param_bytes*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s" if DEBUG>=2 else "")) if DEBUG else None):
-            tok = model(Tensor([[last_tok]], device=device), start_pos, TEMPERATURE, TOP_K, TOP_P, ALPHA_F, ALPHA_P)
-          tok = tok.item()
+          with WallTimeEvent(BenchEvent.STEP):
+            with Timing("enqueue in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on GPU" if DEBUG>=2 else "")+
+                        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, param {param_bytes*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s" if DEBUG>=2 else "")) if DEBUG else None):
+              tok = model(Tensor([[last_tok]], device=device), start_pos, TEMPERATURE, TOP_K, TOP_P, ALPHA_F, ALPHA_P)
+            tok = tok.item()
      start_pos += 1
      last_tok = tok
      generated += tokenizer.decode([tok])
@@ -7,7 +7,7 @@ from train_gpt2 import GPT, GPTConfig
 from tinygrad.helpers import dedup, to_function_name, flatten, getenv, GlobalCounters, ansilen, to_function_name
 from tinygrad.engine.realize import get_kernel, run_schedule
 from tinygrad.engine.memory import memory_planner
-from tinygrad.ops import Ops
+from tinygrad.uop.ops import Ops

 TIMING = getenv("TIMING")

@@ -16,7 +16,7 @@ if __name__ == "__main__":
  #model.load_pretrained()
  for p in nn.state.get_parameters(model): p.replace(Tensor.empty(p.shape, dtype=p.dtype)) # fake load pretrained

-  #early_sched = create_schedule([x.lazydata for x in nn.state.get_parameters(model)])
+  #early_sched = create_schedule([x.uop for x in nn.state.get_parameters(model)])
  #print(f"built model {len(early_sched)}")

  #B, T = Variable("B", 1, 128).bind(4), 64 #Variable("T", 1, 1024).bind(64)
@@ -25,7 +25,7 @@ if __name__ == "__main__":
  Tensor.training = True
  optimizer = nn.optim.Adam(nn.state.get_parameters(model), lr=1e-4)
  warmup_count = getenv("WARMUP", 3)
-  for i in range(warmup_count):  # TODO: why does it take three and not two to stablize
+  for i in range(warmup_count):  # TODO: why does it take three and not two to stabilize
    GlobalCounters.reset()
    X = Tensor.empty(4, 64, dtype=dtypes.int).reshape(B, T)
    Y = Tensor.empty(4, 64, dtype=dtypes.int).reshape(B, T)
@@ -56,7 +56,7 @@ if __name__ == "__main__":
  state_dict.update({'X': X, 'Y': Y, 'loss': loss})
  grad_state_dict = {}
  for k,v in state_dict.items():
-    if v.lazydata.base.buffer not in used_buffers: print(f"UNUSED: {k}")
+    if v.uop.base.buffer not in used_buffers: print(f"UNUSED: {k}")
    if v.grad is not None: grad_state_dict['grad_'+k] = v.grad
  state_dict.update(grad_state_dict)
  state_dict.update({'adam_b1_t': optimizer.b1_t, 'adam_b2_t': optimizer.b2_t, 'adam_lr': optimizer.lr})
@@ -65,7 +65,7 @@ if __name__ == "__main__":
    nm = inverse_state_dict[p]
    state_dict["adam_m_"+nm] = m
    state_dict["adam_v_"+nm] = v
-  named_buffers = {v.lazydata.base.buffer:k.replace(".", "_") for k,v in state_dict.items()}
+  named_buffers = {v.uop.base.buffer:k.replace(".", "_") for k,v in state_dict.items()}

  c_code = ["#include <stdlib.h>", "#include <tgmath.h>", "#include <stdbool.h>"]
  if TIMING: c_code += ["#include <stdio.h>", "#include <time.h>"]
@@ -99,7 +99,7 @@ class GPT:

  def __call__(self, idx:Tensor, targets=None):
    b, t = idx.shape
-    pos = Tensor.arange(0, t)
+    pos = Tensor.arange(0, t, device=idx.device)

    tok_emb = self.wte(idx) # token embeddings of shape (b, t, n_embd)
    pos_emb = self.wpe(pos) # position embeddings of shape (t, n_embd)
@@ -124,6 +124,7 @@ if __name__ == "__main__":
  parser.add_argument("--batch_size", type=int, default=4, help="batch size")
  parser.add_argument("--sequence_length", type=int, default=64, help="sequence length")
  parser.add_argument("--skip_test", action="store_true", help="skip test")
+  parser.add_argument("--gpus", type=int, default=1, help="sequence length")
  args = parser.parse_args()
  B, T = args.batch_size, args.sequence_length
  assert 1 <= T <= 1024
@@ -131,6 +132,10 @@ if __name__ == "__main__":
  model = GPT(GPTConfig(n_layer=12, n_head=12, n_embd=768))
  model.load_pretrained()

+  if args.gpus > 1:
+    GPUS = tuple(f'{Device.DEFAULT}:{i}' for i in range(args.gpus))
+    for x in nn.state.get_parameters(model): x.to_(GPUS)  # we put a copy of the model on every GPU
+
  # init the tokenizer
  enc = tiktoken.get_encoding("gpt2")
  encode = lambda s: enc.encode(s, allowed_special={"<|endoftext|>"})
@@ -165,23 +170,32 @@ if __name__ == "__main__":
  x, y = next(data_iter) # we'll overfit this batch below
  optimizer = nn.optim.AdamW(nn.state.get_parameters(model), lr=1e-4, weight_decay=0)

+  print(f"model state:     {sum(x.nbytes() for x in nn.state.get_parameters(model))/1e9:.2f} GB")
+  print(f"optimizer state: {sum(x.nbytes() for x in nn.state.get_parameters(optimizer))/1e9:.2f} GB")
+
+  # shard the data on axis 0
+  if args.gpus > 1: x, y = x.shard(GPUS, axis=0), y.shard(GPUS, axis=0)
+
  @TinyJit
-  def step(x, y):
+  @Tensor.train()
+  def step(x:Tensor, y:Tensor) -> Tensor:
    _, loss = model(x, y)
    optimizer.zero_grad()
    loss.backward()
    return loss.realize(*optimizer.schedule_step())

-  with Tensor.train():
-    for i in range(args.num_iterations):
-      GlobalCounters.reset()
-      t0 = time.time()
-      loss = step(x.contiguous(), y.contiguous())
-      Device[Device.DEFAULT].synchronize()
-      t1 = time.time()
-      print(f"iteration {i}, loss: {loss.item():.6f}, time: {(t1-t0)*1000:.3f}ms, {int(B*T/(t1-t0))} tok/s")
+  for i in range(args.num_iterations):
+    GlobalCounters.reset()
+    t0 = time.perf_counter()
+    loss = step(x.contiguous(), y.contiguous())
+    Device[Device.DEFAULT].synchronize()
+    t1 = time.perf_counter()
+    print(f"iteration {i}, loss: {loss.item():.6f}, time: {(t1-t0)*1000:.3f}ms, {int(B*T/(t1-t0))} tok/s, {GlobalCounters.global_mem/1e9:.2f} GB")

  if not args.skip_test:
+    # copy back to single gpu for test
+    if args.gpus > 1:
+      for x in nn.state.get_parameters(model): x.to_(Device.DEFAULT)
    start = "<|endoftext|>"
    start_ids = encode(start)
    x = (Tensor(start_ids)[None, ...])
@@ -0,0 +1,156 @@
+# much taken from https://github.com/cloneofsimo/minRF
+from tinygrad import Tensor, nn, GlobalCounters, TinyJit
+from tinygrad.helpers import getenv, trange
+from extra.models.llama import Attention, FeedForward, precompute_freqs_cis
+
+def modulate(x:Tensor, shift:Tensor, scale:Tensor) -> Tensor: return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+# TODO: why doesn't the TimestepEmbedder from minRF work?
+class TimestepEmbedder:
+  def __init__(self, hidden_size): self.mlp = [nn.Linear(1, hidden_size), Tensor.silu, nn.Linear(hidden_size, hidden_size)]
+  def __call__(self, t:Tensor): return t.reshape(-1, 1).sequential(self.mlp)
+
+class TransformerBlock:
+  def __init__(self, dim, n_heads, norm_eps=1e-5):
+    self.attention = Attention(dim, n_heads)
+    self.feed_forward = FeedForward(dim, 4*dim)
+    self.attention_norm = nn.LayerNorm(dim, eps=norm_eps)
+    self.ffn_norm = nn.LayerNorm(dim, eps=norm_eps)
+    self.adaLN_modulation = nn.Linear(dim, 6 * dim, bias=True)
+
+  def __call__(self, x:Tensor, freqs_cis:Tensor, adaln_input:Tensor):
+    shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(adaln_input.silu()).chunk(6, dim=1)
+    x = x + gate_msa.unsqueeze(1) * self.attention(modulate(self.attention_norm(x), shift_msa, scale_msa), 0, freqs_cis)
+    x = x + gate_mlp.unsqueeze(1) * self.feed_forward(modulate(self.ffn_norm(x), shift_mlp, scale_mlp))
+    return x.contiguous().contiguous_backward()
+
+class FinalLayer:
+  def __init__(self, dim, patch_size, out_channels):
+    self.norm_final = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+    self.linear = nn.Linear(dim, patch_size*patch_size*out_channels, bias=True)
+    self.adaLN_modulation = nn.Linear(dim, 2 * dim, bias=True)
+
+    # init weights/bias to 0
+    self.linear.weight.replace(self.linear.weight.zeros_like().contiguous())
+    self.linear.bias.replace(self.linear.bias.zeros_like().contiguous())
+
+  def __call__(self, x:Tensor, c:Tensor):
+    shift, scale = self.adaLN_modulation(c.silu()).chunk(2, dim=1)
+    x = modulate(self.norm_final(x), shift, scale)
+    return self.linear(x)
+
+# channels=1, input_size=32, dim=64, n_layers=6, n_heads=4, num_classes=10
+class DiT_Llama:
+  def __init__(self, in_channels=1, dim=64, n_layers=6, n_heads=4, num_classes=10, patch_size=2):
+    self.patch_size = patch_size
+    self.out_channels = in_channels
+    self.num_classes = num_classes
+
+    self.init_conv_seq = [
+      nn.Conv2d(in_channels, dim // 2, kernel_size=5, padding=2, stride=1), Tensor.silu, nn.GroupNorm(32, dim//2),
+      nn.Conv2d(dim //2, dim // 2, kernel_size=5, padding=2, stride=1), Tensor.silu, nn.GroupNorm(32, dim//2),
+    ]
+
+    self.x_embedder = nn.Linear(self.patch_size * self.patch_size * dim // 2, dim, bias=True)
+    self.t_embedder = TimestepEmbedder(dim)
+    self.y_embedder = nn.Embedding(num_classes+1, dim)
+    self.final_layer = FinalLayer(dim, self.patch_size, self.out_channels)
+
+    self.freqs_cis = precompute_freqs_cis(dim // n_heads, 4096)
+    self.layers = [TransformerBlock(dim, n_heads) for _ in range(n_layers)]
+
+  def unpatchify(self, x:Tensor):
+    c, p = self.out_channels, self.patch_size
+    h = w = int(x.shape[1] ** 0.5)
+    x = x.reshape(shape=(x.shape[0], h, w, p, p, c))
+    x = x.rearrange("n h w p q c -> n c h p w q")
+    return x.reshape(shape=(x.shape[0], c, h * p, h * p))
+
+  def patchify(self, x:Tensor):
+    B, C, H, W = x.shape
+    x = x.reshape(B, C, H // self.patch_size, self.patch_size, W // self.patch_size, self.patch_size)
+    x = x.permute(0, 2, 4, 1, 3, 5).flatten(-3).flatten(1, 2)
+    return x  # B <H*W ish> <C*patch_size*patch_size>
+
+  def __call__(self, x:Tensor, t:Tensor, y:Tensor) -> Tensor:
+    x = x.sequential(self.init_conv_seq)
+    x = self.patchify(x)
+    x = self.x_embedder(x)
+    adaln_input = self.t_embedder(t) + self.y_embedder(y)
+    adaln_input = adaln_input.contiguous()
+    for layer in self.layers:
+      x = layer(x, self.freqs_cis[:, :x.size(1)], adaln_input=adaln_input)
+    x = self.final_layer(x, adaln_input)
+    return self.unpatchify(x)
+
+  def rf(self, x:Tensor, cond:Tensor):
+    b = x.shape[0]
+    # self.ln is True
+    t = Tensor.randn((b,)).sigmoid()
+    texp = t.view([b, *([1] * len(x.shape[1:]))])
+
+    # conditional dropout
+    dropout_prob = 0.1
+    cond = (Tensor.rand(cond.shape[0]) < dropout_prob).where(cond.full_like(self.num_classes), cond)
+
+    # this is rectified flow
+    z1 = x.randn_like()
+    zt = (1 - texp) * x + texp * z1
+    vtheta = self(zt, t, cond)
+
+    # MSE loss
+    return ((z1 - x) - vtheta).square().mean()
+
+  def sample(self, z, cond, null_cond, sample_steps=50, cfg=2.0):
+    b = z.size(0)
+    dt = Tensor.full((b,)+(1,)*len(z.shape[1:]), fill_value=1.0/sample_steps).contiguous()
+    images = [z]
+    for i in range(sample_steps, 0, -1):
+      t = Tensor.full((b,), fill_value=i/sample_steps).contiguous()
+      vc = self(z, t, cond)
+      vu = self(z, t, null_cond)
+      vc = vu + cfg * (vc - vu)
+      z = z - dt * vc
+      z = z.contiguous()
+      images.append(z)
+    return images
+
+def mviz(t:Tensor):
+  assert len(t.shape) == 4 and t.shape[1] == 1
+  ft = t.permute(1,2,0,3).reshape(32, -1)
+  assert ft.shape[-1]%32 == 0
+  print("")
+  for y in ((ft+1)/2).clamp(0,1).tolist():
+    ln = [f"\033[38;5;{232+int(x*23)}m██" for x in y]
+    print(''.join(ln) + "\033[0m")
+
+if __name__ == "__main__":
+  X_train, Y_train, X_test, Y_test = nn.datasets.mnist()
+  X_train = X_train.pad((2,2,2,2))
+  X_train = ((X_train.float()/255)-0.5)/0.5
+  Y_train = Y_train.int()
+
+  model = DiT_Llama(patch_size=getenv("PATCH_SIZE", 2))
+  for r in nn.state.get_parameters(model): r.realize()
+  optimizer = nn.optim.Adam(nn.state.get_parameters(model), lr=5e-4)
+
+  @TinyJit
+  @Tensor.train()
+  def train_step():
+    if getenv("OVERFIT"): samples = Tensor.zeros(getenv("BS", 256), dtype='int')
+    else: samples = Tensor.randint(getenv("BS", 256), high=X_train.shape[0])
+    optimizer.zero_grad()
+    loss = model.rf(X_train[samples], Y_train[samples])
+    loss.backward()
+    optimizer.step()
+    return loss
+
+  @TinyJit
+  def sample(z:Tensor, cond:Tensor) -> Tensor:
+    return model.sample(z, cond, Tensor.full_like(cond, 10), sample_steps=getenv("SAMPLE_STEPS", 20))[-1]
+
+  for steps in (t:=trange(getenv("STEPS", 5000))):
+    if steps%10 == 0: mviz(sample(Tensor.randn(3, 1, 32, 32), Tensor([5,0,4], dtype='int')))
+    GlobalCounters.reset()
+    loss = train_step()
+    t.set_description(f"loss: {loss.item():9.2f}")
@@ -3,6 +3,7 @@ from tinygrad import Tensor, nn, Device, GlobalCounters, Variable
 from tinygrad.helpers import Timing, Profiling, CI, tqdm
 from tinygrad.nn.state import torch_load, get_state_dict
 from extra.models.llama import FeedForward, Transformer
+from extra.bench_log import BenchEvent, WallTimeEvent

 class MixtureFeedForward:
  def __init__(self, num_experts:int, dim:int, hidden_dim:int, linear=nn.Linear):
@@ -30,18 +31,19 @@ if __name__ == "__main__":
                      help="Path to the downloaded weights")
  args = parser.parse_args()

-  state = torch_load(args.weights + "/consolidated.00.pth.b")
-  model = Transformer(n_layers=32, dim=4096, hidden_dim=14336, n_heads=32, n_kv_heads=8, norm_eps=1e-5, vocab_size=32000, feed_forward=functools.partial(MixtureFeedForward, 8), jit=False)
-  model_state_dict = get_state_dict(model)
+  with WallTimeEvent(BenchEvent.LOAD_WEIGHTS):
+    state = torch_load(args.weights + "/consolidated.00.pth.b")
+    model = Transformer(n_layers=32, dim=4096, hidden_dim=14336, n_heads=32, n_kv_heads=8, norm_eps=1e-5, vocab_size=32000, feed_forward=functools.partial(MixtureFeedForward, 8), jit=False)
+    model_state_dict = get_state_dict(model)

-  for k in (t := tqdm(state, disable=CI)):
-    if 'feed_forward.experts.' in k:
-      expert_no = int(k.split('feed_forward.experts.')[1].split('.')[0])
-      device = Device.DEFAULT + ":" + str((expert_no//2)+1)
-    else:
-      device = Device.DEFAULT
-    t.set_description(f"ram used: {GlobalCounters.mem_used/1e9:5.2f} GB, loading {k} to {device}")
-    model_state_dict[k].replace(state[k].to(device).half()).realize()
+    for k in (t := tqdm(state, disable=CI)):
+      if 'feed_forward.experts.' in k:
+        expert_no = int(k.split('feed_forward.experts.')[1].split('.')[0])
+        device = Device.DEFAULT + ":" + str((expert_no//2)+1)
+      else:
+        device = Device.DEFAULT
+      t.set_description(f"ram used: {GlobalCounters.mem_used/1e9:5.2f} GB, loading {k} to {device}")
+      model_state_dict[k].replace(state[k].to(device).half()).realize()
  if CI: print(f"ram used: {GlobalCounters.mem_used/1e9:5.2f} GB")

  from sentencepiece import SentencePieceProcessor
@@ -53,7 +55,8 @@ if __name__ == "__main__":
    GlobalCounters.reset()
    with Profiling(sort="time", frac=0.1, enabled=args.profile):
      with Timing("total ", enabled=args.timing, on_exit=lambda x: f", {1e9/x:.2f} tok/sec"):
-        tok = model(Tensor([toks[start_pos:]]), 0 if start_pos == 0 else Variable("start_pos", 1, 1024).bind(start_pos), args.temperature).item()
+        with WallTimeEvent(BenchEvent.STEP):
+          tok = model(Tensor([toks[start_pos:]]), 0 if start_pos == 0 else Variable("start_pos", 1, 1024-1).bind(start_pos), args.temperature).item()
    toks.append(tok)
    start_pos += 1
    print(spp.decode(toks))
@@ -5,7 +5,7 @@ from multiprocessing import Queue, Process, shared_memory, connection, Lock, cpu

 import numpy as np
 from tinygrad import dtypes, Tensor
-from tinygrad.helpers import getenv, prod, Context, round_up, tqdm
+from tinygrad.helpers import getenv, prod, Context, round_up, tqdm, OSX

 ### ResNet

@@ -71,7 +71,7 @@ def loader_process(q_in, q_out, X:Tensor, seed):
      #storage_tensor._copyin(img_tensor.numpy())

      # faster
-      X[idx].contiguous().realize().lazydata.base.realized.as_buffer(force_zero_copy=True)[:] = img.tobytes()
+      X[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = img.tobytes()

      # ideal
      #X[idx].assign(img.tobytes())   # NOTE: this is slow!
@@ -129,14 +129,15 @@ def batch_load_resnet(batch_size=64, val=False, shuffle=True, seed=None, pad_fir
  q_in, q_out = Queue(), Queue()

  sz = (batch_size*BATCH_COUNT, 224, 224, 3)
-  if os.path.exists("/dev/shm/resnet_X"): os.unlink("/dev/shm/resnet_X")
-  shm = shared_memory.SharedMemory(name="resnet_X", create=True, size=prod(sz))
+  shm_name = "resnet_X_val" if val else "resnet_X_train"
+  if not OSX and os.path.exists(f"/dev/shm/{shm_name}"): os.unlink(f"/dev/shm/{shm_name}")
+  shm = shared_memory.SharedMemory(name=shm_name, create=True, size=prod(sz))
  procs = []

  try:
    # disk:shm is slower
-    #X = Tensor.empty(*sz, dtype=dtypes.uint8, device=f"disk:shm:{shm.name}")
-    X = Tensor.empty(*sz, dtype=dtypes.uint8, device=f"disk:/dev/shm/resnet_X")
+    if OSX: X = Tensor.empty(*sz, dtype=dtypes.uint8, device=f"disk:shm:{shm.name}")
+    else: X = Tensor.empty(*sz, dtype=dtypes.uint8, device=f"disk:/dev/shm/{shm_name}")
    Y = [None] * (batch_size*BATCH_COUNT)

    for _ in range(cpu_count()):
@@ -261,8 +262,8 @@ def load_unet3d_data(preprocessed_dataset_dir, seed, queue_in, queue_out, X:Tens
      x = random_brightness_augmentation(x)
      x = gaussian_noise(x)

-    X[idx].contiguous().realize().lazydata.base.realized.as_buffer(force_zero_copy=True)[:] = x.tobytes()
-    Y[idx].contiguous().realize().lazydata.base.realized.as_buffer(force_zero_copy=True)[:] = y.tobytes()
+    X[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = x.tobytes()
+    Y[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = y.tobytes()

    queue_out.put(idx)
  queue_out.put(None)
@@ -312,7 +313,7 @@ def batch_load_unet3d(preprocessed_dataset_dir:Path, batch_size:int=6, val:bool=
      proc = Process(target=load_unet3d_data, args=(preprocessed_dataset_dir, seed, queue_in, queue_out, X, Y))
      proc.daemon = True
      proc.start()
-      
+
      procs.append(proc)

    for bc in range(batch_count):
@@ -376,12 +377,12 @@ def load_retinanet_data(base_dir:Path, val:bool, queue_in:Queue, queue_out:Queue
      clipped_match_idxs = np.clip(match_idxs, 0, None)
      clipped_boxes, clipped_labels = tgt["boxes"][clipped_match_idxs], tgt["labels"][clipped_match_idxs]

-      boxes[idx].contiguous().realize().lazydata.base.realized.as_buffer(force_zero_copy=True)[:] = clipped_boxes.tobytes()
-      labels[idx].contiguous().realize().lazydata.base.realized.as_buffer(force_zero_copy=True)[:] = clipped_labels.tobytes()
-      matches[idx].contiguous().realize().lazydata.base.realized.as_buffer(force_zero_copy=True)[:] = match_idxs.tobytes()
-      anchors[idx].contiguous().realize().lazydata.base.realized.as_buffer(force_zero_copy=True)[:] = anchor.tobytes()
+      boxes[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = clipped_boxes.tobytes()
+      labels[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = clipped_labels.tobytes()
+      matches[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = match_idxs.tobytes()
+      anchors[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = anchor.tobytes()

-    imgs[idx].contiguous().realize().lazydata.base.realized.as_buffer(force_zero_copy=True)[:] = img.tobytes()
+    imgs[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = img.tobytes()

    queue_out.put(idx)
  queue_out.put(None)
@@ -532,7 +533,7 @@ if __name__ == "__main__":
  def load_retinanet(val):
    from extra.datasets.openimages import BASEDIR, download_dataset
    from pycocotools.coco import COCO
-    dataset = COCO(download_dataset(base_dir:=getenv("BASE_DIR", BASEDIR), "validation" if val else "train"))
+    dataset = COCO(download_dataset(base_dir:=getenv("BASEDIR", BASEDIR), "validation" if val else "train"))
    with tqdm(total=len(dataset.imgs.keys())) as pbar:
      for x in batch_load_retinanet(dataset, val, base_dir):
        pbar.update(x[0].shape[0])
@@ -295,3 +295,62 @@ def generate_anchors(input_size:tuple[int, int], scales:Optional[tuple[Tensor, .
    anchors.append((shifts[:, None] + base_anchors[None, :]).reshape(-1, 4))

  return anchors
+
+
+class BoxCoder(object):
+  def __init__(self, weights, bbox_xform_clip=math.log(1000. / 16), apply_to_remove=True):
+    self.weights = weights
+    self.bbox_xform_clip = bbox_xform_clip
+    self.apply_to_remove = apply_to_remove
+
+  def encode(self, reference_boxes, proposals):
+    TO_REMOVE = self.apply_to_remove  # TODO remove
+    ex_widths = proposals[..., 2] - proposals[..., 0] + TO_REMOVE
+    ex_heights = proposals[..., 3] - proposals[..., 1] + TO_REMOVE
+    ex_ctr_x = proposals[..., 0] + 0.5 * ex_widths
+    ex_ctr_y = proposals[..., 1] + 0.5 * ex_heights
+
+    gt_widths = reference_boxes[..., 2] - reference_boxes[..., 0] + TO_REMOVE
+    gt_heights = reference_boxes[..., 3] - reference_boxes[..., 1] + TO_REMOVE
+    gt_ctr_x = reference_boxes[..., 0] + 0.5 * gt_widths
+    gt_ctr_y = reference_boxes[..., 1] + 0.5 * gt_heights
+
+    wx, wy, ww, wh = self.weights
+    targets_dx = wx * (gt_ctr_x - ex_ctr_x) / ex_widths
+    targets_dy = wy * (gt_ctr_y - ex_ctr_y) / ex_heights
+    targets_dw = ww * Tensor.log(gt_widths / ex_widths)
+    targets_dh = wh * Tensor.log(gt_heights / ex_heights)
+
+    targets = Tensor.stack(targets_dx, targets_dy, targets_dw, targets_dh, dim=-1)
+    return targets
+
+  def decode(self, rel_codes, boxes):
+    boxes = boxes.cast(rel_codes.dtype)
+    rel_codes = rel_codes
+
+    TO_REMOVE = self.apply_to_remove  # TODO remove
+    widths = boxes[:, 2] - boxes[:, 0] + TO_REMOVE
+    heights = boxes[:, 3] - boxes[:, 1] + TO_REMOVE
+    ctr_x = boxes[:, 0] + 0.5 * widths
+    ctr_y = boxes[:, 1] + 0.5 * heights
+
+    wx, wy, ww, wh = self.weights
+    dx = rel_codes[:, 0::4] / wx
+    dy = rel_codes[:, 1::4] / wy
+    dw = rel_codes[:, 2::4] / ww
+    dh = rel_codes[:, 3::4] / wh
+
+    # Prevent sending too large values into Tensor.exp()
+    dw = dw.clip(min_=dw.min(), max_=self.bbox_xform_clip)
+    dh = dh.clip(min_=dh.min(), max_=self.bbox_xform_clip)
+
+    pred_ctr_x = dx * widths[:, None] + ctr_x[:, None]
+    pred_ctr_y = dy * heights[:, None] + ctr_y[:, None]
+    pred_w = dw.exp() * widths[:, None]
+    pred_h = dh.exp() * heights[:, None]
+    x = pred_ctr_x - 0.5 * pred_w
+    y = pred_ctr_y - 0.5 * pred_h
+    w = pred_ctr_x + 0.5 * pred_w - 1
+    h = pred_ctr_y + 0.5 * pred_h - 1
+    pred_boxes = Tensor.stack(x, y, w, h).permute(1,2,0).reshape(rel_codes.shape[0], rel_codes.shape[1])
+    return pred_boxes
@@ -53,7 +53,9 @@ class EmbeddingBert(nn.Embedding):
    arange_shp, weight_shp, big_shp = (1, 1, self.vocab_sz, 1), (1, 1, self.vocab_sz, self.embed_sz), idx.shape+(self.vocab_sz, self.embed_sz,)
    if not hasattr(self, 'arange'): self.arange = Tensor.arange(self.vocab_sz, requires_grad=False, device=self.weight.device).reshape(arange_shp)
    arange, idx, vals = self.arange.expand(big_shp), idx.reshape(idx.shape+(1, 1,)).expand(big_shp), self.weight.cast(dtypes.default_float).reshape(weight_shp).expand(big_shp)
-    return (arange == idx).mul(vals).sum(2, dtype=vals.dtype)
+    # TODO: contiguous() here because the embedding dropout creates different asts on each device, and search becomes very slow.
+    # Should fix with fixing random ast on multi device, and fuse arange to make embedding fast.
+    return (arange == idx).mul(vals).sum(2, dtype=vals.dtype).contiguous()

 class LayerNormBert:
  def __init__(self, normalized_shape:Union[int, tuple[int, ...]], eps:float=1e-12, elementwise_affine:bool=True):
@@ -66,3 +68,62 @@ class LayerNormBert:
    xn = x.cast(dtypes.float32).layernorm(eps=self.eps, axis=self.axis).cast(x.dtype)
    if not self.elementwise_affine: return xn
    return (xn * self.weight.cast(dtypes.default_float) + self.bias.cast(dtypes.default_float))
+
+class FrozenBatchNorm2dRetinaNet(nn.BatchNorm2d):
+  def __init__(self, sz:int, eps=1e-5, affine=True, track_running_stats=True, momentum=0.1):
+    self.eps, self.track_running_stats, self.momentum = eps, track_running_stats, momentum
+
+    self.weight = Tensor.ones(sz, dtype=dtypes.float32, requires_grad=False) if affine else None
+    self.bias = Tensor.zeros(sz, dtype=dtypes.float32, requires_grad=False) if affine else None
+
+    if track_running_stats: self.running_mean, self.running_var = Tensor.zeros(sz, dtype=dtypes.float32, requires_grad=False), Tensor.ones(sz, dtype=dtypes.float32, requires_grad=False)
+    self.num_batches_tracked = Tensor.zeros(1, dtype=dtypes.long, requires_grad=False)
+
+  def __call__(self, x:Tensor) -> Tensor:
+    batch_mean, batch_var = super().calc_stats(x.cast(dtypes.float32))
+    if self.track_running_stats and Tensor.training:
+      self.running_mean.assign((1-self.momentum) * self.running_mean + self.momentum * batch_mean.detach().cast(self.running_mean.dtype))
+      self.running_var.assign((1-self.momentum) * self.running_var + self.momentum * x.numel()/(x.numel()-x.shape[1]) * batch_var.detach().cast(self.running_var.dtype))
+      self.num_batches_tracked += 1
+    return x.cast(dtypes.float32).batchnorm(self.weight, self.bias, batch_mean, batch_var.add(self.eps).rsqrt()).cast(x.dtype)
+
+class Conv2dNormalRetinaNet(nn.Conv2d):
+  def __init__(self, in_channels:int, out_channels:int, kernel_size:int|tuple[int, ...],
+               stride:int=1, padding:int|tuple[int, ...]|str=0, dilation:int=1, groups:int=1,
+               bias:bool=True, prior_prob:float|None=None):
+    super().__init__(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)
+    self.weight = Tensor.normal(*self.weight.shape, std=0.01, dtype=dtypes.float32)
+    if bias:
+      if prior_prob:
+        prior_prob = Tensor(prior_prob, device=self.bias.device, dtype=dtypes.float32).expand(*self.bias.shape)
+        self.bias = -(((1 - prior_prob) / prior_prob).log())
+      else: self.bias = Tensor.zeros_like(self.bias, dtype=dtypes.float32)
+
+  def __call__(self, x:Tensor) -> Tensor:
+    return x.conv2d(self.weight.cast(dtypes.default_float), self.bias.cast(dtypes.default_float) if self.bias is not None else None,
+                    groups=self.groups, stride=self.stride, padding=self.padding)
+
+class Conv2dKaimingUniformRetinaNet(nn.Conv2d):
+  def __init__(self, in_channels:int, out_channels:int, kernel_size:int|tuple[int, ...],
+               stride:int=1, padding:int|tuple[int, ...]|str=0, dilation:int=1, groups:int=1,
+               bias:bool=True):
+    super().__init__(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)
+    self.weight = Tensor.kaiming_uniform(*self.weight.shape, a=1, dtype=dtypes.float32)
+    if bias: self.bias = Tensor.zeros_like(self.bias, dtype=dtypes.float32)
+
+  def __call__(self, x:Tensor) -> Tensor:
+    return x.conv2d(self.weight.cast(dtypes.default_float), self.bias.cast(dtypes.default_float) if self.bias is not None else None,
+                    groups=self.groups, stride=self.stride, padding=self.padding)
+
+class Conv2dRetinaNet(nn.Conv2d):
+  def __init__(self, in_channels:int, out_channels:int, kernel_size:int|tuple[int, ...],
+               stride:int=1, padding:int|tuple[int, ...]|str=0, dilation:int=1, groups:int=1,
+               bias:bool=True):
+    super().__init__(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)
+    scale = 1 / math.sqrt(in_channels * prod(self.kernel_size))
+    self.weight = Tensor.uniform(out_channels, in_channels//groups, *self.kernel_size, low=-scale, high=scale, dtype=dtypes.float32)
+    self.bias: Tensor|None = Tensor.uniform(out_channels, low=-scale, high=scale, dtype=dtypes.float32) if bias else None
+
+  def __call__(self, x:Tensor) -> Tensor:
+    return x.conv2d(self.weight.cast(dtypes.default_float), self.bias.cast(dtypes.default_float) if self.bias is not None else None,
+                    groups=self.groups, stride=self.stride, dilation=self.dilation, padding=self.padding)
@@ -5,61 +5,62 @@ import numpy as np
 from tinygrad import Tensor, Device, dtypes, GlobalCounters, TinyJit
 from tinygrad.nn.state import get_parameters, load_state_dict, safe_load
 from tinygrad.helpers import getenv
+from extra.bench_log import BenchEvent, WallTimeEvent
 def tlog(x): print(f"{x:25s}  @ {time.perf_counter()-start:5.2f}s")

 def eval_resnet():
-  Tensor.no_grad = True
-  # Resnet50-v1.5
-  from extra.models.resnet import ResNet50
-  tlog("imports")
-  GPUS = [f'{Device.DEFAULT}:{i}' for i in range(getenv("GPUS", 6))]
-  for x in GPUS: Device[x]
-  tlog("got devices")    # NOTE: this is faster with rocm-smi running
+  with WallTimeEvent(BenchEvent.FULL):
+    # Resnet50-v1.5
+    from extra.models.resnet import ResNet50
+    tlog("imports")
+    GPUS = [f'{Device.DEFAULT}:{i}' for i in range(getenv("GPUS", 6))]
+    for x in GPUS: Device[x]
+    tlog("got devices")    # NOTE: this is faster with rocm-smi running

-  class ResnetRunner:
-    def __init__(self, device=None):
-      self.mdl = ResNet50()
-      for x in get_parameters(self.mdl) if device else []: x.to_(device)
-      if (fn:=getenv("RESNET_MODEL", "")): load_state_dict(self.mdl, safe_load(fn))
-      else: self.mdl.load_from_pretrained()
-      self.input_mean = Tensor([0.485, 0.456, 0.406], device=device).reshape(1, -1, 1, 1)
-      self.input_std = Tensor([0.229, 0.224, 0.225], device=device).reshape(1, -1, 1, 1)
-    def __call__(self, x:Tensor) -> Tensor:
-      x = x.permute([0,3,1,2]).cast(dtypes.float32) / 255.0
-      x -= self.input_mean
-      x /= self.input_std
-      return self.mdl(x).log_softmax().argmax(axis=1).realize()
+    class ResnetRunner:
+      def __init__(self, device=None):
+        self.mdl = ResNet50()
+        for x in get_parameters(self.mdl) if device else []: x.to_(device)
+        if (fn:=getenv("RESNET_MODEL", "")): load_state_dict(self.mdl, safe_load(fn))
+        else: self.mdl.load_from_pretrained()
+        self.input_mean = Tensor([0.485, 0.456, 0.406], device=device).reshape(1, -1, 1, 1)
+        self.input_std = Tensor([0.229, 0.224, 0.225], device=device).reshape(1, -1, 1, 1)
+      def __call__(self, x:Tensor) -> Tensor:
+        x = x.permute([0,3,1,2]).cast(dtypes.float32) / 255.0
+        x -= self.input_mean
+        x /= self.input_std
+        return self.mdl(x).log_softmax().argmax(axis=1).realize()

-  mdl = TinyJit(ResnetRunner(GPUS))
-  tlog("loaded models")
+    mdl = TinyJit(ResnetRunner(GPUS))
+    tlog("loaded models")

-  # evaluation on the mlperf classes of the validation set from imagenet
-  from examples.mlperf.dataloader import batch_load_resnet
-  iterator = batch_load_resnet(getenv("BS", 128*6), val=getenv("VAL", 1), shuffle=False, pad_first_batch=True)
-  def data_get():
-    x,y,cookie = next(iterator)
-    return x.shard(GPUS, axis=0).realize(), y, cookie
-  n,d = 0,0
-  proc = data_get()
-  tlog("loaded initial data")
-  st = time.perf_counter()
-  while proc is not None:
-    GlobalCounters.reset()
-    proc = (mdl(proc[0]), proc[1], proc[2])  # this frees the images
-    run = time.perf_counter()
-    # load the next data here
-    try: next_proc = data_get()
-    except StopIteration: next_proc = None
-    nd = time.perf_counter()
-    y = np.array(proc[1])
-    proc = (proc[0].numpy() == y) & (y != -1)  # this realizes the models and frees the cookies
-    n += proc.sum()
-    d += (y != -1).sum()
-    et = time.perf_counter()
-    tlog(f"****** {n:5d}/{d:5d}  {n*100.0/d:.2f}% -- {(run-st)*1000:7.2f} ms to enqueue, {(et-run)*1000:7.2f} ms to realize ({(nd-run)*1000:7.2f} ms fetching). {(len(proc))/(et-st):8.2f} examples/sec. {GlobalCounters.global_ops*1e-12/(et-st):5.2f} TFLOPS")
-    st = et
-    proc, next_proc = next_proc, None
-  tlog("done")
+    # evaluation on the mlperf classes of the validation set from imagenet
+    from examples.mlperf.dataloader import batch_load_resnet
+    iterator = batch_load_resnet(getenv("BS", 128*6), val=getenv("VAL", 1), shuffle=False, pad_first_batch=True)
+    def data_get():
+      x,y,cookie = next(iterator)
+      return x.shard(GPUS, axis=0).realize(), y, cookie
+    n,d = 0,0
+    proc = data_get()
+    tlog("loaded initial data")
+    st = time.perf_counter()
+    while proc is not None:
+      GlobalCounters.reset()
+      proc = (mdl(proc[0]), proc[1], proc[2])  # this frees the images
+      run = time.perf_counter()
+      # load the next data here
+      try: next_proc = data_get()
+      except StopIteration: next_proc = None
+      nd = time.perf_counter()
+      y = np.array(proc[1])
+      proc = (proc[0].numpy() == y) & (y != -1)  # this realizes the models and frees the cookies
+      n += proc.sum()
+      d += (y != -1).sum()
+      et = time.perf_counter()
+      tlog(f"****** {n:5d}/{d:5d}  {n*100.0/d:.2f}% -- {(run-st)*1000:7.2f} ms to enqueue, {(et-run)*1000:7.2f} ms to realize ({(nd-run)*1000:7.2f} ms fetching). {(len(proc))/(et-st):8.2f} examples/sec. {GlobalCounters.global_ops*1e-12/(et-st):5.2f} TFLOPS")
+      st = et
+      proc, next_proc = next_proc, None
+    tlog("done")

 def eval_unet3d():
  # UNet3D
@@ -81,47 +82,43 @@ def eval_unet3d():

 def eval_retinanet():
  # RetinaNet with ResNeXt50_32X4D
+  from examples.mlperf.dataloader import batch_load_retinanet
+  from extra.datasets.openimages import normalize, download_dataset, BASEDIR
  from extra.models.resnet import ResNeXt50_32X4D
  from extra.models.retinanet import RetinaNet
-  mdl = RetinaNet(ResNeXt50_32X4D())
-  mdl.load_from_pretrained()
-
-  input_mean = Tensor([0.485, 0.456, 0.406]).reshape(1, -1, 1, 1)
-  input_std = Tensor([0.229, 0.224, 0.225]).reshape(1, -1, 1, 1)
-  def input_fixup(x):
-    x = x.permute([0,3,1,2]) / 255.0
-    x -= input_mean
-    x /= input_std
-    return x
-
-  from extra.datasets.openimages import download_dataset, iterate, BASEDIR
  from pycocotools.coco import COCO
  from pycocotools.cocoeval import COCOeval
  from contextlib import redirect_stdout
-  coco = COCO(download_dataset(base_dir:=getenv("BASE_DIR", BASEDIR), 'validation'))
+  tlog("imports")
+
+  mdl = RetinaNet(ResNeXt50_32X4D())
+  mdl.load_from_pretrained()
+  tlog("loaded models")
+
+  coco = COCO(download_dataset(base_dir:=getenv("BASEDIR", BASEDIR), 'validation'))
  coco_eval = COCOeval(coco, iouType="bbox")
  coco_evalimgs, evaluated_imgs, ncats, narea = [], [], len(coco_eval.params.catIds), len(coco_eval.params.areaRng)
+  tlog("loaded dataset")

-  from tinygrad.engine.jit import TinyJit
-  mdlrun = TinyJit(lambda x: mdl(input_fixup(x)).realize())
-
-  n, bs = 0, 8
+  iterator = batch_load_retinanet(coco, True, Path(base_dir), getenv("BS", 8), shuffle=False)
+  def data_get():
+    x, img_ids, img_sizes, cookie = next(iterator)
+    return x.to(Device.DEFAULT).realize(), img_ids, img_sizes, cookie
+  n = 0
+  proc = data_get()
+  tlog("loaded initial data")
  st = time.perf_counter()
-  for x, targets in iterate(coco, base_dir, bs):
-    dat = Tensor(x.astype(np.float32))
-    mt = time.perf_counter()
-    if dat.shape[0] == bs:
-      outs = mdlrun(dat).numpy()
-    else:
-      mdlrun._jit_cache = []
-      outs =  mdl(input_fixup(dat)).numpy()
-    et = time.perf_counter()
-    predictions = mdl.postprocess_detections(outs, input_size=dat.shape[1:3], orig_image_sizes=[t["image_size"] for t in targets])
-    ext = time.perf_counter()
-    n += len(targets)
-    print(f"[{n}/{len(coco.imgs)}] == {(mt-st)*1000:.2f} ms loading data, {(et-mt)*1000:.2f} ms to run model, {(ext-et)*1000:.2f} ms for postprocessing")
-    img_ids = [t["image_id"] for t in targets]
-    coco_results  = [{"image_id": targets[i]["image_id"], "category_id": label, "bbox": box.tolist(), "score": score}
+  while proc is not None:
+    GlobalCounters.reset()
+    proc = (mdl(normalize(proc[0])), proc[1], proc[2], proc[3])
+    run = time.perf_counter()
+    # load the next data here
+    try: next_proc = data_get()
+    except StopIteration: next_proc = None
+    nd = time.perf_counter()
+    predictions, img_ids = mdl.postprocess_detections(proc[0].numpy(), orig_image_sizes=proc[2]), proc[1]
+    pd = time.perf_counter()
+    coco_results  = [{"image_id": img_ids[i], "category_id": label, "bbox": box.tolist(), "score": score}
      for i, prediction in enumerate(predictions) for box, score, label in zip(*prediction.values())]
    with redirect_stdout(None):
      coco_eval.cocoDt = coco.loadRes(coco_results)
@@ -129,13 +126,18 @@ def eval_retinanet():
      coco_eval.evaluate()
    evaluated_imgs.extend(img_ids)
    coco_evalimgs.append(np.array(coco_eval.evalImgs).reshape(ncats, narea, len(img_ids)))
-    st = time.perf_counter()
+    n += len(proc[0])
+    et = time.perf_counter()
+    tlog(f"****** {(run-st)*1000:7.2f} ms to enqueue, {(et-run)*1000:7.2f} ms to realize ({(nd-run)*1000:7.2f} ms fetching, {(pd-run)*1000:4.2f} ms postprocess_detections). {(len(proc))/(et-st):8.2f} examples/sec. {GlobalCounters.global_ops*1e-12/(et-st):5.2f} TFLOPS")
+    st = et
+    proc, next_proc = next_proc, None

  coco_eval.params.imgIds = evaluated_imgs
  coco_eval._paramsEval.imgIds = evaluated_imgs
  coco_eval.evalImgs = list(np.concatenate(coco_evalimgs, -1).flatten())
  coco_eval.accumulate()
  coco_eval.summarize()
+  tlog("done")

 def eval_rnnt():
  # RNN-T
@@ -242,7 +244,6 @@ def eval_mrcnn():
 if __name__ == "__main__":
  # inference only
  Tensor.training = False
-  Tensor.no_grad = True

  models = getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert,mrcnn").split(",")
  for m in models:
@@ -60,7 +60,6 @@ def spec_mrcnn():
 if __name__ == "__main__":
  # inference only for now
  Tensor.training = False
-  Tensor.no_grad = True

  for m in getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert,mrcnn").split(","):
    nm = f"spec_{m}"
@@ -1,14 +1,15 @@
-import os, time, math, functools
+import os, time, math, functools, random, contextlib
 from pathlib import Path
 import multiprocessing

 from tinygrad import Device, GlobalCounters, Tensor, TinyJit, dtypes
-from tinygrad.helpers import getenv, BEAM, WINO, round_up, diskcache_clear, FUSE_CONV_BW
+from tinygrad.helpers import getenv, BEAM, WINO, round_up, diskcache_clear, FUSE_CONV_BW, Profiling
 from tinygrad.nn.state import get_parameters, get_state_dict, safe_load, safe_save
-from tinygrad.nn.optim import LAMB, LARS, SGD, OptimizerGroup
+from tinygrad.nn.optim import LAMB, LARS, SGD, OptimizerGroup, Adam

 from extra.lr_scheduler import LRSchedulerGroup
 from examples.mlperf.helpers import get_training_state, load_training_state
+from extra.bench_log import BenchEvent, WallTimeEvent
 # TODO: fix benchmark logging and use tinygrad tqdm
 from tqdm import tqdm

@@ -205,24 +206,25 @@ def train_resnet():
    st = time.perf_counter()
    while proc is not None:
      GlobalCounters.reset()
-      (loss, top_1), y, proc = train_step(proc[0], proc[1]), proc[2], proc[3]
+      with WallTimeEvent(BenchEvent.STEP):
+        (loss, top_1), y, proc = train_step(proc[0], proc[1]), proc[2], proc[3]

-      pt = time.perf_counter()
+        pt = time.perf_counter()

-      if len(prev_cookies) == getenv("STORE_COOKIES", 1): prev_cookies = []  # free previous cookies after gpu work has been enqueued
-      try:
-        if INITMLPERF:
-          next_proc = fake_data_get(BS)
-        else:
-          next_proc = data_get(it)
-      except StopIteration:
-        next_proc = None
+        if len(prev_cookies) == getenv("STORE_COOKIES", 1): prev_cookies = []  # free previous cookies after gpu work has been enqueued
+        try:
+          if INITMLPERF:
+            next_proc = fake_data_get(BS)
+          else:
+            next_proc = data_get(it)
+        except StopIteration:
+          next_proc = None

-      dt = time.perf_counter()
+        dt = time.perf_counter()

-      device_str = loss.device if isinstance(loss.device, str) else f"{loss.device[0]} * {len(loss.device)}"
-      loss, top_1 = loss.numpy().item(), top_1.numpy().item()
-      top_1_acc = top_1 / sum(yi != -1 for yi in y)
+        device_str = loss.device if isinstance(loss.device, str) else f"{loss.device[0]} * {len(loss.device)}"
+        loss, top_1 = loss.numpy().item(), top_1.numpy().item()
+        top_1_acc = top_1 / sum(yi != -1 for yi in y)

      cl = time.perf_counter()
      if BENCHMARK:
@@ -343,8 +345,349 @@ def train_resnet():
        safe_save(get_training_state(model, optimizer_group, scheduler_group), fn)

 def train_retinanet():
-  # TODO: Retinanet
-  pass
+  from contextlib import redirect_stdout
+  from examples.mlperf.dataloader import batch_load_retinanet
+  from examples.mlperf.initializers import FrozenBatchNorm2dRetinaNet, Conv2dNormalRetinaNet, Conv2dKaimingUniformRetinaNet, Linear, Conv2dRetinaNet
+  from extra.datasets.openimages import MLPERF_CLASSES, BASEDIR, download_dataset, normalize, get_dataset_count
+  from extra.models import resnet, retinanet
+  from pycocotools.coco import COCO
+  from pycocotools.cocoeval import COCOeval
+  from tinygrad.helpers import colored
+  from typing import Iterator
+
+  import numpy as np
+
+  config, target_metric = {}, 0.34
+
+  config["SEED"] = SEED = getenv("SEED", random.SystemRandom().randint(0, 2**32 - 1))
+  Tensor.manual_seed(SEED)
+
+  NUM_CLASSES = len(MLPERF_CLASSES)
+  BASEDIR = getenv("BASEDIR", BASEDIR)
+  BENCHMARK = getenv("BENCHMARK")
+  INITMLPERF = getenv("INITMLPERF")
+  RUNMLPERF = getenv("RUNMLPERF")
+
+  if INITMLPERF:
+    diskcache_clear()
+
+  if getenv("LOGMLPERF"):
+    from mlperf_logging import mllog
+    import mlperf_logging.mllog.constants as mllog_constants
+
+    mllog.config(filename=f"result_retinanet_{SEED}.log")
+    mllog.config(root_dir=Path(__file__).parents[3].as_posix())
+    MLLOGGER = mllog.get_mllogger()
+    MLLOGGER.logger.propagate = False
+
+    if INITMLPERF:
+      assert BENCHMARK, "BENCHMARK must be set for INITMLPERF"
+      MLLOGGER.event(key=mllog_constants.SUBMISSION_ORG, value="tinycorp")
+      MLLOGGER.event(key=mllog_constants.SUBMISSION_PLATFORM, value=getenv("SUBMISSION_PLATFORM", "tinybox"))
+      MLLOGGER.event(key=mllog_constants.SUBMISSION_DIVISION, value=mllog_constants.CLOSED)
+      MLLOGGER.event(key=mllog_constants.SUBMISSION_STATUS, value=mllog_constants.ONPREM)
+
+      MLLOGGER.event(key=mllog_constants.SUBMISSION_BENCHMARK, value=mllog_constants.RETINANET)
+
+      MLLOGGER.event(key=mllog_constants.CACHE_CLEAR, value=True)
+      MLLOGGER.start(key=mllog_constants.INIT_START)
+
+    if RUNMLPERF:
+      MLLOGGER.start(key=mllog_constants.RUN_START)
+      MLLOGGER.event(key=mllog_constants.SEED, value=SEED)
+  else:
+    MLLOGGER = None
+
+  config["gpus"] = GPUS = [f"{Device.DEFAULT}:{i}" for i in range(getenv("GPUS", 6))]
+
+  for x in GPUS: Device[x]
+  print(f"training on {GPUS}")
+
+  def _freeze_backbone_layers(backbone:resnet.ResNet, trainable_layers:int):
+    layers_to_train = ["layer4", "layer3", "layer2", "layer1", "conv1"][:trainable_layers]
+    for k, v in get_state_dict(backbone).items():
+      if all([not k.startswith(layer) for layer in layers_to_train]):
+        v.requires_grad = False
+
+  def _data_get(it:Iterator[tuple[Tensor, ...]], val:bool=False):
+    if val:
+      x, img_ids, img_sizes, cookie = next(it)
+      return x.shard(GPUS, axis=0), img_ids, img_sizes, cookie
+
+    x, y_boxes, y_labels, matches, anchors, cookie = next(it)
+    return x.shard(GPUS, axis=0), y_boxes.shard(GPUS, axis=0), y_labels.shard(GPUS, axis=0), matches.shard(GPUS, axis=0), anchors.shard(GPUS, axis=0), cookie
+
+  def _fake_data_get(bs:int, val:bool=False):
+    x = Tensor.empty(bs, 800, 800, 3, dtype=dtypes.uint8)
+    if val:
+      img_ids, img_sizes = [0] * bs, [(800, 800)] * bs
+      return x.shard(GPUS, axis=0), img_ids, img_sizes, None
+
+    y_boxes = Tensor.empty(bs, 120087, 4, dtype=dtypes.float32)
+    y_labels = Tensor.empty(bs, 120087, dtype=dtypes.int64)
+    matches = Tensor.empty(bs, 120087, dtype=dtypes.int64)
+    anchors = Tensor.empty(bs, 120087, 4, dtype=dtypes.float64)
+    return x.shard(GPUS, axis=0), y_boxes.shard(GPUS, axis=0), y_labels.shard(GPUS, axis=0), matches.shard(GPUS, axis=0), anchors.shard(GPUS, axis=0), None
+
+  @TinyJit
+  def _train_step(model, optim, loss_scaler, x, **kwargs):
+    optim.zero_grad()
+
+    losses = model(normalize(x, GPUS), **kwargs)
+    loss = sum(losses.values())
+
+    (loss * loss_scaler).backward()
+    for t in optim.params: t.grad = t.grad / loss_scaler
+
+    optim.step()
+
+    return loss.realize(), losses
+
+  @TinyJit
+  def _eval_step(model, x, **kwargs):
+    out = model(normalize(x, GPUS), **kwargs)
+    # reassemble on GPUS[0] before sending back to CPU for speed
+    return out.to(GPUS[0]).realize()
+
+  # ** hyperparameters **
+  config["BS"] = BS = getenv("BS", 16 * len(GPUS) if dtypes.default_float == dtypes.float16 else 12 * len(GPUS))
+  config["EVAL_BS"] = EVAL_BS = getenv("EVAL_BS", BS)
+  config["EPOCHS"] = EPOCHS = getenv("EPOCHS", 4)
+  config["TRAIN_BEAM"] = TRAIN_BEAM = getenv("TRAIN_BEAM", BEAM.value)
+  config["EVAL_BEAM"] = EVAL_BEAM = getenv("EVAL_BEAM", BEAM.value)
+  config["LR"] = lr = getenv("LR", 9.5e-5 * (BS / 96))
+  config["LOSS_SCALER"] = loss_scaler = getenv("LOSS_SCALER", 2**11 if dtypes.default_float == dtypes.float16 else 1.0)
+  config["DEFAULT_FLOAT"] = dtypes.default_float.name
+  config["EVAL_FREQ"] = eval_freq = getenv("EVAL_FREQ", 1)
+
+  # ** initialize wandb **
+  if (WANDB:=getenv("WANDB")):
+    import wandb
+    wandb.init(config=config, project="MLPerf-RetinaNet")
+
+  # ** model initializers **
+  resnet.BatchNorm = FrozenBatchNorm2dRetinaNet
+  resnet.Linear = Linear
+  resnet.Conv2d = Conv2dRetinaNet
+
+  retinanet.ConvHead = Conv2dNormalRetinaNet
+  retinanet.ConvClassificationHeadLogits = functools.partial(Conv2dNormalRetinaNet, prior_prob=0.01)
+  retinanet.ConvFPN = Conv2dKaimingUniformRetinaNet
+
+  # ** model setup **
+  backbone = resnet.ResNeXt50_32X4D(num_classes=None)
+  if RUNMLPERF:
+    backbone.load_from_pretrained()
+  _freeze_backbone_layers(backbone, 3)
+
+  model = retinanet.RetinaNet(backbone, num_classes=NUM_CLASSES)
+  params = get_parameters(model)
+
+  if not RUNMLPERF:
+    # for init, zero out all weights
+    for p in params:
+      p = p.assign(Tensor.zeros_like(p).contiguous()).realize()
+
+  if len(GPUS) > 1:
+    for p in params: p.to_(GPUS)
+
+  step_times, start_epoch = [], 0
+
+  # ** optimizer **
+  optim = Adam(params, lr=lr)
+
+  # ** dataset **
+  config["STEPS_IN_TRAIN_EPOCH"] = steps_in_train_epoch = round_up(get_dataset_count((base_dir_path:=Path(BASEDIR)), False), BS) // BS
+  config["STEPS_IN_VAL_EPOCH"] = steps_in_val_epoch = (round_up(get_dataset_count(base_dir_path, True), EVAL_BS) // EVAL_BS)
+
+  # log mlperf hparams
+  if MLLOGGER:
+    if RUNMLPERF:
+      MLLOGGER.event(key=mllog_constants.GLOBAL_BATCH_SIZE, value=config["BS"])
+      MLLOGGER.event(key=mllog_constants.TRAIN_SAMPLES, value=config["STEPS_IN_TRAIN_EPOCH"])
+      MLLOGGER.event(key=mllog_constants.EVAL_SAMPLES, value=config["STEPS_IN_VAL_EPOCH"])
+      MLLOGGER.event(key=mllog_constants.EPOCH_COUNT, value=config["EPOCHS"])
+      MLLOGGER.event(key=mllog_constants.FIRST_EPOCH_NUM, value=start_epoch)
+
+      MLLOGGER.event(key=mllog_constants.OPT_NAME, value=mllog_constants.ADAM)
+      MLLOGGER.event(key=mllog_constants.OPT_BASE_LR, value=config["LR"])
+      MLLOGGER.event(key=mllog_constants.OPT_WEIGHT_DECAY, value=0)
+      MLLOGGER.event(key=mllog_constants.OPT_LR_WARMUP_EPOCHS, value=0)
+      MLLOGGER.event(key=mllog_constants.OPT_LR_WARMUP_FACTOR, value=0)
+      MLLOGGER.event(key=mllog_constants.GRADIENT_ACCUMULATION_STEPS, value=1)
+
+  if RUNMLPERF:
+    train_dataset = COCO(download_dataset(BASEDIR, "train"))
+    val_dataset = COCO(download_dataset(BASEDIR, "validation"))
+    coco_val = COCOeval(cocoGt=val_dataset, iouType="bbox")
+
+  print(f"training with batch size {BS} for {EPOCHS} epochs")
+
+  for e in range(start_epoch, EPOCHS):
+    # ** training loop **
+    if MLLOGGER and RUNMLPERF:
+      MLLOGGER.start(key=mllog_constants.EPOCH_START, value=e + 1, metadata={"epoch_num": e + 1})
+
+    BEAM.value = TRAIN_BEAM
+
+    if not RUNMLPERF:
+      i, proc = 0, _fake_data_get(BS)
+    else:
+      train_dataloader = batch_load_retinanet(train_dataset, False, base_dir_path, batch_size=BS, seed=SEED)
+      it = iter(tqdm(train_dataloader, total=steps_in_train_epoch, desc=f"epoch {e + 1}", disable=BENCHMARK))
+      i, proc = 0, _data_get(it)
+
+    prev_cookies = []
+    st = time.perf_counter()
+
+    while proc is not None:
+      GlobalCounters.reset()
+
+      x, y_bboxes, y_labels, matches, anchors, proc = proc
+      loss, losses = _train_step(model, optim, loss_scaler, x, labels=y_labels, matches=matches, anchors=anchors, bboxes=y_bboxes)
+
+      pt = time.perf_counter()
+
+      if len(prev_cookies) == getenv("STORE_COOKIES", 1): prev_cookies = []  # free previous cookies after gpu work has been enqueued
+      try:
+        if not RUNMLPERF:
+          next_proc = _fake_data_get(BS)
+        else:
+          next_proc = _data_get(it)
+      except StopIteration:
+        next_proc = None
+
+      dt = time.perf_counter()
+
+      device_str = loss.device if isinstance(loss.device, str) else f"{loss.device[0]} * {len(loss.device)}"
+      loss = loss.item()
+
+      cl = time.perf_counter()
+      if BENCHMARK: step_times.append(cl - st)
+
+      if not math.isfinite(loss):
+        print("loss is nan")
+        return
+
+      tqdm.write(
+        f"{i:5} {((cl - st)) * 1000.0:7.2f} ms run, {(pt - st) * 1000.0:7.2f} ms python, {(dt - pt) * 1000.0:6.2f} ms fetch data, "
+        f"{(cl - dt) * 1000.0:7.2f} ms {device_str}, {loss:5.2f} loss, {losses['classification_loss'].item():5.4f} classification loss, {losses['regression_loss'].item():5.4f} regression loss, "
+        f"{optim.lr.numpy()[0]:.6f} LR, {GlobalCounters.mem_used / 1e9:.2f} GB used, {GlobalCounters.global_ops * 1e-9 / (cl - st):9.2f} GFLOPS"
+      )
+
+      if WANDB:
+        wandb.log({"lr": optim.lr.numpy(), "train/loss": loss, "train/classification_loss": losses["classification_loss"].item(), "train/regression_loss": losses["regression_loss"].item(),
+                  "train/step_time": cl - st, "train/python_time": pt - st, "train/data_time": dt - pt, "train/cl_time": cl - dt,
+                  "train/GFLOPS": GlobalCounters.global_ops * 1e-9 / (cl - st), "epoch": e + (i + 1) / steps_in_train_epoch})
+
+      st = cl
+      prev_cookies.append(proc)
+      proc, next_proc = next_proc, None  # return old cookie
+      i += 1
+
+      if i == BENCHMARK:
+        assert not math.isnan(loss)
+        median_step_time = sorted(step_times)[(BENCHMARK + 1) // 2]  # in seconds
+        estimated_total_minutes = int(median_step_time * steps_in_train_epoch * EPOCHS / 60)
+        print(f"Estimated training time: {estimated_total_minutes // 60}h{estimated_total_minutes % 60}m")
+        print(f"epoch global_ops: {steps_in_train_epoch * GlobalCounters.global_ops:_}, "
+              f"epoch global_mem: {steps_in_train_epoch * GlobalCounters.global_mem:_}")
+        # if we are doing beam search, run the first eval too
+        if (TRAIN_BEAM or EVAL_BEAM) and e == start_epoch: break
+        return
+
+    if MLLOGGER and RUNMLPERF:
+      MLLOGGER.event(key=mllog_constants.EPOCH_STOP, value=e + 1, metadata={"epoch_num": e + 1})
+
+    # ** eval loop **
+    if (e + 1) % eval_freq == 0:
+      if MLLOGGER and RUNMLPERF:
+        MLLOGGER.start(key=mllog_constants.EVAL_START, value=e + 1, metadata={"epoch_num": e + 1})
+
+      BEAM.value = EVAL_BEAM
+
+      if getenv("RESET_STEP", 1): _train_step.reset()
+
+      with Tensor.train(mode=False):
+        if not RUNMLPERF:
+          i, proc = 0, _fake_data_get(EVAL_BS, val=(val:=True))
+        else:
+          val_dataloader = batch_load_retinanet(val_dataset, (val:=True), Path(BASEDIR), batch_size=EVAL_BS, shuffle=False, seed=SEED)
+          it = iter(tqdm(val_dataloader, total=steps_in_val_epoch))
+          i, proc = 0, _data_get(it, val=val)
+          val_img_ids, val_imgs, ncats, narea = [], [], len(coco_val.params.catIds), len(coco_val.params.areaRng)
+
+        eval_times, prev_cookies = [], []
+
+        while proc is not None:
+          GlobalCounters.reset()
+          st = time.time()
+
+          out, img_ids, img_sizes, proc = _eval_step(model, (x:=proc[0])).numpy(), proc[1], proc[2], proc[3]
+
+          if RUNMLPERF:
+            out = model.postprocess_detections(out, input_size=x.shape[1:3], orig_image_sizes=img_sizes)
+            coco_results  = [{"image_id": img_ids[i], "category_id": label, "bbox": box.tolist(), "score": score}
+              for i, prediction in enumerate(out) for box, score, label in zip(*prediction.values())]
+
+            with redirect_stdout(None):
+              coco_val.cocoDt = val_dataset.loadRes(coco_results)
+              coco_val.params.imgIds = img_ids
+              coco_val.evaluate()
+
+            val_img_ids.extend(img_ids)
+            val_imgs.append(np.array(coco_val.evalImgs).reshape(ncats, narea, len(img_ids)))
+
+          if len(prev_cookies) == getenv("STORE_COOKIES", 1): prev_cookies = []  # free previous cookies after gpu work has been enqueued
+          try:
+            if not RUNMLPERF:
+              next_proc = _fake_data_get(EVAL_BS, val=val)
+            else:
+              next_proc = _data_get(it, val=val)
+          except StopIteration:
+            next_proc = None
+
+          prev_cookies.append(proc)
+          proc, next_proc = next_proc, None
+          i += 1
+
+          et = time.time()
+          eval_times.append(et - st)
+
+          if i == BENCHMARK:
+            # assume INITMLPERF has BENCHMARK set
+            if MLLOGGER and INITMLPERF:
+              MLLOGGER.event(key=mllog_constants.INIT_STOP)
+            return
+
+        if getenv("RESET_STEP", 1): _eval_step.reset()
+        total_fw_time = sum(eval_times) / len(eval_times)
+
+        if RUNMLPERF:
+          coco_val.params.imgIds = val_img_ids
+          coco_val._paramsEval.imgIds = val_img_ids
+          coco_val.evalImgs = list(np.concatenate(val_imgs, -1).flatten())
+          coco_val.accumulate()
+          coco_val.summarize()
+
+          val_metric = coco_val.stats[0]
+
+          tqdm.write(f"eval time: {total_fw_time:.2f}, eval metric: {val_metric:.4f}")
+
+          if WANDB:
+            wandb.log({"eval/forward_time": total_fw_time, "eval/metric": val_metric, "epoch": e + 1})
+
+          if MLLOGGER:
+            MLLOGGER.event(key=mllog_constants.EVAL_ACCURACY, value=val_metric, metadata={"epoch_num": e + 1}, clear_line=True)
+            MLLOGGER.end(key=mllog_constants.EVAL_STOP, value=e + 1, metadata={"epoch_num": e + 1})
+
+          if val_metric >= target_metric:
+            print(colored(f"target metric reached: {val_metric:.2f}/{target_metric:.2f}", color="green"))
+
+            if MLLOGGER:
+              MLLOGGER.end(key=mllog_constants.RUN_STOP, metadata={"status": mllog_constants.SUCCESS})
+
+            break

 def train_unet3d():
  """
@@ -448,7 +791,6 @@ def train_unet3d():
    return loss.realize()

  @Tensor.train(mode=False)
-  @Tensor.test()
  def eval_step(model, x, y):
    y_hat, y = sliding_window_inference(model, x, y, gpus=GPUS)
    y_hat, y = Tensor(y_hat), Tensor(y, requires_grad=False)
@@ -551,7 +893,7 @@ def train_unet3d():

        if mean_dice >= TARGET_METRIC:
          is_successful = True
-          save_checkpoint(get_state_dict(model), f"./ckpts/unet3d.safe")
+          save_checkpoint(get_state_dict(model), "./ckpts/unet3d.safe")
        elif mean_dice < 1e-6:
          print("Model diverging. Aborting.")
          diverged = True
@@ -583,12 +925,13 @@ def train_step_bert(model, optimizer, scheduler, loss_scaler:float, input_ids:Te
  loss = model.loss(lm_logits, seq_relationship_logits, masked_lm_ids, masked_lm_weights, next_sentence_labels)
  (loss * loss_scaler).backward()

-  global_norm = Tensor([0.0], dtype=dtypes.float32, device=optimizer[0].device).realize()
+  global_norm = Tensor([0.0], dtype=dtypes.float32, device=optimizer[0].device)
  for p in optimizer.params:
    p.grad = p.grad / loss_scaler
    global_norm += p.grad.float().square().sum()
-  global_norm = global_norm.sqrt()
-  for p in optimizer.params: p.grad = (p.grad / Tensor.where(global_norm > 1.0, global_norm, 1.0)).cast(p.grad.dtype)
+  global_norm = global_norm.sqrt().contiguous()
+  for p in optimizer.params:
+    p.grad = (global_norm > 1.0).where((p.grad/global_norm).cast(p.grad.dtype), p.grad)

  optimizer.step()
  scheduler.step()
@@ -637,7 +980,7 @@ def train_bert():
    MLLOGGER.logger.propagate = False

    if INITMLPERF:
-      assert BENCHMARK, f"BENCHMARK must be set for INITMLPERF"
+      assert BENCHMARK, "BENCHMARK must be set for INITMLPERF"
      MLLOGGER.event(key=mllog_constants.SUBMISSION_ORG, value="tinycorp")
      MLLOGGER.event(key=mllog_constants.SUBMISSION_PLATFORM, value=getenv("SUBMISSION_PLATFORM", "tinybox"))
      MLLOGGER.event(key=mllog_constants.SUBMISSION_DIVISION, value=mllog_constants.CLOSED)
@@ -659,8 +1002,10 @@ def train_bert():
  BS                 = config["GLOBAL_BATCH_SIZE"]      = getenv("BS", 11 * len(GPUS) if dtypes.default_float in (dtypes.float16, dtypes.bfloat16) else 8 * len(GPUS))
  EVAL_BS            = config["EVAL_BS"]                = getenv("EVAL_BS", 1 * len(GPUS))
  max_lr             = config["OPT_BASE_LEARNING_RATE"] = getenv("OPT_BASE_LEARNING_RATE", 0.000175 * math.sqrt(BS/96))
+  opt_lamb_beta_1    = config["OPT_LAMB_BETA_1"]        = getenv("OPT_LAMB_BETA_1", 0.9)
+  opt_lamb_beta_2    = config["OPT_LAMB_BETA_2"]        = getenv("OPT_LAMB_BETA_2", 0.999)

-  train_steps        = config["TRAIN_STEPS"]            = getenv("TRAIN_STEPS", 3300000 // BS)
+  train_steps        = config["TRAIN_STEPS"]            = getenv("TRAIN_STEPS", 3600000 // BS)
  warmup_steps       = config["NUM_WARMUP_STEPS"]       = getenv("NUM_WARMUP_STEPS", 1)
  max_eval_steps     = config["MAX_EVAL_STEPS"]         = getenv("MAX_EVAL_STEPS", (10000 + EVAL_BS - 1) // EVAL_BS) # EVAL_BS * MAX_EVAL_STEPS >= 10000
  eval_step_freq     = config["EVAL_STEP_FREQ"]         = getenv("EVAL_STEP_FREQ", int((math.floor(0.05 * (230.23 * BS + 3000000) / 25000) * 25000) / BS)) # Round down
@@ -713,8 +1058,8 @@ def train_bert():
  # ** Optimizer **
  parameters_no_wd = [v for k, v in get_state_dict(model).items() if "bias" in k or "LayerNorm" in k]
  parameters = [x for x in parameters if x not in set(parameters_no_wd)]
-  optimizer_wd = LAMB(parameters, lr=max_lr, eps=epsilon, weight_decay=decay, adam=False)
-  optimizer_no_wd = LAMB(parameters_no_wd, lr=max_lr, eps=epsilon, weight_decay=0.0, adam=False)
+  optimizer_wd = LAMB(parameters, lr=max_lr, b1=opt_lamb_beta_1, b2=opt_lamb_beta_2, eps=epsilon, weight_decay=decay, adam=False)
+  optimizer_no_wd = LAMB(parameters_no_wd, lr=max_lr, b1=opt_lamb_beta_1, b2=opt_lamb_beta_2, eps=epsilon, weight_decay=0.0, adam=False)
  optimizer_group = OptimizerGroup(optimizer_wd, optimizer_no_wd)

  # ** LR scheduler **
@@ -733,8 +1078,8 @@ def train_bert():
      MLLOGGER.event(key=mllog_constants.OPT_NAME, value="LAMB")
      MLLOGGER.event(key=mllog_constants.OPT_BASE_LR, value=config["OPT_BASE_LEARNING_RATE"])
      MLLOGGER.event(key=mllog_constants.OPT_LAMB_WEIGHT_DECAY, value=config["DECAY"])
-      MLLOGGER.event(key=mllog_constants.OPT_LAMB_BETA_1, value=optimizer_wd.b1)
-      MLLOGGER.event(key=mllog_constants.OPT_LAMB_BETA_2, value=optimizer_wd.b2)
+      MLLOGGER.event(key=mllog_constants.OPT_LAMB_BETA_1, value=config["OPT_LAMB_BETA_1"])
+      MLLOGGER.event(key=mllog_constants.OPT_LAMB_BETA_2, value=config["OPT_LAMB_BETA_2"])
      MLLOGGER.event(key=mllog_constants.OPT_LAMB_LR_DECAY_POLY_POWER, value=config["POLY_POWER"])
      MLLOGGER.event(key=mllog_constants.OPT_LAMB_EPSILON, value=config["EPSILON"])

@@ -782,22 +1127,24 @@ def train_bert():
      BEAM.value = TRAIN_BEAM
      st = time.perf_counter()
      GlobalCounters.reset()
-      loss, global_norm, lr = train_step_bert(model, optimizer_group, scheduler_group, loss_scaler,
-        train_data["input_ids"], train_data["segment_ids"], train_data["input_mask"], train_data["masked_lm_positions"], \
-        train_data["masked_lm_ids"], train_data["masked_lm_weights"], train_data["next_sentence_labels"], GPUS)
+      with WallTimeEvent(BenchEvent.STEP):
+        loss, global_norm, lr = train_step_bert(model, optimizer_group, scheduler_group, loss_scaler,
+          train_data["input_ids"], train_data["segment_ids"], train_data["input_mask"], train_data["masked_lm_positions"], \
+          train_data["masked_lm_ids"], train_data["masked_lm_weights"], train_data["next_sentence_labels"], GPUS)

-      pt = time.perf_counter()
+        pt = time.perf_counter()

-      try:
-        next_data = next(train_it)
-      except StopIteration:
-        next_data = None
+        try:
+          next_data = next(train_it)
+        except StopIteration:
+          next_data = None

-      dt = time.perf_counter()
+        dt = time.perf_counter()

-      device_str = parameters[0].device if isinstance(parameters[0].device, str) else f"{parameters[0].device[0]} * {len(parameters[0].device)}"
-      loss = loss.item()
-      lr = lr.item()
+        device_str = parameters[0].device if isinstance(parameters[0].device, str) else f"{parameters[0].device[0]} * {len(parameters[0].device)}"
+        loss = loss.item()
+        assert not math.isnan(loss)
+        lr = lr.item()

      cl = time.perf_counter()
      if BENCHMARK: step_times.append(cl - st)
@@ -825,8 +1172,8 @@ def train_bert():
    if i % eval_step_freq == 0 or (BENCHMARK and i == BENCHMARK) or i == train_steps:
      if MLLOGGER and RUNMLPERF:
        MLLOGGER.start(key=mllog_constants.EVAL_START, value=None, metadata={"epoch_num": i*BS, "step_num": i})
-      if getenv("RESET_STEP", 0): train_step_bert.reset()
-      elif train_step_bert.captured is not None: train_step_bert.captured.free_intermediates()
+      if getenv("RESET_STEP"): train_step_bert.reset()
+      elif getenv("FREE_INTERMEDIATE", 1) and train_step_bert.captured is not None: train_step_bert.captured.free_intermediates()
      eval_lm_losses = []
      eval_clsf_losses = []
      eval_lm_accs = []
@@ -853,14 +1200,15 @@ def train_bert():
        et = time.time()
        eval_times.append(et - st)

-        if BENCHMARK and j == BENCHMARK:
+        if BENCHMARK and (j+1) == min(BENCHMARK, max_eval_steps):
          # assume INITMLPERF has BENCHMARK set
          if MLLOGGER and INITMLPERF:
            MLLOGGER.event(key=mllog_constants.INIT_STOP, value=None)
          return

-      if getenv("RESET_STEP", 0): eval_step_bert.reset()
-      elif eval_step_bert.captured is not None: eval_step_bert.captured.free_intermediates()
+      if getenv("RESET_STEP"): eval_step_bert.reset()
+      elif getenv("FREE_INTERMEDIATE", 1) and eval_step_bert.captured is not None: eval_step_bert.captured.free_intermediates()
+
      del eval_data
      avg_lm_loss = sum(eval_lm_losses) / len(eval_lm_losses)
      avg_clsf_loss = sum(eval_clsf_losses) / len(eval_clsf_losses)
@@ -873,7 +1221,7 @@ def train_bert():

      if WANDB:
        wandb.log({"eval/lm_loss": avg_lm_loss, "eval/clsf_loss": avg_clsf_loss, "eval/lm_accuracy": avg_lm_acc, \
-                    "eval/clsf_accuracy": avg_clsf_acc, "eval/forward_time": avg_fw_time})
+                    "eval/clsf_accuracy": avg_clsf_acc, "eval/forward_time": avg_fw_time, "epoch": (i+1)*BS})

      if MLLOGGER and RUNMLPERF:
        MLLOGGER.end(key=mllog_constants.EVAL_STOP, value=i*BS, metadata={"epoch_count": i*BS, "step_num": i, "samples_count": config["EVAL_BS"] * config["MAX_EVAL_STEPS"]})
@@ -900,11 +1248,14 @@ def train_bert():
        # stop once hitting the target
        break

+    # should not happen, BENCHMARK not properly terminated
+    if BENCHMARK: assert i < BENCHMARK, i
+
    if getenv("CKPT") and i % save_ckpt_freq == 0:
      if MLLOGGER and RUNMLPERF:
        if previous_step:
          MLLOGGER.end(key=mllog_constants.BLOCK_STOP, value=None, metadata={"first_epoch_num": 1, "epoch_num": 1, "first_step_num": i, "step_num": i, "step_count": i - previous_step})
-        MLLOGGER.start(key="checkpoint_start", value=None, metadata={"step_num" : i})
+        MLLOGGER.start(key="checkpoint_start", value=None, metadata={"step_num": i})
      if not os.path.exists(ckpt_dir := save_ckpt_dir): os.mkdir(ckpt_dir)
      if WANDB and wandb.run is not None:
        fn = f"{ckpt_dir}/{time.strftime('%Y%m%d_%H%M%S')}_{wandb.run.id}.safe"
@@ -929,9 +1280,15 @@ def train_maskrcnn():

 if __name__ == "__main__":
  multiprocessing.set_start_method('spawn')
+
+  if getenv("INITMLPERF"): bench_log_manager = WallTimeEvent(BenchEvent.MLPERF_INIT)
+  elif getenv("RUNMLPERF"): bench_log_manager = WallTimeEvent(BenchEvent.MLPERF_RUN)
+  else: bench_log_manager = contextlib.nullcontext()
+
  with Tensor.train():
    for m in getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert,maskrcnn").split(","):
      nm = f"train_{m}"
      if nm in globals():
        print(f"training {m}")
-        globals()[nm]()
+        with bench_log_manager:
+          with Profiling(enabled=getenv("PYPROFILE")): globals()[nm]()
@@ -0,0 +1,15 @@
+#!/bin/bash
+
+export PYTHONPATH="." AMD=1
+export MODEL="bert"
+export DEFAULT_FLOAT="HALF" GPUS=1 BS=128 EVAL_BS=128
+
+export BEAM=3 BEAM_UOPS_MAX=4000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export IGNORE_JIT_FIRST_BEAM=1
+# export BEAM_LOG_SURPASS_MAX=1
+# export BASEDIR="/raid/datasets/wiki"
+
+export RESET_STEP=1
+export BENCHMARK=10 BERT_LAYERS=2 DEBUG=2
+
+python3 examples/mlperf/model_train.py
@@ -0,0 +1,69 @@
+# 1. Problem
+
+This problem uses BERT for NLP.
+
+## Requirements
+
+Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0.
+```
+git clone https://github.com/tinygrad/tinygrad.git
+python3 -m pip install -e ".[mlperf]"
+```
+Also install gdown (for dataset), numpy, tqdm and tensorflow.
+```
+pip install gdown numpy tqdm tensorflow
+```
+
+### tinybox_green
+Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
+This is the default on production tinybox green.
+
+# 2. Directions
+
+## Steps to download and verify data
+
+### 1. Download raw data
+
+```
+BASEDIR="/raid/datasets/wiki" WIKI_TRAIN=1 VERIFY_CHECKSUM=1 python3 extra/datasets/wikipedia_download.py
+```
+
+### 2. Preprocess train and validation data
+
+Note: The number of threads used for preprocessing is limited by available memory. With 128GB of RAM, a maximum of 16 threads is recommended. 
+
+#### Training:
+```
+BASEDIR="/raid/datasets/wiki" NUM_WORKERS=16 python3 extra/datasets/wikipedia.py pre-train all
+```
+
+Generating a specific topic (Between 0 and 499)
+```
+BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-train 42
+```
+
+#### Validation:
+```
+BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-eval
+```
+## Running
+
+### tinybox_green
+
+#### Steps to run benchmark
+```
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh
+```
+
+### tinybox_red
+
+#### Steps to run benchmark
+```
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh
+```
+### tinybox_8xMI300X
+
+#### Steps to run benchmark
+```
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/run_and_time.sh
+```
@@ -0,0 +1,14 @@
+#!/bin/bash
+
+export PYTHONPATH="." AMD=1
+export MODEL="bert"
+export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
+export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
+
+export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
+export BASEDIR="/raid/datasets/wiki"
+
+export BENCHMARK=10 BERT_LAYERS=2 DEBUG=2
+
+python3 examples/mlperf/model_train.py
@@ -0,0 +1,17 @@
+#!/bin/bash
+
+export PYTHONPATH="." AMD=1
+export MODEL="bert"
+export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
+
+# similar to https://github.com/mlcommons/training_results_v3.1/blob/d06288b2bd675a9d88e0e6181f5bb5626b71ec19/Quanta_Cloud_Technology/results/D54U-3U/bert/result_1.txt#L54
+export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
+export TRAIN_STEPS=3900
+
+export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
+export BASEDIR="/raid/datasets/wiki"
+
+export WANDB=1 PARALLEL=0
+
+RUNMLPERF=1 python3 examples/mlperf/model_train.py
@@ -0,0 +1,28 @@
+#!/bin/bash
+set -e  # Exit on any error
+
+export PYTHONPATH="." AMD=1
+export MODEL="bert"
+export SUBMISSION_PLATFORM="tinybox_8xMI300X"
+export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
+
+# similar to https://github.com/mlcommons/training_results_v3.1/blob/d06288b2bd675a9d88e0e6181f5bb5626b71ec19/Quanta_Cloud_Technology/results/D54U-3U/bert/result_1.txt#L54
+export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
+export TRAIN_STEPS=3900
+
+export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
+export BASEDIR="/raid/datasets/wiki"
+
+# pip install -e ".[mlperf]"
+export LOGMLPERF=1
+
+export SEED=$RANDOM
+DATETIME=$(date "+%m%d%H%M")
+LOGFILE="bert_8xMI300x_${DATETIME}_${SEED}.log"
+
+# init  # TODO: without DEBUG=2 it hangs
+BENCHMARK=10 INITMLPERF=1 BERT_LAYERS=2 DEBUG=2 python3 examples/mlperf/model_train.py | tee $LOGFILE
+
+# run
+PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
@@ -4,24 +4,20 @@ This problem uses BERT for NLP.

 ## Requirements

-Install tinygrad and mlperf-logging from master.
+Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0.
 ```
 git clone https://github.com/tinygrad/tinygrad.git
 python3 -m pip install -e ".[mlperf]"
 ```
-Also install tqdm and tensorflow.
+Also install gdown (for dataset), numpy, tqdm and tensorflow.
 ```
-pip install tqdm tensorflow
+pip install gdown numpy tqdm tensorflow
 ```

 ### tinybox_green
 Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
 This is the default on production tinybox green.

-### tinybox_red
-Disable cwsr + increase mes timeout.
-Install the custom amdgpu driver per [README](https://github.com/nimlgen/amdgpu_ubuntu_22_04/blob/v6.1.3/readme.md)
-
 # 2. Directions

 ## Steps to download and verify data
@@ -56,18 +52,18 @@ BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-eval

 #### Steps to run benchmark
 ```
-examples/mlperf/training_submission_v4.1/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh
 ```

 ### tinybox_red

-#### One time setup
-
+#### Steps to run benchmark
 ```
-examples/mlperf/training_submission_v4.1/tinycorp/benchmarks/bert/implementations/tinybox_red/setup.sh
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh
 ```
+### tinybox_8xMI300X

 #### Steps to run benchmark
 ```
-examples/mlperf/training_submission_v4.1/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/run_and_time.sh
 ```
@@ -1,14 +1,16 @@
 #!/bin/bash

-export PYTHONPATH="."
+export PYTHONPATH="." NV=1
 export MODEL="bert"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=24
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96

-export BEAM=4 BEAM_UOPS_MAX=4000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
+
+export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 export IGNORE_JIT_FIRST_BEAM=1
 export BEAM_LOG_SURPASS_MAX=1
 export BASEDIR="/raid/datasets/wiki"

-export BENCHMARK=10 DEBUG=2
+export BENCHMARK=10 BERT_LAYERS=2 DEBUG=2

 python3 examples/mlperf/model_train.py
@@ -1,10 +1,12 @@
 #!/bin/bash

-export PYTHONPATH="."
+export PYTHONPATH="." NV=1
 export MODEL="bert"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=24
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96

-export BEAM=4 BEAM_UOPS_MAX=4000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
+
+export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 export IGNORE_JIT_FIRST_BEAM=1
 export BASEDIR="/raid/datasets/wiki"

@@ -1,11 +1,14 @@
 #!/bin/bash
+set -e  # Exit on any error

-export PYTHONPATH="."
+export PYTHONPATH="." NV=1
 export MODEL="bert"
 export SUBMISSION_PLATFORM="tinybox_green"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=24
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96

-export BEAM=4 BEAM_UOPS_MAX=4000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
+
+export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 export IGNORE_JIT_FIRST_BEAM=1
 export BASEDIR="/raid/datasets/wiki"

@@ -17,7 +20,7 @@ DATETIME=$(date "+%m%d%H%M")
 LOGFILE="bert_green_${DATETIME}_${SEED}.log"

 # init
-BENCHMARK=10 INITMLPERF=1 BEAM_LOG_SURPASS_MAX=1 python3 examples/mlperf/model_train.py | tee $LOGFILE
+BENCHMARK=10 INITMLPERF=1 BERT_LAYERS=2 python3 examples/mlperf/model_train.py | tee $LOGFILE

 # run
 PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
@@ -4,24 +4,20 @@ This problem uses BERT for NLP.

 ## Requirements

-Install tinygrad and mlperf-logging from master.
+Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0.
 ```
 git clone https://github.com/tinygrad/tinygrad.git
 python3 -m pip install -e ".[mlperf]"
 ```
-Also install tqdm and tensorflow.
+Also install gdown (for dataset), numpy, tqdm and tensorflow.
 ```
-pip install tqdm tensorflow
+pip install gdown numpy tqdm tensorflow
 ```

 ### tinybox_green
 Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
 This is the default on production tinybox green.

-### tinybox_red
-Disable cwsr + increase mes timeout.
-Install the custom amdgpu driver per [README](https://github.com/nimlgen/amdgpu_ubuntu_22_04/blob/v6.1.3/readme.md)
-
 # 2. Directions

 ## Steps to download and verify data
@@ -56,18 +52,18 @@ BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-eval

 #### Steps to run benchmark
 ```
-examples/mlperf/training_submission_v4.1/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh
 ```

 ### tinybox_red

-#### One time setup
-
+#### Steps to run benchmark
 ```
-examples/mlperf/training_submission_v4.1/tinycorp/benchmarks/bert/implementations/tinybox_red/setup.sh
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh
 ```
+### tinybox_8xMI300X

 #### Steps to run benchmark
 ```
-examples/mlperf/training_submission_v4.1/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/run_and_time.sh
 ```
@@ -1,15 +1,17 @@
 #!/bin/bash

-export PYTHONPATH="."
+export PYTHONPATH="." AMD=1
 export MODEL="bert"
 export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96

-export BEAM=3 BEAM_UOPS_MAX=4000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
+
+export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 export IGNORE_JIT_FIRST_BEAM=1
 export BEAM_LOG_SURPASS_MAX=1
 export BASEDIR="/raid/datasets/wiki"

 export RESET_STEP=1
-export BENCHMARK=10 DEBUG=2
+export BENCHMARK=10 BERT_LAYERS=2 DEBUG=2

 python3 examples/mlperf/model_train.py
@@ -1,10 +1,12 @@
 #!/bin/bash

-export PYTHONPATH="."
+export PYTHONPATH="." AMD=1
 export MODEL="bert"
 export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96

-export BEAM=3 BEAM_UOPS_MAX=4000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
+
+export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 export IGNORE_JIT_FIRST_BEAM=1
 export BASEDIR="/raid/datasets/wiki"

@@ -1,11 +1,14 @@
 #!/bin/bash
+set -e  # Exit on any error

-export PYTHONPATH="."
+export PYTHONPATH="." AMD=1
 export MODEL="bert"
 export SUBMISSION_PLATFORM="tinybox_red"
 export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96

-export BEAM=3 BEAM_UOPS_MAX=4000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
+
+export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 export IGNORE_JIT_FIRST_BEAM=1
 export BASEDIR="/raid/datasets/wiki"

@@ -16,12 +19,13 @@ export SEED=$RANDOM
 DATETIME=$(date "+%m%d%H%M")
 LOGFILE="bert_red_${DATETIME}_${SEED}.log"

+export HCQDEV_WAIT_TIMEOUT_MS=100000  # prevents hang?
+
 # init
-sudo rmmod amdgpu || true
-BENCHMARK=10 INITMLPERF=1 BEAM_LOG_SURPASS_MAX=1 python3 examples/mlperf/model_train.py | tee $LOGFILE
+sleep 5 && sudo rmmod amdgpu || true
+BENCHMARK=10 INITMLPERF=1 BERT_LAYERS=2 python3 examples/mlperf/model_train.py | tee $LOGFILE

 # run
-# TODO: AMD driver hangs during init, but is 5% faster per step in real run.
+# TODO: AM driver resulted in nan
 sudo modprobe amdgpu
 PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
-sudo rmmod amdgpu || true
@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="."
+export PYTHONPATH="." NV=1
 export MODEL="resnet"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192

@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="."
+export PYTHONPATH="." NV=1
 export MODEL="resnet"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192

@@ -1,6 +1,7 @@
 #!/bin/bash
+set -e  # Exit on any error

-export PYTHONPATH="."
+export PYTHONPATH="." NV=1
 export MODEL="resnet"
 export SUBMISSION_PLATFORM="tinybox_green"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
@@ -10,7 +11,7 @@ export RESET_STEP=0
 export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=10 BEAM_PADTO=0

 # pip install -e ".[mlperf]"
-export LOGMLPERF=1
+export LOGMLPERF=${LOGMLPERF:-1}

 export SEED=$RANDOM
 DATETIME=$(date "+%m%d%H%M")
@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="."
+export PYTHONPATH="." AMD=1
 export MODEL="resnet"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192

@@ -8,6 +8,6 @@ export RESET_STEP=0

 export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=96 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0

-export BENCHMARK=10 DEBUG=2
+export BENCHMARK=10 DEBUG=${DEBUG:-2}

 python3 examples/mlperf/model_train.py
@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="."
+export PYTHONPATH="." AMD=1
 export MODEL="resnet"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192

@@ -1,6 +1,7 @@
 #!/bin/bash
+set -e  # Exit on any error

-export PYTHONPATH="."
+export PYTHONPATH="." AMD=1
 export MODEL="resnet"
 export SUBMISSION_PLATFORM="tinybox_red"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
@@ -10,13 +11,14 @@ export RESET_STEP=0
 export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=96 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0

 # pip install -e ".[mlperf]"
-export LOGMLPERF=1
+export LOGMLPERF=${LOGMLPERF:-1}

 export SEED=$RANDOM
 DATETIME=$(date "+%m%d%H%M")
 LOGFILE="resnet_red_${DATETIME}_${SEED}.log"

 # init
+sleep 5 && sudo rmmod amdgpu || true
 BENCHMARK=10 INITMLPERF=1 python3 examples/mlperf/model_train.py | tee $LOGFILE

 # run
@@ -0,0 +1,38 @@
+# 1. Problem
+
+This problem uses RetinaNet for SSD.
+
+## Requirements
+
+Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0.
+```
+git clone https://github.com/tinygrad/tinygrad.git
+python3 -m pip install -e ".[mlperf]"
+```
+
+Also install the following dependencies:
+```
+pip install tqdm numpy pycocotools boto3 pandas torch torchvision
+```
+
+### tinybox_green
+Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
+This is the default on production tinybox green.
+
+# 2. Directions
+
+## Steps to download data
+
+Run the following:
+```
+BASEDIR=/raid/datasets/openimages python3 extra/datasets/openimages.py
+```
+
+## Running
+
+### tinybox_green
+
+#### Steps to run benchmark
+```
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/retinanet/implementations/tinybox_green/run_and_time.sh
+```
@@ -0,0 +1,14 @@
+#!/bin/bash
+
+export PYTHONPATH="." NV=1
+export MODEL="retinanet"
+export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
+export BASEDIR="/raid/datasets/openimages"
+
+# export RESET_STEP=0
+
+export TRAIN_BEAM=2 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
+
+export BENCHMARK=5 DEBUG=2
+
+python examples/mlperf/model_train.py
@@ -0,0 +1,15 @@
+#!/bin/bash
+
+export PYTHONPATH="." NV=1
+export MODEL="retinanet"
+export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
+export BASEDIR="/raid/datasets/openimages"
+
+# export RESET_STEP=0
+
+export TRAIN_BEAM=2 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
+
+export WANDB=1 PARALLEL=0
+export RUNMLPERF=1
+
+python examples/mlperf/model_train.py
@@ -0,0 +1,24 @@
+#!/bin/bash
+set -e  # Exit on any error
+
+export PYTHONPATH="." NV=1
+export MODEL="retinanet"
+export SUBMISSION_PLATFORM="tinybox_green"
+export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
+
+export TRAIN_BEAM=2 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
+export IGNORE_JIT_FIRST_BEAM=1
+export BASEDIR="/raid/datasets/openimages"
+
+# pip install -e ".[mlperf]"
+export LOGMLPERF=1
+
+export SEED=$RANDOM
+DATETIME=$(date "+%m%d%H%M")
+LOGFILE="retinanet_green_${DATETIME}_${SEED}.log"
+
+# init
+BENCHMARK=10 INITMLPERF=1 python3 examples/mlperf/model_train.py | tee $LOGFILE
+
+# run
+PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
@@ -0,0 +1,14 @@
+#!/bin/bash
+
+export PYTHONPATH="." AMD=1
+export MODEL="retinanet"
+export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
+export BASEDIR="/raid/datasets/openimages"
+
+# export RESET_STEP=0
+
+export TRAIN_BEAM=2 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
+
+export BENCHMARK=5 DEBUG=2
+
+python examples/mlperf/model_train.py
@@ -0,0 +1,15 @@
+#!/bin/bash
+
+export PYTHONPATH="." AMD=1
+export MODEL="retinanet"
+export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
+export BASEDIR="/raid/datasets/openimages"
+
+# export RESET_STEP=0
+
+export TRAIN_BEAM=2 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
+
+export WANDB=1 PARALLEL=0
+export RUNMLPERF=1
+
+python examples/mlperf/model_train.py
@@ -0,0 +1,38 @@
+{
+    "submitter": "tinycorp",
+    "division": "closed",
+    "status": "Available on-premise",
+    "system_name": "tinybox 8xMI300X",
+    "number_of_nodes": "1",
+    "host_processors_per_node": "2",
+    "host_processor_model_name": "AMD EPYC 9354",
+    "host_processor_core_count": "32",
+    "host_processor_vcpu_count": "64",
+    "host_processor_frequency": "",
+    "host_processor_caches": "",
+    "host_processor_interconnect": "",
+    "host_memory_capacity": "2304GB",
+    "host_storage_type": "NVMe SSD",
+    "host_storage_capacity": "3x 4TB raid array",
+    "host_networking": "",
+    "host_networking_topology": "",
+    "host_memory_configuration": "24x 96GB DDR5",
+    "accelerators_per_node": "8",
+    "accelerator_model_name": "AMD Instinct MI300X 192GB HBM3",
+    "accelerator_host_interconnect": "PCIe 5.0 x16",
+    "accelerator_frequency": "",
+    "accelerator_on-chip_memories": "",
+    "accelerator_memory_configuration": "HBM3",
+    "accelerator_memory_capacity": "192GB",
+    "accelerator_interconnect": "",
+    "accelerator_interconnect_topology": "",
+    "cooling": "air",
+    "hw_notes": "",
+    "framework": "tinygrad, branch mlperf_training_v5.0",
+    "other_software_stack": {
+        "python": "3.10.16",
+        "ROCm": "3.0.0+94441cb"
+    },
+    "operating_system": "Ubuntu 24.04.1 LTS",
+    "sw_notes": ""
+  }
@@ -5,7 +5,7 @@
  "system_name": "tinybox green",
  "number_of_nodes": "1",
  "host_processors_per_node": "1",
-  "host_processor_model_name": "AMD EPYC 7532 32-Core Processor",
+  "host_processor_model_name": "AMD EPYC 7532",
  "host_processor_core_count": "32",
  "host_processor_vcpu_count": "64",
  "host_processor_frequency": "",
@@ -28,11 +28,11 @@
  "accelerator_interconnect_topology": "",
  "cooling": "air",
  "hw_notes": "",
-  "framework": "tinygrad, commit b5546912e24e0a864b35924da4efa5d71cfe368b",
+  "framework": "tinygrad, branch mlperf_training_v5.0",
  "other_software_stack": {
    "python": "3.10.12",
    "CUDA": "12.4"
  },
  "operating_system": "Ubuntu 22.04.4",
  "sw_notes": ""
-}
+}
@@ -5,7 +5,7 @@
  "system_name": "tinybox red",
  "number_of_nodes": "1",
  "host_processors_per_node": "1",
-  "host_processor_model_name": "AMD EPYC 7532 32-Core Processor",
+  "host_processor_model_name": "AMD EPYC 7532",
  "host_processor_core_count": "32",
  "host_processor_vcpu_count": "64",
  "host_processor_frequency": "",
@@ -28,11 +28,10 @@
  "accelerator_interconnect_topology": "",
  "cooling": "air",
  "hw_notes": "",
-  "framework": "tinygrad, commit b5546912e24e0a864b35924da4efa5d71cfe368b",
+  "framework": "tinygrad, branch mlperf_training_v5.0",
  "other_software_stack": {
-    "python": "3.10.12",
-    "ROCm": "6.1.3"
+    "python": "3.10.12"
  },
  "operating_system": "Ubuntu 22.04.4",
  "sw_notes": ""
-}
+}
@@ -0,0 +1,15 @@
+#!/bin/bash
+
+export PYTHONPATH="." AMD=1
+export MODEL="bert"
+export DEFAULT_FLOAT="HALF" GPUS=1 BS=128 EVAL_BS=128
+
+export BEAM=3 BEAM_UOPS_MAX=4000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export IGNORE_JIT_FIRST_BEAM=1
+# export BEAM_LOG_SURPASS_MAX=1
+# export BASEDIR="/raid/datasets/wiki"
+
+export RESET_STEP=1
+export BENCHMARK=10 BERT_LAYERS=2 DEBUG=2
+
+python3 examples/mlperf/model_train.py
@@ -0,0 +1,69 @@
+# 1. Problem
+
+This problem uses BERT for NLP.
+
+## Requirements
+
+Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0.
+```
+git clone https://github.com/tinygrad/tinygrad.git
+python3 -m pip install -e ".[mlperf]"
+```
+Also install gdown (for dataset), numpy, tqdm and tensorflow.
+```
+pip install gdown numpy tqdm tensorflow
+```
+
+### tinybox_green
+Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
+This is the default on production tinybox green.
+
+# 2. Directions
+
+## Steps to download and verify data
+
+### 1. Download raw data
+
+```
+BASEDIR="/raid/datasets/wiki" WIKI_TRAIN=1 VERIFY_CHECKSUM=1 python3 extra/datasets/wikipedia_download.py
+```
+
+### 2. Preprocess train and validation data
+
+Note: The number of threads used for preprocessing is limited by available memory. With 128GB of RAM, a maximum of 16 threads is recommended. 
+
+#### Training:
+```
+BASEDIR="/raid/datasets/wiki" NUM_WORKERS=16 python3 extra/datasets/wikipedia.py pre-train all
+```
+
+Generating a specific topic (Between 0 and 499)
+```
+BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-train 42
+```
+
+#### Validation:
+```
+BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-eval
+```
+## Running
+
+### tinybox_green
+
+#### Steps to run benchmark
+```
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh
+```
+
+### tinybox_red
+
+#### Steps to run benchmark
+```
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh
+```
+### tinybox_8xMI300X
+
+#### Steps to run benchmark
+```
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/run_and_time.sh
+```
@@ -0,0 +1,14 @@
+#!/bin/bash
+
+export PYTHONPATH="." AMD=1
+export MODEL="bert"
+export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
+export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
+
+export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
+export BASEDIR="/raid/datasets/wiki"
+
+export BENCHMARK=10 BERT_LAYERS=2 DEBUG=2
+
+python3 examples/mlperf/model_train.py
@@ -0,0 +1,17 @@
+#!/bin/bash
+
+export PYTHONPATH="." AMD=1
+export MODEL="bert"
+export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
+
+# similar to https://github.com/mlcommons/training_results_v3.1/blob/d06288b2bd675a9d88e0e6181f5bb5626b71ec19/Quanta_Cloud_Technology/results/D54U-3U/bert/result_1.txt#L54
+export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
+export TRAIN_STEPS=3900
+
+export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
+export BASEDIR="/raid/datasets/wiki"
+
+export WANDB=1 PARALLEL=0
+
+RUNMLPERF=1 python3 examples/mlperf/model_train.py
@@ -0,0 +1,29 @@
+#!/bin/bash
+set -e  # Exit on any error
+set -o pipefail  # Make pipeline fail if any command fails
+
+export PYTHONPATH="." AMD=1
+export MODEL="bert"
+export SUBMISSION_PLATFORM="tinybox_8xMI300X"
+export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
+
+# similar to https://github.com/mlcommons/training_results_v3.1/blob/d06288b2bd675a9d88e0e6181f5bb5626b71ec19/Quanta_Cloud_Technology/results/D54U-3U/bert/result_1.txt#L54
+export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
+export TRAIN_STEPS=3900
+
+export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
+export BASEDIR="/raid/datasets/wiki"
+
+# pip install -e ".[mlperf]"
+export LOGMLPERF=1
+
+export SEED=$RANDOM
+DATETIME=$(date "+%m%d%H%M")
+LOGFILE="bert_8xMI300x_${DATETIME}_${SEED}.log"
+
+# init  # TODO: without DEBUG=2 it hangs
+BENCHMARK=10 INITMLPERF=1 BERT_LAYERS=2 DEBUG=2 python3 examples/mlperf/model_train.py | tee $LOGFILE
+
+# run
+PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
@@ -0,0 +1,69 @@
+# 1. Problem
+
+This problem uses BERT for NLP.
+
+## Requirements
+
+Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0.
+```
+git clone https://github.com/tinygrad/tinygrad.git
+python3 -m pip install -e ".[mlperf]"
+```
+Also install gdown (for dataset), numpy, tqdm and tensorflow.
+```
+pip install gdown numpy tqdm tensorflow
+```
+
+### tinybox_green
+Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
+This is the default on production tinybox green.
+
+# 2. Directions
+
+## Steps to download and verify data
+
+### 1. Download raw data
+
+```
+BASEDIR="/raid/datasets/wiki" WIKI_TRAIN=1 VERIFY_CHECKSUM=1 python3 extra/datasets/wikipedia_download.py
+```
+
+### 2. Preprocess train and validation data
+
+Note: The number of threads used for preprocessing is limited by available memory. With 128GB of RAM, a maximum of 16 threads is recommended. 
+
+#### Training:
+```
+BASEDIR="/raid/datasets/wiki" NUM_WORKERS=16 python3 extra/datasets/wikipedia.py pre-train all
+```
+
+Generating a specific topic (Between 0 and 499)
+```
+BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-train 42
+```
+
+#### Validation:
+```
+BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-eval
+```
+## Running
+
+### tinybox_green
+
+#### Steps to run benchmark
+```
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh
+```
+
+### tinybox_red
+
+#### Steps to run benchmark
+```
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh
+```
+### tinybox_8xMI300X
+
+#### Steps to run benchmark
+```
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/run_and_time.sh
+```
@@ -0,0 +1,16 @@
+#!/bin/bash
+
+export PYTHONPATH="." NV=1
+export MODEL="bert"
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
+
+export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
+
+export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export IGNORE_JIT_FIRST_BEAM=1
+export BEAM_LOG_SURPASS_MAX=1
+export BASEDIR="/raid/datasets/wiki"
+
+export BENCHMARK=10 BERT_LAYERS=2 DEBUG=2
+
+python3 examples/mlperf/model_train.py
@@ -0,0 +1,15 @@
+#!/bin/bash
+
+export PYTHONPATH="." NV=1
+export MODEL="bert"
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
+
+export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
+
+export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export IGNORE_JIT_FIRST_BEAM=1
+export BASEDIR="/raid/datasets/wiki"
+
+export WANDB=1 PARALLEL=0
+
+RUNMLPERF=1 python3 examples/mlperf/model_train.py
@@ -0,0 +1,27 @@
+#!/bin/bash
+set -e  # Exit on any error
+set -o pipefail  # Make pipeline fail if any command fails
+
+export PYTHONPATH="." NV=1
+export MODEL="bert"
+export SUBMISSION_PLATFORM="tinybox_green"
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
+
+export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
+
+export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export IGNORE_JIT_FIRST_BEAM=1
+export BASEDIR="/raid/datasets/wiki"
+
+# pip install -e ".[mlperf]"
+export LOGMLPERF=1
+
+export SEED=$RANDOM
+DATETIME=$(date "+%m%d%H%M")
+LOGFILE="bert_green_${DATETIME}_${SEED}.log"
+
+# init
+BENCHMARK=10 INITMLPERF=1 BERT_LAYERS=2 python3 examples/mlperf/model_train.py | tee $LOGFILE
+
+# run
+PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
@@ -0,0 +1,69 @@
+# 1. Problem
+
+This problem uses BERT for NLP.
+
+## Requirements
+
+Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0.
+```
+git clone https://github.com/tinygrad/tinygrad.git
+python3 -m pip install -e ".[mlperf]"
+```
+Also install gdown (for dataset), numpy, tqdm and tensorflow.
+```
+pip install gdown numpy tqdm tensorflow
+```
+
+### tinybox_green
+Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
+This is the default on production tinybox green.
+
+# 2. Directions
+
+## Steps to download and verify data
+
+### 1. Download raw data
+
+```
+BASEDIR="/raid/datasets/wiki" WIKI_TRAIN=1 VERIFY_CHECKSUM=1 python3 extra/datasets/wikipedia_download.py
+```
+
+### 2. Preprocess train and validation data
+
+Note: The number of threads used for preprocessing is limited by available memory. With 128GB of RAM, a maximum of 16 threads is recommended. 
+
+#### Training:
+```
+BASEDIR="/raid/datasets/wiki" NUM_WORKERS=16 python3 extra/datasets/wikipedia.py pre-train all
+```
+
+Generating a specific topic (Between 0 and 499)
+```
+BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-train 42
+```
+
+#### Validation:
+```
+BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-eval
+```
+## Running
+
+### tinybox_green
+
+#### Steps to run benchmark
+```
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh
+```
+
+### tinybox_red
+
+#### Steps to run benchmark
+```
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh
+```
+### tinybox_8xMI300X
+
+#### Steps to run benchmark
+```
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/run_and_time.sh
+```
@@ -0,0 +1,17 @@
+#!/bin/bash
+
+export PYTHONPATH="." AMD=1
+export MODEL="bert"
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
+
+export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
+
+export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export IGNORE_JIT_FIRST_BEAM=1
+export BEAM_LOG_SURPASS_MAX=1
+export BASEDIR="/raid/datasets/wiki"
+
+export RESET_STEP=1
+export BENCHMARK=10 BERT_LAYERS=2 DEBUG=2
+
+python3 examples/mlperf/model_train.py
@@ -0,0 +1,15 @@
+#!/bin/bash
+
+export PYTHONPATH="." AMD=1
+export MODEL="bert"
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
+
+export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
+
+export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export IGNORE_JIT_FIRST_BEAM=1
+export BASEDIR="/raid/datasets/wiki"
+
+export WANDB=1 PARALLEL=0
+
+RUNMLPERF=1 python3 examples/mlperf/model_train.py
@@ -0,0 +1,32 @@
+#!/bin/bash
+set -e  # Exit on any error
+set -o pipefail  # Make pipeline fail if any command fails
+
+export PYTHONPATH="." AMD=1
+export MODEL="bert"
+export SUBMISSION_PLATFORM="tinybox_red"
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
+
+export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
+
+export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export IGNORE_JIT_FIRST_BEAM=1
+export BASEDIR="/raid/datasets/wiki"
+
+# pip install -e ".[mlperf]"
+export LOGMLPERF=1
+
+export SEED=$RANDOM
+DATETIME=$(date "+%m%d%H%M")
+LOGFILE="bert_red_${DATETIME}_${SEED}.log"
+
+export HCQDEV_WAIT_TIMEOUT_MS=100000  # prevents hang?
+
+# init
+sleep 5 && sudo rmmod amdgpu || true
+BENCHMARK=10 INITMLPERF=1 BERT_LAYERS=2 python3 examples/mlperf/model_train.py | tee $LOGFILE
+
+# run
+# TODO: AM driver resulted in nan
+sudo modprobe amdgpu
+PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
@@ -0,0 +1,50 @@
+# 1. Problem
+
+This problem uses the ResNet-50 CNN to do image classification.
+
+## Requirements
+
+Install tinygrad and mlperf-logging from master.
+```
+git clone https://github.com/tinygrad/tinygrad.git
+python3 -m pip install -e ".[mlperf]"
+```
+
+### tinybox_green
+Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
+This is the default on production tinybox green.
+
+### tinybox_red
+Disable cwsr
+This is the default on production tinybox red.
+```
+sudo vi /etc/modprobe.d/amdgpu.conf
+cat <<EOF > /etc/modprobe.d/amdgpu.conf
+options amdgpu cwsr_enable=0
+EOF
+sudo update-initramfs -u
+sudo reboot
+
+# validate
+sudo cat /sys/module/amdgpu/parameters/cwsr_enable #= 0
+```
+
+# 2. Directions
+
+## Steps to download and verify data
+
+```
+IMGNET_TRAIN=1 python3 extra/datasets/imagenet_download.py
+```
+
+## Steps for one time setup
+
+### tinybox_red
+```
+examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/setup.sh
+```
+
+## Steps to run benchmark
+```
+examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/run_and_time.sh
+```
@@ -0,0 +1,13 @@
+#!/bin/bash
+
+export PYTHONPATH="." NV=1
+export MODEL="resnet"
+export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
+
+export RESET_STEP=0
+
+export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=10 BEAM_PADTO=0
+
+export BENCHMARK=10 DEBUG=2
+
+python3 examples/mlperf/model_train.py
@@ -0,0 +1,15 @@
+#!/bin/bash
+
+export PYTHONPATH="." NV=1
+export MODEL="resnet"
+export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
+
+export RESET_STEP=0
+
+export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=10 BEAM_PADTO=0
+
+export EVAL_START_EPOCH=3 EVAL_FREQ=4
+
+export WANDB=1 PARALLEL=0
+
+python3 examples/mlperf/model_train.py
@@ -0,0 +1,25 @@
+#!/bin/bash
+set -e  # Exit on any error
+set -o pipefail  # Make pipeline fail if any command fails
+
+export PYTHONPATH="." NV=1
+export MODEL="resnet"
+export SUBMISSION_PLATFORM="tinybox_green"
+export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
+
+export RESET_STEP=0
+
+export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=10 BEAM_PADTO=0
+
+# pip install -e ".[mlperf]"
+export LOGMLPERF=${LOGMLPERF:-1}
+
+export SEED=$RANDOM
+DATETIME=$(date "+%m%d%H%M")
+LOGFILE="resnet_green_${DATETIME}_${SEED}.log"
+
+# init
+BENCHMARK=10 INITMLPERF=1 python3 examples/mlperf/model_train.py | tee $LOGFILE
+
+# run
+PARALLEL=0 RUNMLPERF=1 EVAL_START_EPOCH=3 EVAL_FREQ=4 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
@@ -0,0 +1,50 @@
+# 1. Problem
+
+This problem uses the ResNet-50 CNN to do image classification.
+
+## Requirements
+
+Install tinygrad and mlperf-logging from master.
+```
+git clone https://github.com/tinygrad/tinygrad.git
+python3 -m pip install -e ".[mlperf]"
+```
+
+### tinybox_green
+Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
+This is the default on production tinybox green.
+
+### tinybox_red
+Disable cwsr
+This is the default on production tinybox red.
+```
+sudo vi /etc/modprobe.d/amdgpu.conf
+cat <<EOF > /etc/modprobe.d/amdgpu.conf
+options amdgpu cwsr_enable=0
+EOF
+sudo update-initramfs -u
+sudo reboot
+
+# validate
+sudo cat /sys/module/amdgpu/parameters/cwsr_enable #= 0
+```
+
+# 2. Directions
+
+## Steps to download and verify data
+
+```
+IMGNET_TRAIN=1 python3 extra/datasets/imagenet_download.py
+```
+
+## Steps for one time setup
+
+### tinybox_red
+```
+examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/setup.sh
+```
+
+## Steps to run benchmark
+```
+examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/run_and_time.sh
+```
@@ -0,0 +1,13 @@
+#!/bin/bash
+
+export PYTHONPATH="." AMD=1
+export MODEL="resnet"
+export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
+
+export RESET_STEP=0
+
+export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=96 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
+
+export BENCHMARK=10 DEBUG=${DEBUG:-2}
+
+python3 examples/mlperf/model_train.py
@@ -0,0 +1,15 @@
+#!/bin/bash
+
+export PYTHONPATH="." AMD=1
+export MODEL="resnet"
+export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
+
+export RESET_STEP=0
+
+export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=96 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
+
+export EVAL_START_EPOCH=3 EVAL_FREQ=4
+
+export WANDB=1 PARALLEL=0
+
+python3 examples/mlperf/model_train.py
@@ -0,0 +1,26 @@
+#!/bin/bash
+set -e  # Exit on any error
+set -o pipefail  # Make pipeline fail if any command fails
+
+export PYTHONPATH="." AMD=1
+export MODEL="resnet"
+export SUBMISSION_PLATFORM="tinybox_red"
+export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
+
+export RESET_STEP=0
+
+export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=96 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
+
+# pip install -e ".[mlperf]"
+export LOGMLPERF=${LOGMLPERF:-1}
+
+export SEED=$RANDOM
+DATETIME=$(date "+%m%d%H%M")
+LOGFILE="resnet_red_${DATETIME}_${SEED}.log"
+
+# init
+sleep 5 && sudo rmmod amdgpu || true
+BENCHMARK=10 INITMLPERF=1 python3 examples/mlperf/model_train.py | tee $LOGFILE
+
+# run
+PARALLEL=0 RUNMLPERF=1 EVAL_START_EPOCH=3 EVAL_FREQ=4 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
@@ -5,4 +5,4 @@ rocm-smi --setmclk 3
 rocm-smi --setperflevel high

 # power cap to 350W
-# echo "350000000" | sudo tee /sys/class/drm/card{1..6}/device/hwmon/hwmon*/power1_cap
+echo "350000000" | sudo tee /sys/class/drm/card{1..6}/device/hwmon/hwmon*/power1_cap
@@ -0,0 +1,38 @@
+# 1. Problem
+
+This problem uses RetinaNet for SSD.
+
+## Requirements
+
+Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0.
+```
+git clone https://github.com/tinygrad/tinygrad.git
+python3 -m pip install -e ".[mlperf]"
+```
+
+Also install the following dependencies:
+```
+pip install tqdm numpy pycocotools boto3 pandas torch torchvision
+```
+
+### tinybox_green
+Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
+This is the default on production tinybox green.
+
+# 2. Directions
+
+## Steps to download data
+
+Run the following:
+```
+BASEDIR=/raid/datasets/openimages python3 extra/datasets/openimages.py
+```
+
+## Running
+
+### tinybox_green
+
+#### Steps to run benchmark
+```
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/retinanet/implementations/tinybox_green/run_and_time.sh
+```
@@ -0,0 +1,14 @@
+#!/bin/bash
+
+export PYTHONPATH="." NV=1
+export MODEL="retinanet"
+export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
+export BASEDIR="/raid/datasets/openimages"
+
+# export RESET_STEP=0
+
+export TRAIN_BEAM=2 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
+
+export BENCHMARK=5 DEBUG=2
+
+python examples/mlperf/model_train.py
@@ -0,0 +1,15 @@
+#!/bin/bash
+
+export PYTHONPATH="." NV=1
+export MODEL="retinanet"
+export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
+export BASEDIR="/raid/datasets/openimages"
+
+# export RESET_STEP=0
+
+export TRAIN_BEAM=2 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
+
+export WANDB=1 PARALLEL=0
+export RUNMLPERF=1
+
+python examples/mlperf/model_train.py
@@ -0,0 +1,25 @@
+#!/bin/bash
+set -e  # Exit on any error
+set -o pipefail  # Make pipeline fail if any command fails
+
+export PYTHONPATH="." NV=1
+export MODEL="retinanet"
+export SUBMISSION_PLATFORM="tinybox_green"
+export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
+
+export TRAIN_BEAM=2 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
+export IGNORE_JIT_FIRST_BEAM=1
+export BASEDIR="/raid/datasets/openimages"
+
+# pip install -e ".[mlperf]"
+export LOGMLPERF=1
+
+export SEED=$RANDOM
+DATETIME=$(date "+%m%d%H%M")
+LOGFILE="retinanet_green_${DATETIME}_${SEED}.log"
+
+# init
+BENCHMARK=10 INITMLPERF=1 python3 examples/mlperf/model_train.py | tee $LOGFILE
+
+# run
+PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
--- a/Show More
+++ b/Show More