Progress with pack/unpack wgsl built-ins

f16 comp/decomp WIP
2026-08-20 19:26:07 +00:00 · 2024-12-05 16:12:21 +01:00 · 2024-12-05 14:14:23 +01:00
584 changed files with 18241 additions and 1706004 deletions
@@ -1,15 +0,0 @@
-name: Run process replay tests
-description: Verify process replay compared to master
-runs:
-  using: "composite"
-  steps:
-    - name: Run process replay tests
-      shell: bash
-      run: |
-        export PR_TITLE=$(jq -r .pull_request.title "$GITHUB_EVENT_PATH")
-        export CURRENT_SHA=${{ github.event.pull_request && github.event.pull_request.head.sha || github.sha }}
-        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
-        git checkout $CURRENT_HEAD  # restore to branch
@@ -1,224 +0,0 @@
-name: Setup Python & Install
-description: Sets up Python and installs project dependencies.
-inputs:
-  python-version:
-    description: 'Python version to use'
-    required: false
-    default: '3.12'
-  key:
-    description: 'Key for the python cache'
-    required: false
-    default: ''  # if you don't set a key, it doesn't cache
-  deps:
-    description: 'Extra dependency groups (comma separated)'
-    required: false
-    default: ''
-  pydeps:
-    description: 'Extra Python dependency groups (space separated)'
-    required: false
-    default: ''
-  opencl:
-    description: "Install OpenCL?"
-    required: false
-    default: 'false'
-  amd:
-    description: "Install AMD?"
-    required: false
-    default: 'false'
-  cuda:
-    description: "Install CUDA?"
-    required: false
-    default: 'false'
-  webgpu:
-    description: "Install webgpu?"
-    required: false
-    default: 'false'
-  llvm:
-    description: "Install LLVM?"
-    required: false
-    default: 'false'
-runs:
-  using: "composite"
-  steps:
-    - name: Set up Python ${{ inputs.python-version }}
-      uses: actions/setup-python@v5
-      with:
-        python-version: ${{ inputs.python-version }}
-    - name: Upgrade pip
-      shell: bash
-      run: python -m pip install --upgrade pip
-
-    # **** 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'
-      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') }}
-
-    # **** Caching downloads ****
-
-    - name: Cache downloads (Linux)
-      if: inputs.key != '' && runner.os == 'Linux'
-      uses: actions/cache@v4
-      with:
-        path: ~/.cache/tinygrad/downloads/
-        key: downloads-cache-${{ inputs.key }}-${{ env.DOWNLOAD_CACHE_VERSION }}
-    - name: Cache downloads (macOS)
-      if: inputs.key != '' && runner.os == 'macOS'
-      uses: actions/cache@v4
-      with:
-        path: ~/Library/Caches/tinygrad/downloads/
-        key: osx-downloads-cache-${{ inputs.key }}-${{ env.DOWNLOAD_CACHE_VERSION }}
-
-    # **** 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'
-      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
-
-    # **** AMD ****
-
-    - name: Install AMD (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
-        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
-        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)
-      if: inputs.amd == 'true' && runner.os == 'macOS'
-      shell: bash
-      run: |
-        sudo mkdir -p /usr/local/lib
-        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
-        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 --quiet cmake ninja llvm@15 zlib glew flex bison boost zstd ncurses
-    - name: Cache gpuocelot
-      if: inputs.cuda == 'true'
-      id: cache-build
-      uses: actions/cache@v4
-      env:
-        cache-name: cache-gpuocelot-build
-      with:
-        path: ${{ github.workspace }}/gpuocelot/ocelot
-        key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-0
-    - name: Clone/compile gpuocelot
-      if: inputs.cuda == 'true' && steps.cache-build.outputs.cache-hit != 'true'
-      shell: bash
-      run: |
-        git clone --recurse-submodules https://github.com/gpuocelot/gpuocelot.git ${{ github.workspace }}/gpuocelot
-        cd ${{ github.workspace }}/gpuocelot/ocelot
-        git checkout b16039dc940dc6bc4ea0a98380495769ff35ed99
-        mkdir build
-        cd build
-        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'
-      shell: bash
-      run: |
-        cd ${{ github.workspace }}/gpuocelot/ocelot/build
-        sudo cp libgpuocelot.${{ runner.os == 'macOS' && 'dylib' || 'so' }} /usr/${{ runner.os == 'macOS' && 'local/' || ''}}lib/
-
-    # **** WebGPU ****
-
-    - name: Install WebGPU dawn (Linux)
-      if: inputs.webgpu == 'true' && runner.os == 'Linux'
-      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
-
-    - name: Install WebGPU dawn (macOS)
-      if: inputs.webgpu == 'true' && runner.os == 'macOS'
-      shell: bash
-      run: |
-        brew tap wpmed92/dawn
-        brew install dawn
-
-    # **** 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@19
@@ -2,7 +2,7 @@ name: Benchmarks
 env:
  # TODO: this rescheduling makes gpt2, mixtral and llama unjitted slower
  # TODO: very slow for llama 70B and resnet training 6 GPU
-  CAPTURE_PROCESS_REPLAY: "1"
+  RUN_PROCESS_REPLAY: "1"
  ASSERT_PROCESS_REPLAY: "0"
  PYTHONPATH: .
  GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
@@ -24,14 +24,11 @@ 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 -e -o pipefail {0}
+        shell: bash -o pipefail {0}
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -52,76 +49,57 @@ jobs:
    - name: reset process replay
      run: python3.11 test/external/process_replay/reset.py
    - name: Run Stable Diffusion
-      run: BENCHMARK_LOG=stable_diffusion JIT=1 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
+      run: JIT=1 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
    - name: Run Stable Diffusion without fp16
-      run: BENCHMARK_LOG=stable_diffusion_fp32 JIT=1 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd_no_fp16.txt
+      run: JIT=1 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd_no_fp16.txt
    - name: Run Stable Diffusion v2
-      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
+      run: JIT=1 python3.11 examples/sdv2.py --fp16 --seed 0 --noshow --timing | tee sdv2.txt
    - name: Run SDXL
-      run: BENCHMARK_LOG=stable_diffusion_xl CAPTURE_PROCESS_REPLAY=0 JIT=1 python3.11 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
+      run: 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: 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_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_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
+      run: METAL=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded
    - name: Run Tensor Core GEMM (float)
-      run: DEBUG=2 SHOULD_USE_TC=1 python3.11 extra/gemm/simple_matmul.py | tee matmul.txt
+      run: DEBUG=2 python3.11 extra/gemm/simple_matmul.py | tee matmul.txt
    - name: Run Tensor Core GEMM (half)
-      run: DEBUG=2 SHOULD_USE_TC=1 HALF=1 python3.11 extra/gemm/simple_matmul.py | tee matmul_half.txt
+      run: DEBUG=2 HALF=1 python3.11 extra/gemm/simple_matmul.py | tee matmul_half.txt
    - name: Run Tensor Core GEMM (bfloat16)
-      run: DEBUG=2 SHOULD_USE_TC=1 BFLOAT16=1 python3.11 extra/gemm/simple_matmul.py | tee matmul_bfloat16.txt
+      run: DEBUG=2 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: |
-        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
+        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
    - name: Run LLaMA with BEAM
-      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
+      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
    - name: Run quantized LLaMA
      run: |
-        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
+        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
+    - 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: |
-        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
+        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
    - name: Run GPT2 w HALF
-      run: BENCHMARK_LOG=gpt2_half HALF=1 python3.11 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
+      run: HALF=1 python3.11 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
    - name: Run GPT2 w HALF/BEAM
-      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
+      run: HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 CAST_BEFORE_VIEW=0 python3.11 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half_beam.txt
    - 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: BENCHMARK_LOG=cifar_10steps JIT=1 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar.txt
+      run: JIT=1 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar.txt
    - name: Run 10 CIFAR training steps w HALF
-      run: BENCHMARK_LOG=cifar_10steps_half JIT=2 STEPS=10 DEFAULT_FLOAT=HALF python3.11 examples/hlb_cifar10.py | tee train_cifar_half.txt
+      run: 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: 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
+      run: JIT=1 WINO=1 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar_wino.txt
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (Mac)
@@ -133,8 +111,6 @@ 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
@@ -158,10 +134,10 @@ jobs:
  testnvidiabenchmark:
    name: tinybox green Benchmark
    runs-on: [self-hosted, Linux, tinyboxgreen]
-    timeout-minutes: 30
+    timeout-minutes: 20
    defaults:
      run:
-        shell: bash -e -o pipefail {0}
+        shell: bash -o pipefail {0}
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -185,62 +161,59 @@ jobs:
    - name: reset process replay
      run: test/external/process_replay/reset.py
    - name: Run model inference benchmark
-      run: NV=1 CAPTURE_PROCESS_REPLAY=0 NOCLANG=1 python3 test/external/external_model_benchmark.py
+      run: NV=1 RUN_PROCESS_REPLAY=0 NOCLANG=1 python3 test/external/external_model_benchmark.py
    - name: Test speed vs torch
-      run: NV=1 CAPTURE_PROCESS_REPLAY=0 HALF=1 BIG=2 TORCHCUDA=1 python3 test/test_speed_v_torch.py | tee torch_speed.txt
+      run: NV=1 RUN_PROCESS_REPLAY=0 HALF=1 BIG=2 TORCHCUDA=1 python3 test/test_speed_v_torch.py | tee torch_speed.txt
    - name: Test speed vs theoretical
      run: NV=1 IGNORE_BEAM_CACHE=1 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
-    - name: Test benchmark allreduce
-      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_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
+        NV=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded
+        PTX=1 NV=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded
    - name: Run Tensor Core GEMM (CUDA)
      run: |
-        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
+        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
    - name: Run Tensor Core GEMM (PTX)
-      run: NV=1 PTX=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_ptx.txt
+      run: NV=1 PTX=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_ptx.txt
    - name: Run Tensor Core GEMM (NV)
-      run: NV=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_nv.txt
+      run: NV=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: BENCHMARK_LOG=stable_diffusion NV=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
+      run: NV=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
    - name: Run SDXL
-      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
+      run: NV=1 RUN_PROCESS_REPLAY=0 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
    - name: Run LLaMA
      run: |
-        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
+        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
    - name: Run LLaMA with BEAM
-      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
+      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
    # - 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
+    #   run: NV=1 RUN_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
+    #   run: NV=1 RUN_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: 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: 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
+      run: NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_beam.txt
+    - name: Run LLaMA-3 8B on 4 GPUs
+      run: NV=1 RUN_PROCESS_REPLAY=0 python3 examples/llama3.py --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 RUN_PROCESS_REPLAY=0 python3 examples/llama3.py --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 RUN_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 BENCHMARK_LOG=mixtral NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/mixtral.py --temperature 0 --count 10 --timing | tee mixtral.txt
+      run: time NV=1 RUN_PROCESS_REPLAY=0 python3 examples/mixtral.py --temperature 0 --count 10 --timing | tee mixtral.txt
    - name: Run GPT2
      run: |
-        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
+        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
    - name: Run GPT2 w HALF
-      run: BENCHMARK_LOG=gpt2_half NV=1 HALF=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
+      run: NV=1 HALF=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
    - name: Run GPT2 w HALF/BEAM
-      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
+      run: NV=1 HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 CAST_BEFORE_VIEW=0 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half_beam.txt
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (NVIDIA)
@@ -249,7 +222,6 @@ jobs:
          torch_speed.txt
          matmul.txt
          matmul_bfloat16.txt
-          matmul_tf32.txt
          matmul_ptx.txt
          matmul_nv.txt
          sd.txt
@@ -275,7 +247,7 @@ jobs:
    timeout-minutes: 20
    defaults:
      run:
-        shell: bash -e -o pipefail {0}
+        shell: bash -o pipefail {0}
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -304,26 +276,23 @@ 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: BENCHMARK_LOG=cifar_10steps NV=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
+      run: NV=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
    - name: Run 10 CIFAR training steps w HALF
-      run: BENCHMARK_LOG=cifar_10steps_half NV=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
+      run: 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: BENCHMARK_LOG=cifar_10steps_bf16 NV=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
+      run: 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: 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
+      run: NV=1 RUN_PROCESS_REPLAY=0 WINO=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
    - name: Run full CIFAR training w 1 GPU
-      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
+      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
    - name: Run full CIFAR training steps w 6 GPUS
-      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
+      run: time RUN_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 BENCHMARK_LOG=resnet_eval NV=1 MODEL=resnet python3 examples/mlperf/model_eval.py
+      run: time NV=1 MODEL=resnet python3 examples/mlperf/model_eval.py
    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
-      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
+      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
    - name: Run 10 MLPerf ResNet50 training steps (6 gpu)
-      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: 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
+      run: NV=1 RUN_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (NVIDIA Training)
@@ -334,10 +303,9 @@ jobs:
          train_cifar_bf16.txt
          train_cifar_wino.txt
          train_cifar_one_gpu.txt
-          train_cifar_six_gpu.txt
          train_resnet.txt
          train_resnet_one_gpu.txt
-          train_bert.txt
+          train_cifar_six_gpu.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

@@ -347,13 +315,11 @@ jobs:
    timeout-minutes: 20
    defaults:
      run:
-        shell: bash -e -o pipefail {0}
+        shell: bash -o pipefail {0}
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
-    - name: Insert amdgpu
-      run: sudo modprobe amdgpu
    - name: Symlink models and datasets
      run: |
        mkdir -p weights
@@ -389,12 +355,9 @@ 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_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
+      run: AMD=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded
    - name: Run Tensor Core GEMM (AMD)
-      run: AMD=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py | tee matmul_amd.txt
+      run: AMD=1 HALF=1 DEBUG=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
@@ -402,46 +365,38 @@ jobs:
    # 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 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: BENCHMARK_LOG=stable_diffusion AMD=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
+      run: AMD=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
    - name: Run SDXL
-      run: BENCHMARK_LOG=stable_diffusion_xl CAPTURE_PROCESS_REPLAY=0 AMD=1 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
+      run: AMD=1 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
    - name: Run LLaMA 7B
      run: |
-        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
+        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
    - name: Run LLaMA 7B with BEAM
-      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
+      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
    # - 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
+    #   run: AMD=1 RUN_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
+    #   run: AMD=1 RUN_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: 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: 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: 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
+      run: AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_beam.txt
+    - name: Run LLaMA-3 8B on 4 GPUs
+      run: AMD=1 RUN_PROCESS_REPLAY=0 python3 examples/llama3.py --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 RUN_PROCESS_REPLAY=0 python3 examples/llama3.py --shard 6 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_six_gpu.txt
+    - name: Run LLaMA-2 70B
+      run: AMD=1 RUN_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 BENCHMARK_LOG=mixtral AMD=1 python3 examples/mixtral.py --temperature 0 --count 10 --timing | tee mixtral.txt
+      run: time AMD=1 python3 examples/mixtral.py --temperature 0 --count 10 --timing | tee mixtral.txt
    - name: Run GPT2
      run: |
-        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
+        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
    - name: Run GPT2 w HALF
-      run: BENCHMARK_LOG=gpt2_half AMD=1 HALF=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
+      run: AMD=1 HALF=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
    - name: Run GPT2 w HALF/BEAM
-      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
+      run: AMD=1 HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 CAST_BEFORE_VIEW=0 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half_beam.txt
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (AMD)
@@ -470,16 +425,14 @@ jobs:
  testmoreamdbenchmark:
    name: tinybox red Training Benchmark
    runs-on: [self-hosted, Linux, tinybox]
-    timeout-minutes: 30
+    timeout-minutes: 20
    defaults:
      run:
-        shell: bash -e -o pipefail {0}
+        shell: bash -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: Symlink models and datasets
      run: |
        mkdir -p weights
@@ -497,29 +450,30 @@ 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: 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: BENCHMARK_LOG=cifar_10steps AMD=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
+      run: AMD=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
    - name: Run 10 CIFAR training steps w HALF
-      run: BENCHMARK_LOG=cifar_10steps_half AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
+      run: 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: BENCHMARK_LOG=cifar_10steps_bf16 AMD=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
+      run: 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: 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
+      run: AMD=1 WINO=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
    - name: Run full CIFAR training w 1 GPU
-      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
+      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
    - name: Run full CIFAR training steps w 6 GPUS
-      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
+      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 BENCHMARK_LOG=resnet_eval AMD=1 MODEL=resnet python3 examples/mlperf/model_eval.py
+      run: time 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
+      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: 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
+      run: AMD=1 RUN_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (AMD Training)
@@ -530,10 +484,9 @@ jobs:
          train_cifar_bf16.txt
          train_cifar_wino.txt
          train_cifar_one_gpu.txt
-          train_cifar_six_gpu.txt
          train_resnet.txt
          train_resnet_one_gpu.txt
-          train_bert.txt
+          train_cifar_six_gpu.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

@@ -543,7 +496,7 @@ jobs:
    timeout-minutes: 20
    defaults:
      run:
-        shell: bash -e -o pipefail {0}
+        shell: bash -o pipefail {0}
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -555,16 +508,20 @@ 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: openpilot compile 0.9.4
+      run: PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python examples/openpilot/compile2.py | tee openpilot_compile_0_9_4.txt
+    - name: openpilot compile 0.9.7
+      run: PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python examples/openpilot/compile2.py https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx | tee openpilot_compile_0_9_7.txt
    - 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: 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
+      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
    - name: benchmark openpilot 0.9.7
-      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
+      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
    - name: benchmark openpilot w IMAGE=2 0.9.4
-      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
+      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
    - name: benchmark openpilot w IMAGE=2 0.9.7
-      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
+      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
    - 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
@@ -1,21 +0,0 @@
-name: Run MLPerf Training
-
-on:
-  # schedule:
-  #   - cron: '0 */1 * * *'
-  workflow_dispatch:
-
-jobs:
-  run_script_job:
-    runs-on: [self-hosted, Linux, tinyboxgreen]
-
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v4
-    - name: Symlink datasets
-      run: |
-        mkdir -p extra/datasets
-        ln -s /raid/datasets/imagenet extra/datasets/imagenet
-    - name: Run resnet
-      run: |
-        LOGMLPERF=0 examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/resnet/implementations/tinybox_green/run_and_time.sh
@@ -10,7 +10,6 @@ notebooks
 *.so
 *.txt
 build
-!examples/tinychat/assets/cdn.jsdelivr.net/npm/[email protected]/build/
 /dist
 *.egg-info
 /env
@@ -34,8 +33,6 @@ extra/datasets/open-images-v6-mlperf
 extra/datasets/kits/
 extra/datasets/COCO/
 extra/datasets/audio*
-extra/huggingface_onnx/models/*
-extra/huggingface_onnx/*.yaml
 extra/weights
 venv
 examples/**/net.*[js,json]
@@ -58,7 +55,4 @@ weights
 comgr_*
 *.pkl
 site/
-profile_stats
-*.log
-target
-.mypy_cache
+master_schedule.py
@@ -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

 # Add files or directories matching the regex patterns to the blacklist. The
 # regex matches against base names, not paths.
@@ -81,7 +81,7 @@ See [examples/beautiful_mnist.py](examples/beautiful_mnist.py) for the full vers
 tinygrad already supports numerous accelerators, including:

 - [x] [GPU (OpenCL)](tinygrad/runtime/ops_gpu.py)
- [x] [CPU (C Code)](tinygrad/runtime/ops_cpu.py)
+- [x] [CLANG (C Code)](tinygrad/runtime/ops_clang.py)
 - [x] [LLVM](tinygrad/runtime/ops_llvm.py)
 - [x] [METAL](tinygrad/runtime/ops_metal.py)
 - [x] [CUDA](tinygrad/runtime/ops_cuda.py)
@@ -151,7 +151,7 @@ We'll start with what will get your PR closed with a pointer to this section:

 - No code golf! While low line count is a guiding light of this project, anything that remotely looks like code golf will be closed. The true goal is reducing complexity and increasing readability, and deleting `\n`s does nothing to help with that.
 - All docs and whitespace changes will be closed unless you are a well-known contributor. The people writing the docs should be those who know the codebase the absolute best. People who have not demonstrated that shouldn't be messing with docs. Whitespace changes are both useless *and* carry a risk of introducing bugs.
- Anything you claim is a "speedup" must be benchmarked. In general, the goal is simplicity, so even if your PR makes things marginally faster, you have to consider the tradeoff with maintainability and readability.
+- Anything you claim is a "speedup" must be benchmarked. In general, the goal is simplicity, so even if your PR makes things marginally faster, you have to consider the tradeoff with maintainablity and readablity.
 - In general, the code outside the core `tinygrad/` folder is not well tested, so unless the current code there is broken, you shouldn't be changing it.
 - If your PR looks "complex", is a big diff, or adds lots of lines, it won't be reviewed or merged. Consider breaking it up into smaller PRs that are individually clear wins. A common pattern I see is prerequisite refactors before adding new functionality. If you can (cleanly) refactor to the point that the feature is a 3 line change, this is great, and something easy for us to review.

@@ -35,7 +35,7 @@ def _try_dlopen_$name():
  for candidate in PATHS_TO_TRY:
    try: return ctypes.CDLL(candidate)
    except OSError: pass
-  return None
+  raise RuntimeError("library $name not found")
 EOF
 }

@@ -69,7 +69,7 @@ generate_comgr() {
  --clang-args="-D__HIP_PLATFORM_AMD__ -I/opt/rocm/include -x c++" -o $BASE/comgr.py -l /opt/rocm/lib/libamd_comgr.so
  fixup $BASE/comgr.py
  sed -i "s\import ctypes\import ctypes, ctypes.util, os\g" $BASE/comgr.py
-  patch_dlopen $BASE/comgr.py amd_comgr "'/opt/rocm/lib/libamd_comgr.so'" "os.getenv('ROCM_PATH', '')+'/lib/libamd_comgr.so'" "'/usr/local/lib/libamd_comgr.dylib'" "'/opt/homebrew/lib/libamd_comgr.dylib'"
+  patch_dlopen $BASE/comgr.py amd_comgr "'/opt/rocm/lib/libamd_comgr.so'" "os.getenv('ROCM_PATH', '')+'/lib/libamd_comgr.so'"
  sed -i "s\ctypes.CDLL('/opt/rocm/lib/libamd_comgr.so')\_try_dlopen_amd_comgr()\g" $BASE/comgr.py
  python3 -c "import tinygrad.runtime.autogen.comgr"
 }
@@ -78,11 +78,7 @@ 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/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\import ctypes\import ctypes, os\g" $BASE/kfd.py
  python3 -c "import tinygrad.runtime.autogen.kfd"
 }

@@ -171,7 +167,6 @@ generate_amd() {
    extra/hip_gpu_driver/sdma_v6_0_0_pkt_open.h \
    extra/hip_gpu_driver/gc_11_0_0_offset.h \
    extra/hip_gpu_driver/gc_10_3_0_offset.h \
-    extra/hip_gpu_driver/sienna_cichlid_ip_offset.h \
    --clang-args="-I/opt/rocm/include -x c++" \
    -o $BASE/amd_gpu.py

@@ -212,10 +207,8 @@ generate_libc() {
  clang2py -k cdefstum \
    $(dpkg -L libc6-dev | grep sys/mman.h) \
    $(dpkg -L libc6-dev | grep sys/syscall.h) \
-    /usr/include/string.h \
    /usr/include/elf.h \
    /usr/include/unistd.h \
-    /usr/include/asm-generic/mman-common.h \
    -o $BASE/libc.py

  sed -i "s\import ctypes\import ctypes, ctypes.util, os\g" $BASE/libc.py
@@ -225,30 +218,11 @@ generate_libc() {
  fixup $BASE/libc.py
 }

-generate_llvm() {
-  INC="$(llvm-config-14 --includedir)"
-  clang2py -k cdefstum \
-    $(find "$INC/llvm-c/" -type f -name '*.h' | sort) \
-    "$INC/llvm/Config/Targets.def" \
-    "$INC/llvm/Config/AsmPrinters.def" \
-    "$INC/llvm/Config/AsmParsers.def" \
-    "$INC/llvm/Config/Disassemblers.def" \
-    --clang-args="$(llvm-config-14 --cflags)" \
-    -o "$BASE/llvm.py"
-
-  sed -i "s\import ctypes\import ctypes, tinygrad.runtime.support.llvm as llvm_support\g" "$BASE/llvm.py"
-  sed -i "s\FIXME_STUB\llvm\g" "$BASE/llvm.py"
-  sed -i "s\FunctionFactoryStub()\ctypes.CDLL(llvm_support.LLVM_PATH)\g" "$BASE/llvm.py"
-
-  fixup "$BASE/llvm.py"
-}
-
 generate_kgsl() {
  clang2py extra/qcom_gpu_driver/msm_kgsl.h -o $BASE/kgsl.py -k cdefstum
  fixup $BASE/kgsl.py
  sed -i "s\import ctypes\import ctypes, os\g" $BASE/kgsl.py
  sed -nE 's/#define ([A-Za-z0-9_]+)_SHIFT\s*[^\S\r\n]*[0-9]*$/def \1(val): return (val << \1_SHIFT) \& \1_MASK/p' extra/qcom_gpu_driver/msm_kgsl.h >> $BASE/kgsl.py
-  sed -i "s\fcntl.ioctl(__fd, (__idir<<30)\__fd.ioctl((__idir<<30)\g" $BASE/kgsl.py
  python3 -c "import tinygrad.runtime.autogen.kgsl"
 }

@@ -273,274 +247,6 @@ generate_qcom() {
  python3 -c "import tinygrad.runtime.autogen.qcom_dsp"
 }

-generate_pci() {
-  clang2py -k cdefstum \
-    /usr/include/linux/pci_regs.h \
-    -o $BASE/pci.py
-  fixup $BASE/pci.py
-}
-
-generate_vfio() {
-  clang2py -k cdefstum \
-    /usr/include/linux/vfio.h \
-    -o $BASE/vfio.py
-  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 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 \
-    extra/amdpci/headers/amdgpu_vm.h \
-    extra/amdpci/headers/discovery.h \
-    extra/amdpci/headers/amdgpu_ucode.h \
-    extra/amdpci/headers/psp_gfx_if.h \
-    extra/amdpci/headers/amdgpu_psp.h \
-    extra/amdpci/headers/amdgpu_irq.h \
-    extra/amdpci/headers/amdgpu_doorbell.h \
-    $AMKERN_INC/soc15_ih_clientid.h \
-    --clang-args="-include stdint.h" \
-    -o $BASE/am/am.py
-  fixup $BASE/am/am.py
-  sed -i "s\(int64_t)\ \g" $BASE/am/am.py
-  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 \
-    $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 \
-    $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 \
-    $AMKERN_INC/vega10_enum.h \
-    -o $BASE/am/vega10.py
-  fixup $BASE/am/vega10.py
-
-  clang2py -k cdefstum \
-    $AMKERN_INC/navi10_enum.h \
-    -o $BASE/am/navi10.py
-  fixup $BASE/am/navi10.py
-
-  clang2py -k cdefstum \
-    $AMKERN_INC/soc21_enum.h \
-    -o $BASE/am/soc21.py
-  fixup $BASE/am/soc21.py
-
-  clang2py -k cdefstum \
-    $AMKERN_INC/soc24_enum.h \
-    -o $BASE/am/soc24.py
-  fixup $BASE/am/soc24.py
-
-  clang2py -k cdefstum \
-    $AMKERN_INC/asic_reg/mp/mp_13_0_0_offset.h \
-    $AMKERN_INC/asic_reg/mp/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 \
-    $AMKERN_INC/asic_reg/mp/mp_14_0_2_offset.h \
-    $AMKERN_INC/asic_reg/mp/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 \
-    $AMKERN_INC/asic_reg/mp/mp_11_0_offset.h \
-    $AMKERN_INC/asic_reg/mp/mp_11_0_sh_mask.h \
-    -o $BASE/am/mp_11_0.py
-  fixup $BASE/am/mp_11_0.py
-
-  clang2py -k cdefstum \
-    $AMKERN_INC/asic_reg/gc/gc_9_4_3_offset.h \
-    $AMKERN_INC/asic_reg/gc/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 \
-    $AMKERN_INC/asic_reg/gc/gc_10_3_0_offset.h \
-    $AMKERN_INC/asic_reg/gc/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 \
-    $AMKERN_INC/asic_reg/gc/gc_11_0_0_offset.h \
-    $AMKERN_INC/asic_reg/gc/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 \
-    $AMKERN_INC/asic_reg/gc/gc_12_0_0_offset.h \
-    $AMKERN_INC/asic_reg/gc/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 \
-    $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 \
-    $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 \
-    $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 \
-    $AMKERN_INC/asic_reg/mmhub/mmhub_3_0_0_offset.h \
-    $AMKERN_INC/asic_reg/mmhub/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 \
-    $AMKERN_INC/asic_reg/mmhub/mmhub_3_0_2_offset.h \
-    $AMKERN_INC/asic_reg/mmhub/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 \
-    $AMKERN_INC/asic_reg/nbio/nbio_2_3_offset.h \
-    $AMKERN_INC/asic_reg/nbio/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 \
-    $AMKERN_INC/asic_reg/nbio/nbio_7_2_0_offset.h \
-    $AMKERN_INC/asic_reg/nbio/nbio_7_2_0_sh_mask.h \
-    -o $BASE/am/nbio_7_2_0.py
-  fixup $BASE/am/nbio_7_2_0.py
-
-  clang2py -k cdefstum \
-    $AMKERN_INC/asic_reg/mmhub/mmhub_4_1_0_offset.h \
-    $AMKERN_INC/asic_reg/mmhub/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 \
-    $AMKERN_INC/asic_reg/nbio/nbio_4_3_0_offset.h \
-    $AMKERN_INC/asic_reg/nbio/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 \
-    $AMKERN_INC/asic_reg/nbif/nbif_6_3_1_offset.h \
-    $AMKERN_INC/asic_reg/nbif/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 \
-    $AMKERN_INC/asic_reg/nbio/nbio_7_9_0_offset.h \
-    $AMKERN_INC/asic_reg/nbio/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 \
-    $AMKERN_INC/asic_reg/nbio/nbio_7_11_0_offset.h \
-    $AMKERN_INC/asic_reg/nbio/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 \
-    $AMKERN_INC/asic_reg/oss/osssys_6_0_0_offset.h \
-    $AMKERN_INC/asic_reg/oss/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 \
-    $AMKERN_INC/asic_reg/oss/osssys_7_0_0_offset.h \
-    $AMKERN_INC/asic_reg/oss/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 \
-    $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 \
-    $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 \
-    $AMKERN_INC/asic_reg/hdp/hdp_6_0_0_offset.h \
-    $AMKERN_INC/asic_reg/hdp/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 \
-    $AMKERN_INC/asic_reg/hdp/hdp_7_0_0_offset.h \
-    $AMKERN_INC/asic_reg/hdp/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() {
-  clang2py -k cdefstum \
-    extra/sqtt/sqtt.h \
-    -o $BASE/sqtt.py
-
-  fixup $BASE/sqtt.py
-  sed -i "s\import ctypes\import ctypes, os\g" $BASE/sqtt.py
-  python3 -c "import tinygrad.runtime.autogen.sqtt"
-}
-
-generate_webgpu() {
-  clang2py extra/webgpu/webgpu.h -o $BASE/webgpu.py
-  fixup $BASE/webgpu.py
-  sed -i "s/FIXME_STUB/webgpu/g" "$BASE/webgpu.py"
-  sed -i "s/FunctionFactoryStub()/ctypes.CDLL(webgpu_support.WEBGPU_PATH)/g" "$BASE/webgpu.py"
-  sed -i "s/import ctypes/import ctypes, tinygrad.runtime.support.webgpu as webgpu_support/g" "$BASE/webgpu.py"
-  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
@@ -550,18 +256,11 @@ elif [ "$1" == "hsa" ]; then generate_hsa
 elif [ "$1" == "kfd" ]; then generate_kfd
 elif [ "$1" == "nv" ]; then generate_nv
 elif [ "$1" == "amd" ]; then generate_amd
-elif [ "$1" == "am" ]; then generate_am
-elif [ "$1" == "sqtt" ]; then generate_sqtt
 elif [ "$1" == "qcom" ]; then generate_qcom
 elif [ "$1" == "io_uring" ]; then generate_io_uring
 elif [ "$1" == "libc" ]; then generate_libc
-elif [ "$1" == "llvm" ]; then generate_llvm
 elif [ "$1" == "kgsl" ]; then generate_kgsl
 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
+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
 else echo "usage: $0 <type>"
 fi
@@ -1,13 +1,13 @@
 # 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 "UOp" that fuses the compute into kernels, using memory only when needed
+# 3. a "LazyBuffer" that fuses the compute into kernels, using memory only when needed
 # 4. a "Tensor" that provides an easy to use frontend with autograd ".backward()"


 print("******** first, the runtime ***********")

-from tinygrad.runtime.ops_cpu import ClangJITCompiler, MallocAllocator, CPUProgram
+from tinygrad.runtime.ops_clang import ClangProgram, ClangCompiler, MallocAllocator

 # allocate some buffers
 out = MallocAllocator.alloc(4)
@@ -19,10 +19,10 @@ MallocAllocator._copyin(a, memoryview(bytearray([2,0,0,0])))
 MallocAllocator._copyin(b, memoryview(bytearray([3,0,0,0])))

 # compile a program to a binary
-lib = ClangJITCompiler().compile("void add(int *out, int *a, int *b) { out[0] = a[0] + b[0]; }")
+lib = ClangCompiler().compile("void add(int *out, int *a, int *b) { out[0] = a[0] + b[0]; }")

-# create a runtime for the program
-fxn = CPUProgram("add", lib)
+# create a runtime for the program (ctypes.CDLL)
+fxn = ClangProgram("add", lib)

 # run the program
 fxn(out, a, b)
@@ -34,12 +34,12 @@ assert val == 5

 print("******** second, the Device ***********")

-DEVICE = "CPU"   # NOTE: you can change this!
+DEVICE = "CLANG"   # NOTE: you can change this!

 import struct
 from tinygrad.dtype import dtypes
 from tinygrad.device import Buffer, Device
-from tinygrad.uop.ops import UOp, Ops
+from tinygrad.ops import UOp, Ops
 from tinygrad.shape.shapetracker import ShapeTracker

 # allocate some buffers + load in values
@@ -65,7 +65,7 @@ kernel = get_kernel(Device[DEVICE].renderer, s).linearize()
 # compile a program (and print the source)
 fxn = CompiledRunner(kernel.to_program())
 print(fxn.p.src)
-# NOTE: fxn.clprg is the CPUProgram
+# NOTE: fxn.clprg is the ClangProgram

 # run the program
 fxn.exec([out, a, b])
@@ -74,52 +74,36 @@ fxn.exec([out, a, b])
 assert out.as_buffer().cast('I')[0] == 5


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

+from tinygrad.engine.lazy import LazyBuffer
 from tinygrad.engine.realize import run_schedule
-from tinygrad.engine.schedule import create_schedule_with_vars
-from tinygrad.engine.grouper import get_becomes_map
+from tinygrad.engine.schedule import create_schedule

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

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

-# group the computation into kernels
-becomes_map = get_becomes_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, _, becomes_map = create_schedule_with_vars(s)
-assert len(sched) == 1
+# schedule the computation as a list of kernels
+sched = create_schedule([out])
+for si in sched: print(si.ast.op)  # NOTE: the first two convert it to CLANG

 # 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 = becomes_map[assign]
-assert out.op is Ops.BUFFER and not out.buffer.is_allocated()
-print(out)
-
 # run that schedule
 run_schedule(sched)

 # check the data out
-assert out.is_realized and out.buffer.as_buffer().cast('I')[0] == 5
+assert out.realized is not None and out.realized.as_buffer().cast('I')[0] == 5


 print("******** fourth, the Tensor ***********")
@@ -26,11 +26,10 @@ l1n, l2n = l1.numpy(), l2.numpy()
 from tinygrad.nn.optim import SGD
 optim = SGD([l1, l2])

-Tensor.training = True
 X, Y = X_train[(samples:=Tensor.randint(128, high=X_train.shape[0]))], Y_train[samples]
 optim.zero_grad()
 model(X).sparse_categorical_crossentropy(Y).backward()
-optim.schedule_step()   # this will step the optimizer without running realize
+optim._step()   # this will step the optimizer without running realize

 # *****
 # 3. Create a schedule.
@@ -48,7 +47,7 @@ for si in schedule: print(str(si)[:80])
 # 4. Lower a schedule.

 from tinygrad.engine.realize import lower_schedule_item, ExecItem
-lowered: List[ExecItem] = [lower_schedule_item(si) for si in tqdm(schedule)]
+lowered: List[ExecItem] = [ExecItem(lower_schedule_item(si).prg, list(si.bufs)) for si in tqdm(schedule)]

 # *****
 # 5. Run the schedule
@@ -1,39 +0,0 @@
-# AM Driver
-
-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`!
-
-Optional requirements:
-
-* System without IOMMU for P2P / SDMA support
-* vfio-pci module for IRQ handling
-
-## Environment Variables
-
-| Variable | Possible Value(s) | Description |
-|----------|------------------|-------------|
-| AM_RESET | [1] | Performs a full GPU reset (reloading all firmware and IP blocks) |
-| AM_DEBUG | [0-4] | Sets the level of additional debugging information |
-
-## AM Driver Details
-
-### Compute & SDMA Queues
-
-AM binds compute queues directly to MEC (bypassing MES). Tinygrad uses only one compute queue, which is bound at `pipe=0 queue=0`. Similarly, the single SDMA queue is bound at `engine=0 queue=0`.
-
-### Boot
-
-The GPU being passed can be in one of several states:
-1. Not initialized
-2. Initialized by amdgpu
-3. Initialized by AM
-
-The first and second states require a full GPU setup since their states are unknown. The second state also requires a mode1 reset to reinitialize all components.
-
-The third state can be set up partially to optimize boot time. In this case, only the GFX and SDMA IPs need to be initialized. To enable this, AM uses a separate boot memory that is guaranteed not to be overwritten. This physical memory is utilized for all blocks that are initialized only during the initial AM boot. To determine if the GPU is in the third state, AM uses `regSCRATCH_REG7` as a flag.
-
-### VM Management
-
-Each AM device sets up only a single `VMID=0` and one page directory. The page directory used is 3-level and thus supports up to 512GB of virtual addresses. All AM devices are located in one virtual address space.
@@ -7,17 +7,19 @@ The tinygrad framework has four pieces

 There is a good [bunch of tutorials](https://mesozoic-egg.github.io/tinygrad-notes/) by Di Zhu that go over tinygrad internals.

-There's also a [doc describing speed](../developer/speed.md)
-
 ## Frontend

-Everything in [Tensor](../tensor/index.md) is syntactic sugar around constructing a graph of [UOps](../developer/uop.md).
+Everything in [Tensor](../tensor/index.md) is syntactic sugar around [function.py](function.md), where the forwards and backwards passes are implemented for the different functions. There's about 25 of them, implemented using about 20 basic ops. Those basic ops go on to construct a graph of:

-The `UOp` graph specifies the compute in terms of low level tinygrad ops. Not all UOps will actually become realized. There's two types of UOps, base and view. base contains compute into a contiguous buffer, and view is a view (specified by a ShapeTracker). Inputs to a base can be either base or view, inputs to a view can only be a single base.
+::: tinygrad.engine.lazy.LazyBuffer
+    options:
+        show_source: false
+
+The `LazyBuffer` graph specifies the compute in terms of low level tinygrad ops. Not all LazyBuffers will actually become realized. There's two types of LazyBuffers, base and view. base contains compute into a contiguous buffer, and view is a view (specified by a ShapeTracker). Inputs to a base can be either base or view, inputs to a view can only be a single base.

 ## Scheduling

-The [scheduler](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/engine/schedule.py) converts the graph of UOps into a list of `ScheduleItem`. One `ScheduleItem` is one kernel on the GPU, and the scheduler is responsible for breaking the large compute graph into subgraphs that can fit in a kernel. `ast` specifies what compute to run, and `bufs` specifies what buffers to run it on.
+The [scheduler](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/engine/schedule.py) converts the graph of LazyBuffers into a list of `ScheduleItem`. One `ScheduleItem` is one kernel on the GPU, and the scheduler is responsible for breaking the large compute graph into subgraphs that can fit in a kernel. `ast` specifies what compute to run, and `bufs` specifies what buffers to run it on.

 ::: tinygrad.engine.schedule.ScheduleItem

@@ -0,0 +1,33 @@
+::: tinygrad.function
+    options:
+        members: [
+            "Contiguous",
+            "ContiguousBackward",
+            "Cast",
+            "Neg",
+            "Reciprocal",
+            "Sin",
+            "Relu",
+            "Log",
+            "Exp",
+            "Sqrt",
+            "Sigmoid",
+            "Sign",
+            "Less",
+            "Eq",
+            "Xor",
+            "Add",
+            "Sub",
+            "Mul",
+            "Div",
+            "Where",
+            "Sum",
+            "Max",
+            "Expand",
+            "Reshape",
+            "Permute",
+            "Pad",
+            "Shrink",
+            "Flip",
+        ]
+        show_source: false
@@ -115,8 +115,9 @@ HCQ-compatible devices use a global timeline signal for synchronizing all operat
 ```python
 HWQueue().wait(your_device.timeline_signal, your_device.timeline_value - 1) \
         .exec(...)
-         .signal(your_device.timeline_signal, your_device.next_timeline()) \
+         .signal(your_device.timeline_signal, your_device.timeline_value) \
         .submit(your_device)
+your_device.timeline_value += 1

 # Optionally wait for execution
 your_device.timeline_signal.wait(your_device.timeline_value - 1)
@@ -1,109 +0,0 @@
-# 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.lazydata)
-```
-
-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.lazydata.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`.
@@ -36,9 +36,9 @@ The `Allocator` class is responsible for managing memory on the device. There is

 ### Program

-The `Program` class is created for each loaded program. It is responsible for executing the program on the device. As an example, here is a `CPUProgram` implementation which loads program and runs it.
+The `Program` class is created for each loaded program. It is responsible for compiling and executing the program on the device. As an example, here is a `ClangProgram` implementation which loads program and runs it.

-::: tinygrad.runtime.ops_cpu.CPUProgram
+::: tinygrad.runtime.ops_clang.ClangProgram
    options:
        members: true

@@ -1,71 +0,0 @@
-# speed in tinygrad
-
-## Overview
-
-Speed refers to many different things. To break it down to four, there's:
-
- Compile Speed (Python)
- Execution Speed (driver)
- Model Speed (scheduler)
- Kernel Speed (codegen)
-
-## Compile Speed (Python)
-
-This is how long the first run of your model takes. It's limited largely by the runtime of the Python doing UOp rewrites. Currently it's a bit slow, but on par with torch.compile. It gets even slower if you are using BEAM, since that's compiling many variants of each kernel.
-
-This will be improved by writing faster graph_rewrite, doing less graph_rewrite, and better parallelization.
-
-## Execution Speed (driver)
-
-After your model is compiled, you are often using the `TinyJIT`. tinygrad has the best execution speed of any framework because it usually bypasses the GPU driver and prebuilds the command queue. It's tons faster than normal CUDA, and often even faster than CUDA Graph.
-
-There's very little to improve here, as this is almost never the bottleneck.
-
-## Model Speed (scheduler)
-
-The scheduler determines how operations are grouped into kernels and which Tensors are written to memory. This is currently a big bottleneck of training speed.
-
-The decisions are often not obvious. For example, when is it worth recomputing an arithmetic operation instead of storing and loading from memory? Example:
-
-```python
-from tinygrad import Tensor
-a = Tensor.rand(100)
-b = Tensor.rand(100)
-c = Tensor.rand(100)
-d = Tensor.rand(100)
-out1 = a+b+c
-out2 = a+b+d
-Tensor.realize(out1, out2)
-```
-
-The real answer is obvious, compute both `out1` and `out2` in the same kernel. But you can't always do that. If you can't, should `a+b` first be saved to a subbuffer? Or should both the `out1` and `out2` kernels recompute `a+b`?
-
-In this case: with recompute (6 reads + 2 writes), no recompute (6 reads + 3 writes), so we should probably recompute. However, once you add movement ops and casts this is even harder to figure out. tinygrad doesn't yet have a systematic way to do it.
-
-## Kernel Speed (codegen)
-
-Given that you have decided how the model ops will be grouped and what will be written to memory, kernel speed determines how fast that operation is done. This is what BEAM changes, it searches over a set of equivalent kernels which all perform the same operation and finds the one which performs the task the fastest.
-
-In `kernel.py` we have a set of `OptOps`, these control the parameters of the speed optimizations applied to the kernel.
-
-### Memory
-
-The main bottleneck in most kernels is accessing memory. In a freshman algorithms class, you'll learn about cache aware matrix multiplication, and this is all forms of that. While the same math is run, the order in which you run it can have large impacts on the speed depending on if the data you are loading. OptOps will change this order.
-
-Memory, even cache, is often much slower than accessing the register file. The amount of times data is used in math is called the "arithmetic intensity". For operations like BS=1 GEMV, the arithmetic intensity is 1, but for GEMMs and convs it can be much higher. OptOps like UPCAST and UNROLL can increase this, but be careful of making them too large, as if there's too much register pressure on the GPU the warp scheduler may not be able to fit many warps, or even worse, it could be spilling to local memory.
-
-4090s have 1 TB/s of ram bandwidth and ~160 TFLOPS of compute, so you need to use each loaded value ~100 times. The L1 cache has around 40 TB/s of bandwidth, so in order to get full compute utilization you need to use each value ~4 times.
-
-A lot of work can still be done here. For example, we never copy the inputs to on chip SRAM, but this is often quite helpful for kernel speed. Also, we aren't doing a good job with L2 cache awareness (the locals handle L1 quite well)
-
-### Tensor Cores
-
-Many accelerators have Tensor Cores / MAC arrays / systolic arrays. The main value of these is that, since they are 2-D, they create an n^2 ratio between the compute and the input data.
-
-GPUs use Tensor Cores instead of MAC arrays to fit better in the GPU warp paradigm. This is because the output of Tensor Cores is O(n) wrt the input, while the output of MAC arrays like the AMX is O(n^2)
-
-We have a simple framework in tinygrad for adding these ALU blocks and achieving good performance from them.
-
-### 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 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.uop.ops.UOp
+::: tinygrad.ops.UOp
    options:
        members: false
        members_order: source
        show_labels: false

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

 Variable | Possible Value(s) | Description
 ---|---|---
-DEBUG               | [1-7]      | enable debugging output (operations, timings, speed, generated code and more)
-GPU                 | [1]        | enable the GPU (OpenCL) backend
+DEBUG               | [1-6]      | enable debugging output, with 4 you get operations, timings, speed, generated code and more
+GPU                 | [1]        | enable the GPU 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
+CLANG               | [1]        | enable Clang backend
 LLVM                | [1]        | enable LLVM backend
 BEAM                | [#]        | number of beams in kernel beam search
 DEFAULT_FLOAT       | [HALF, ...]| specify the default float dtype (FLOAT32, HALF, BFLOAT16, FLOAT64, ...), default to FLOAT32
 IMAGE               | [1-2]      | enable 2d specific optimizations
 FLOAT16             | [1]        | use float16 for images instead of float32
 PTX                 | [1]        | enable the specialized [PTX](https://docs.nvidia.com/cuda/parallel-thread-execution/) assembler for Nvidia GPUs. If not set, defaults to generic CUDA codegen backend.
-PROFILE             | [1]        | enable profiling. This feature is supported in NV, AMD, QCOM and METAL backends.
+PROFILE             | [1]        | enable output of [perfetto](https://ui.perfetto.dev/) compatible profile. This feature is supported in NV and AMD backends.
 VISIBLE_DEVICES     | [list[int]]| restricts the NV/AMD devices that are available. The format is a comma-separated list of identifiers (indexing starts with 0).
 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...)
-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
+VIZ                 | [1]        | 0=disabled, 1=[viz enabled](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/viz)
@@ -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_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.
+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.

 ### tinygrad is lazy

@@ -17,7 +17,7 @@ from tinygrad import Device
 print(Device.DEFAULT)
 ```

-You will see `CUDA` here on a GPU instance, or `CPU` here on a CPU instance.
+You will see `CUDA` here on a GPU instance, or `CLANG` here on a CPU instance.

 ## A simple model

@@ -29,12 +29,4 @@
 ::: tinygrad.nn.state.get_state_dict
 ::: tinygrad.nn.state.get_parameters
 ::: tinygrad.nn.state.load_state_dict
-::: tinygrad.nn.state.tar_extract
-    options:
-        show_signature: false
-        separate_signature: false
 ::: tinygrad.nn.state.torch_load
-    options:
-        show_signature: false
-        separate_signature: false
-::: tinygrad.nn.state.gguf_load
@@ -110,7 +110,7 @@ class TinyNet:

  def __call__(self, x):
    x = self.l1(x)
-    x = x.leaky_relu()
+    x = x.leakyrelu()
    x = self.l2(x)
    return x

@@ -118,7 +118,7 @@ net = TinyNet()
 ```

 We can see that the forward pass of our neural network is just the sequence of operations performed on the input tensor `x`.
-We can also see that functional operations like `leaky_relu` are not defined as classes and instead are just methods we can just call.
+We can also see that functional operations like `leakyrelu` are not defined as classes and instead are just methods we can just call.
 Finally, we just initialize an instance of our neural network, and we are ready to start training it.

 ## Training
@@ -1,74 +1,14 @@
 # Runtimes

-tinygrad supports various runtimes, enabling your code to scale across a wide range of devices. The default runtime can be automatically selected based on the available hardware, or you can force a specific runtime to be default using environment variables (e.g., `CPU=1`).
+tinygrad supports various runtimes, enabling your code to scale across a wide range of devices. The default runtime can be automatically selected based on the available hardware, or you can force a specific runtime to be default using environment variables (e.g., `CLANG=1`).

 | 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/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. |
+| [AMD](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_amd.py) | Provides acceleration for AMD GPUs | RDNA2/RDNA3 series GPUs |
 | [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 |
 | [GPU (OpenCL)](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_gpu.py) | Accelerates computations using OpenCL on GPUs | OpenCL 2.0 compatible device |
-| [CPU (C Code)](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cpu.py) | Runs on CPU using the clang compiler | `clang` compiler in system `PATH` |
-| [LLVM (LLVM IR)](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_llvm.py) | Runs on CPU using the LLVM compiler infrastructure | llvm libraries installed and findable |
-| [WEBGPU](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_webgpu.py) | Runs on GPU using the Dawn WebGPU engine (used in Google Chrome) | Dawn library installed and findable. Download binaries [here](https://github.com/wpmed92/pydawn/releases/tag/v0.1.6). |
-
-## Interoperability
-
-tinygrad provides interoperability with OpenCL and PyTorch, allowing efficient tensor data sharing between frameworks through the `Tensor.from_blob` API. This enables zero-copy operations by working directly with external memory pointers.
-
-**Important**: When using external memory pointers with tinygrad tensors, you must ensure these pointers remain valid throughout the entire lifetime of the tinygrad tensor to prevent memory corruption.
-
-### `CUDA`/`METAL` PyTorch Interoperability
-
-You can seamlessly work with CUDA/MPS tensors between PyTorch and tinygrad without data copying:
-```python
-from tinygrad.dtype import _from_torch_dtype
-tensor1 = torch.tensor([1.0, 2.0, 3.0], device=torch.device("cuda"))
-tiny_tensor1 = Tensor.from_blob(tensor1.data_ptr(), tensor1.shape, dtype=_from_torch_dtype(tensor1.dtype), device='CUDA')
-
-# Before tinygrad calculations, mps needs to be synchronized to make sure data is valid.
-if data.device.type == "mps": torch.mps.synchronize()
-else: torch.cuda.synchronize()
-
-x = (tiny_tensor1 + 1).realize()
-```
-
-### `QCOM` OpenCL Interoperability
-
-tinygrad supports OpenCL interoperability on `QCOM` backend.
-
-Buffer interop allows direct access to OpenCL memory buffers:
-```python
-# create raw opencl buffer.
-cl_buf = cl.clCreateBuffer(cl_context, cl.CL_MEM_READ_WRITE, 0x100, None, status := ctypes.c_int32())
-
-# extract pointers
-cl_buf_desc_ptr = to_mv(ctypes.addressof(cl_buf), 8).cast('Q')[0]
-rawbuf_ptr = to_mv(cl_buf_desc_ptr, 0x100).cast('Q')[20] # offset 0xA0 is a raw gpu pointer.
-
-# create tiny tensor
-tiny = Tensor.from_blob(rawbuf_ptr, (8, 8), dtype=dtypes.int, device='QCOM')
-```
-
-And the same for the images:
-```python
-# create cl image.
-cl_img = cl.clCreateImage2D(cl_context, cl.CL_MEM_READ_WRITE, cl.cl_image_format(cl.CL_RGBA, cl.CL_FLOAT), w, h, 0, None, status := ctypes.c_int32())
-
-# extract pointers
-cl_buf_desc_ptr = to_mv(ctypes.addressof(cl_img), 8).cast('Q')[0]
-rawbuf_ptr = to_mv(cl_buf_desc_ptr, 0x100).cast('Q')[20] # offset 0xA0 is a raw gpu pointer.
-
-# 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.
+| [CLANG (C Code)](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_clang.py) | Runs on CPU using the clang compiler | `clang` compiler in system `PATH` |
+| [LLVM](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_llvm.py) | Runs on CPU using the LLVM compiler infrastructure | `llvmlite` package installed |
@@ -20,9 +20,7 @@

 ::: tinygrad.Tensor.manual_seed
 ::: tinygrad.Tensor.rand
-::: tinygrad.Tensor.rand_like
 ::: tinygrad.Tensor.randn
-::: tinygrad.Tensor.randn_like
 ::: tinygrad.Tensor.randint
 ::: tinygrad.Tensor.normal
 ::: tinygrad.Tensor.uniform
@@ -22,7 +22,6 @@ Elementwise ops operate on a per element basis. They don't change the shape of t
 ::: tinygrad.Tensor.round
 ::: tinygrad.Tensor.isinf
 ::: tinygrad.Tensor.isnan
-::: tinygrad.Tensor.isfinite
 ::: tinygrad.Tensor.lerp
 ::: tinygrad.Tensor.square
 ::: tinygrad.Tensor.clamp
@@ -53,7 +52,7 @@ Elementwise ops operate on a per element basis. They don't change the shape of t
 ::: tinygrad.Tensor.erf
 ::: tinygrad.Tensor.gelu
 ::: tinygrad.Tensor.quick_gelu
-::: tinygrad.Tensor.leaky_relu
+::: tinygrad.Tensor.leakyrelu
 ::: tinygrad.Tensor.mish
 ::: tinygrad.Tensor.softplus
 ::: tinygrad.Tensor.softsign
@@ -64,19 +63,13 @@ Elementwise ops operate on a per element basis. They don't change the shape of t
 ::: tinygrad.Tensor.sub
 ::: tinygrad.Tensor.mul
 ::: tinygrad.Tensor.div
-::: tinygrad.Tensor.idiv
-::: tinygrad.Tensor.mod
-::: tinygrad.Tensor.bitwise_xor
-::: tinygrad.Tensor.bitwise_and
-::: tinygrad.Tensor.bitwise_or
-::: tinygrad.Tensor.bitwise_not
+::: tinygrad.Tensor.xor
 ::: tinygrad.Tensor.lshift
 ::: tinygrad.Tensor.rshift
 ::: tinygrad.Tensor.pow
 ::: tinygrad.Tensor.maximum
 ::: tinygrad.Tensor.minimum
 ::: tinygrad.Tensor.where
-::: tinygrad.Tensor.copysign

 ## Casting Ops

@@ -10,7 +10,6 @@

 ## Movement (high level)

-::: tinygrad.Tensor.__getitem__
 ::: tinygrad.Tensor.gather
 ::: tinygrad.Tensor.cat
 ::: tinygrad.Tensor.stack
@@ -25,5 +24,3 @@
 ::: tinygrad.Tensor.transpose
 ::: tinygrad.Tensor.flatten
 ::: tinygrad.Tensor.unflatten
-::: tinygrad.Tensor.roll
-::: tinygrad.Tensor.rearrange
@@ -6,10 +6,8 @@
 ::: tinygrad.Tensor.min
 ::: tinygrad.Tensor.any
 ::: tinygrad.Tensor.all
-::: tinygrad.Tensor.isclose
 ::: tinygrad.Tensor.mean
 ::: tinygrad.Tensor.var
-::: tinygrad.Tensor.var_mean
 ::: tinygrad.Tensor.std
 ::: tinygrad.Tensor.std_mean
 ::: tinygrad.Tensor.softmax
@@ -23,7 +21,6 @@

 ::: tinygrad.Tensor.avg_pool2d
 ::: tinygrad.Tensor.max_pool2d
-::: tinygrad.Tensor.max_unpool2d
 ::: tinygrad.Tensor.conv2d
 ::: tinygrad.Tensor.conv_transpose2d
 ::: tinygrad.Tensor.dot
@@ -35,10 +32,6 @@
 ::: tinygrad.Tensor.tril
 ::: tinygrad.Tensor.interpolate
 ::: tinygrad.Tensor.scatter
-::: tinygrad.Tensor.scatter_reduce
-::: tinygrad.Tensor.masked_select
-::: tinygrad.Tensor.sort
-::: tinygrad.Tensor.topk

 ## Neural Network (functional)

@@ -25,15 +25,10 @@
 ::: tinygrad.Tensor.replace
 ::: tinygrad.Tensor.assign
 ::: tinygrad.Tensor.detach
-::: tinygrad.Tensor.clone
 ::: tinygrad.Tensor.to
 ::: tinygrad.Tensor.to_
 ::: tinygrad.Tensor.shard
 ::: tinygrad.Tensor.shard_
 ::: tinygrad.Tensor.contiguous
 ::: tinygrad.Tensor.contiguous_backward
-
-## Gradient
-
-::: tinygrad.Tensor.gradient
 ::: tinygrad.Tensor.backward
@@ -1,4 +1,4 @@
-# model based off https://medium.com/data-science/going-beyond-99-mnist-handwritten-digits-recognition-cfff96337392
+# model based off https://towardsdatascience.com/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
@@ -1,36 +0,0 @@
-import sys, onnx, time, pickle
-from tinygrad import TinyJit, GlobalCounters, fetch, getenv
-from tinygrad.frontend.onnx import OnnxRunner
-from extra.onnx_helpers import get_example_inputs, validate
-
-def load_onnx_model(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(None) for k,v in kwargs.items()}).values())), prune=True, optimize=True)
-  return run_onnx_jit, run_onnx.graph_inputs
-
-if __name__ == "__main__":
-  onnx_file = fetch(sys.argv[1])
-  run_onnx_jit, input_specs = load_onnx_model(onnx_file)
-  print("loaded model")
-
-  for i in range(3):
-    new_inputs = get_example_inputs(input_specs)
-    GlobalCounters.reset()
-    print(f"run {i}")
-    run_onnx_jit(**new_inputs)
-
-  # run 20 times
-  for _ in range(20):
-    new_inputs = get_example_inputs(input_specs)
-    GlobalCounters.reset()
-    st = time.perf_counter()
-    out = run_onnx_jit(**new_inputs)
-    mt = time.perf_counter()
-    val = out.numpy()
-    et = time.perf_counter()
-    print(f"enqueue {(mt-st)*1e3:6.2f} ms -- total run {(et-st)*1e3:6.2f} ms")
-
-  if getenv("ORT"):
-    validate(onnx_file, new_inputs, rtol=1e-3, atol=1e-3)
-    print("model validated")
@@ -34,8 +34,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, 32, 32, 8)), strict=False)
-    nn.state.load_state_dict(model, fix_bf16(convert_from_huggingface(part2, 32, 32, 8)), strict=False)
+    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)

  if not os.path.isfile("/tmp/tokenizer.model"): create_fixed_tokenizer("/tmp/tokenizer.model")
  spp = SentencePieceProcessor(model_file="/tmp/tokenizer.model")
@@ -15,9 +15,9 @@ if __name__ == "__main__":
  if getenv("WEBGPU"):
    safe_save(get_state_dict(model), (dirname / "net.safetensors").as_posix())
    load_state_dict(model, safe_load(str(dirname / "net.safetensors")))
-  mode = "clang" if getenv("CPU", "") != "" else "webgpu" if getenv("WEBGPU", "") != "" else ""
+  mode = "clang" if getenv("CLANG", "") != "" else "webgpu" if getenv("WEBGPU", "") != "" else ""
  prg, inp_sizes, out_sizes, state = export_model(model, mode, Tensor.randn(1,3,224,224))
-  if getenv("CPU", "") == "":
+  if getenv("CLANG", "") == "":
    ext = "js" if getenv("WEBGPU", "") != "" else "json"
    with open(dirname / f"net.{ext}", "w") as text_file:
      text_file.write(prg)
@@ -68,6 +68,6 @@ if __name__ == "__main__":
    else printf("%s\\n", lbls[best_idx]);
  }""")

-    # CPU=1 python3 examples/compile_efficientnet.py | clang -O2 -lm -x c - -o recognize && DEBUG=1 time ./recognize docs/showcase/stable_diffusion_by_tinygrad.jpg
+    # CLANG=1 python3 examples/compile_efficientnet.py | clang -O2 -lm -x c - -o recognize && DEBUG=1 time ./recognize docs/showcase/stable_diffusion_by_tinygrad.jpg
    # category : 281 (tabby, tabby cat) with 9.452788
    print('\n'.join(cprog))
@@ -1,14 +1,14 @@
 # An example to compile a small Tensorflow model to extremely portable C code

 import os, sys
-os.environ["CPU"] = '1'
+os.environ["CLANG"] = '1'
 os.environ["JIT"] = '2'

 import numpy as np
 import subprocess
 import tensorflow as tf
 import tf2onnx
-from tinygrad.frontend.onnx import OnnxRunner
+from extra.onnx import get_run_onnx
 from tinygrad.tensor import Tensor
 from extra.export_model import export_model_clang, compile_net, jit_model

@@ -25,7 +25,7 @@ class TinyOnnx:
  def __init__(self, keras_model):
    input_signature = [tf.TensorSpec([1,32], tf.float32, name='x')]
    onnx_model, _ = tf2onnx.convert.from_keras(keras_model, input_signature, opset=13)
-    self.run_onnx = OnnxRunner(onnx_model)
+    self.run_onnx = get_run_onnx(onnx_model)

  def forward(self, x):
    return self.run_onnx({"x": x}, debug=False)['predictions']
@@ -117,7 +117,7 @@ def tts(
  stn_tst = text_mapper.get_text(text_to_synthesize, hps.data.add_blank, hps.data.text_cleaners)
  init_shape = stn_tst.shape
  assert init_shape[0] < pad_length, "text is too long"
-  x_tst, x_tst_lengths = stn_tst.pad(((0, pad_length - init_shape[0]),), value=1).unsqueeze(0), Tensor([init_shape[0]], dtype=dtypes.int64)
+  x_tst, x_tst_lengths = stn_tst.pad(((0, pad_length - init_shape[0]),), 1).unsqueeze(0), Tensor([init_shape[0]], dtype=dtypes.int64)
  sid = Tensor([speaker_id], dtype=dtypes.int64) if model_has_multiple_speakers else None

  # Perform inference.
@@ -1,13 +1,12 @@
 #!/usr/bin/env python3
-import os, argparse, contextlib
+import os, argparse
 from typing import Optional, Union
-with contextlib.suppress(ImportError): import tiktoken
+import tiktoken
 from tinygrad import Tensor, TinyJit, Device, GlobalCounters, Variable, dtypes
-from tinygrad.uop.ops import UOp
+from tinygrad.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")
@@ -135,12 +134,11 @@ class GPT2:
    # lm head and wte are tied
    weights['lm_head.weight'] = weights['wte.weight']

-    with WallTimeEvent(BenchEvent.LOAD_WEIGHTS):
-      load_state_dict(model, 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)

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

  def __init__(self, model, tokenizer):
@@ -188,12 +185,11 @@ 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):
-        with WallTimeEvent(BenchEvent.STEP):
-          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-1).bind(start_pos), temperature).tolist()
+        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()
      start_pos = len(toks[0])
      for i,t in enumerate(tok): toks[i].append(t)
    return [self.tokenizer.decode(x) for x in toks]
@@ -1,14 +1,14 @@
+from typing import List, Tuple
 from extra.models.resnet import ResNet50
 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.codegen.heuristic import hand_coded_optimizations
-from tinygrad.uop.ops import Ops, sym_infer
+from tinygrad.ops import Ops, sym_infer
 from tinygrad.device import Compiled
-from tinygrad.engine.search import beam_search, bufs_from_lin
+from tinygrad.engine.schedule import create_schedule
+from tinygrad.engine.search import time_linearizer, beam_search, bufs_from_lin
 from tinygrad.helpers import DEBUG, ansilen, getenv, colored, TRACEMETA
-from extra.optimization.helpers import time_linearizer

 def get_sched_resnet():
  mdl = ResNet50()
@@ -18,12 +18,12 @@ def get_sched_resnet():
  # run model twice to get only what changes, these are the kernels of the model
  for _ in range(2):
    out = mdl(Tensor.empty(BS, 3, 224, 224))
-    targets = [out]
+    targets = [out.lazydata]
    if getenv("BACKWARD"):
      optim.zero_grad()
      out.sparse_categorical_crossentropy(Tensor.empty(BS, dtype=dtypes.int)).backward()
-      targets += [x for x in optim.schedule_step()]
-    sched = Tensor.schedule(*targets)
+      targets += [x.lazydata for x in optim.schedule_step()]
+    sched = create_schedule(targets)
    print(f"schedule length {len(sched)}")
  return sched

@@ -42,16 +42,17 @@ def get_sched_bert():
  next_sentence_labels = Tensor.empty((BS, 1), dtype=dtypes.float32)

  # run model twice to get only what changes, these are the kernels of the model
+  seen = set()
  for _ in range(2):
    lm_logits, seq_relationship_logits = mdl(input_ids, attention_mask, masked_positions, segment_ids)
-    targets = [lm_logits, seq_relationship_logits]
+    targets = [lm_logits.lazydata, seq_relationship_logits.lazydata]
    if getenv("BACKWARD"):
      optim.zero_grad()
      loss = mdl.loss(lm_logits, seq_relationship_logits, masked_lm_ids, masked_lm_weights, next_sentence_labels)
      # ignore grad norm and loss scaler for now
      loss.backward()
-      targets += [x for x in optim.schedule_step()]
-    sched = Tensor.schedule(*targets)
+      targets += [x.lazydata for x in optim.schedule_step()]
+    sched = create_schedule(targets)
    print(f"schedule length {len(sched)}")
  return sched

@@ -80,11 +81,11 @@ if __name__ == "__main__":
    rawbufs = bufs_from_lin(Kernel(si.ast))

    # "linearize" the op into uops in different ways
-    lins: list[tuple[Kernel, str]] = []
+    lins: List[Tuple[Kernel, str]] = []

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

    # maybe try tensor cores
@@ -108,7 +109,7 @@ if __name__ == "__main__":
    choices = []
    for lin, nm in lins:
      tm = time_linearizer(lin, rawbufs, allow_test_size=False, cnt=10, disable_cache=True)
-      ops = (prg:=lin.to_program()).estimates.ops
+      ops = (prg:=lin.to_program()).op_estimate
      gflops = sym_infer(ops, {k:k.min for k in lin.ast.variables()})*1e-9/tm
      choices.append((tm, gflops, lin, prg, nm))

@@ -11,7 +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
+from tinygrad.multi import MultiLazyBuffer

 cifar_mean = [0.4913997551666284, 0.48215855929893703, 0.4465309133731618]
 cifar_std = [0.24703225141799082, 0.24348516474564, 0.26158783926049628]
@@ -35,6 +35,8 @@ class UnsyncedBatchNorm:
    self.num_batches_tracked = Tensor.zeros(1, dtype=dtypes.int, requires_grad=False)

  def __call__(self, x:Tensor):
+    if isinstance(x.lazydata, MultiLazyBuffer): assert x.lazydata.axis is None or x.lazydata.axis == 0 and len(x.lazydata.lbs) == self.num_devices
+
    xr = x.reshape(self.num_devices, -1, *x.shape[1:]).cast(dtypes.float32)
    batch_mean, batch_invstd = self.calc_stats(xr)
    ret = xr.batchnorm(
@@ -396,23 +398,20 @@ 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 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']]))
-
+      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
@@ -428,5 +427,4 @@ def train_cifar():
      raise ValueError(colored(f"{eval_acc_pct=} < {target}", "red"))

 if __name__ == "__main__":
-  with WallTimeEvent(BenchEvent.FULL):
-    train_cifar()
+  train_cifar()
@@ -17,6 +17,7 @@ canvas {  display: none; }
 * {  text-align: center;  font-family: monospace; }
 </style>
 <title>tinygrad has WebGPU</title>
+<script src="./net.js"></script>
 <link rel="icon" type="image/x-icon" href="https://raw.githubusercontent.com/tinygrad/tinygrad/master/docs/logo.png">
 </head>
 <body>
@@ -45,10 +46,7 @@ canvas {  display: none; }
 	const getDevice = async () => {
 		if (!navigator.gpu) error("WebGPU not supported.");
 		const adapter = await navigator.gpu.requestAdapter();
-		return await adapter.requestDevice({
-			requiredFeatures: ["shader-f16"],
-			powerPreference: "high-performance"
-		});
+		return await adapter.requestDevice();
 	};

 	const timer = async (func, label = "") => {
@@ -63,6 +61,8 @@ canvas {  display: none; }

 	const getLabels = async () => (await fetch("https://raw.githubusercontent.com/anishathalye/imagenet-simple-labels/master/imagenet-simple-labels.json")).json();

+	const getSavetensorBuffer = async () => new Uint8Array(await (await fetch("./net.safetensors")).arrayBuffer());
+
 	const reorderChannelsAndRemoveAlpha = (data) => {
 		const out = [];
 		let i = 0;
@@ -97,9 +97,9 @@ canvas {  display: none; }
 		try {
 			resultText.innerHTML = "loading..."
 			labels = await getLabels();
+			const safetensor = await getSavetensorBuffer();
 			const device = await getDevice();
-			const model = (await import("../../net.js")).default;
-			net = await timer(() => model.load(device, '../../net.safetensors'), "(compilation)");
+			net = await timer(() => setupNet(device, safetensor), "(compilation)");
 			resultText.innerHTML = "ready"
 		} catch (e) {
 			error(e)
@@ -13,7 +13,6 @@ 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)

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

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

-    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"]))
+    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"]))

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

-      # 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()
+    with Context(BEAM=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)
+          #print(k, v.shape, v.lazydata.axis)

-        # 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)

@@ -479,12 +475,11 @@ 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 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()
+          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,7 +7,6 @@ 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+"
@@ -48,7 +47,7 @@ def concat_weights(models, device=None):
    disk_tensors: List[Tensor] = [model[name] for model in models]
    if len(disk_tensors) == 1 or len(disk_tensors[0].shape) == 1:
      return disk_tensors[0].to(device=device)
-    axis = 1 if name.endswith((".attention.wo.weight", ".feed_forward.w2.weight")) else 0
+    axis = 1 if name.endswith(".attention.wo.weight") or name.endswith(".feed_forward.w2.weight") else 0
    lazy_tensors = [data.to(device=device) for data in disk_tensors]
    return lazy_tensors[0].cat(*lazy_tensors[1:], dim=axis)
  return {name: convert(name) for name in {name: None for model in models for name in model}}
@@ -74,17 +73,16 @@ class Int8Linear:
    self.scale = Tensor.ones(out_features, dtype=dtypes.half)

  def __call__(self, x):
-    return x.dot(self.weight.cast(self.scale.dtype).T*self.scale)
+    return x.dot(self.weight.cast(dtype=dtypes.half).T*self.scale)

  @staticmethod
-  def quantize(tensors, device, scale_dtype=dtypes.float16, quantize_embeds=False):
+  def quantize(tensors, device):
    new_tensors = {}
    for name,v in tensors.items():
-      if "feed_forward" in name or "attention.w" in name or (quantize_embeds and "tok_embeddings.weight" in name):
+      if "feed_forward" in name or "attention.w" in name:
        assert "weight" in name, name
-        v = v.cast(scale_dtype)
        scale = v.abs().max(axis=1) / 127.0
-        int8_weight = (v.T/scale).T.round().cast(dtype=dtypes.int8) # without round(), cast truncates -34.9 to -34
+        int8_weight = (v.T/scale).T.cast(dtype=dtypes.int8)
        new_tensors[name] = int8_weight
        new_tensors[name.replace('weight', 'scale')] = scale
        if isinstance(device, tuple):
@@ -92,20 +90,8 @@ class Int8Linear:
          new_tensors[name.replace('weight', 'scale')].shard_(device, axis=None)
      else:
        new_tensors[name] = v
-    if quantize_embeds: new_tensors.update({"output.weight": new_tensors["tok_embeddings.weight"], "output.scale": new_tensors["tok_embeddings.scale"]})
    return new_tensors

-class Int8Embedding:
-  def __init__(self, vocab_size:int, embed_size:int):
-    self.vocab_sz, self.embed_sz = vocab_size, embed_size
-    self.weight, self.scale = Tensor.ones(vocab_size, embed_size, dtype=dtypes.int8), Tensor.ones(vocab_size, dtype=dtypes.half)
-
-  def __call__(self, idx:Tensor) -> Tensor:
-    if not hasattr(self, 'arange'): self.arange = Tensor.arange(self.vocab_sz, requires_grad=False, device=self.weight.device).unsqueeze(-1)
-    big_shp = idx.shape+(self.vocab_sz, self.embed_sz)
-    arange, idx, vals = self.arange.expand(big_shp), idx.reshape(idx.shape+(1, 1)).expand(big_shp), (self.weight.cast(self.scale.dtype).T*self.scale).T
-    return (arange == idx).mul(vals).sum(-2, dtype=vals.dtype)
-
 def NF4Linear(block_size):
  _CODE = [
    -1.0, -0.6961928009986877, -0.5250730514526367, -0.39491748809814453, -0.28444138169288635, -0.18477343022823334, -0.09105003625154495, 0.0,
@@ -127,8 +113,7 @@ 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, quantize_embeds=False) -> dict[str, Tensor]:
-      assert not quantize_embeds  # TODO: support this?
+    def quantize(state_dict: dict[str, Tensor], device) -> dict[str, Tensor]:
      new_state_dict = {}
      for k, v in state_dict.items():
        if "feed_forward" in k or "attention.w" in k:
@@ -136,7 +121,7 @@ def NF4Linear(block_size):
          scale = (grouped.abs().max(axis=1, keepdim=True))
          coded = ((grouped / scale).unsqueeze(-1) - CODE.to(v.device)).abs().argmin(axis=-1).cast(dtypes.uint8).flatten()
          new_state_dict[k] = coded[::2] * 2 ** 4 + coded[1::2]
-          new_state_dict[k.replace(".weight", ".scale")] = scale.cast(scale_dtype)
+          new_state_dict[k.replace(".weight", ".scale")] = scale.cast(dtypes.float16)
          if isinstance(device, tuple):
            new_state_dict[k].shard_(device, axis=-1)
            new_state_dict[k.replace('weight', 'scale')].shard_(device, axis=None)
@@ -159,50 +144,47 @@ MODEL_PARAMS = {
    "files": 8
  }
 }
-def build_transformer(model_path: Path, model_size="8B", quantize=None, scale_dtype=dtypes.float16, device=None, max_context=8192, load_weights=True):
+def build_transformer(model_path: Path, model_size="8B", quantize=None, device=None):
  # build model
-  if quantize == "int8": linear, embedding, quantize_embeds = Int8Linear, Int8Embedding, True
-  elif quantize == "nf4": linear, embedding, quantize_embeds = NF4Linear(64), nn.Embedding, False
-  else: linear, embedding, quantize_embeds = nn.Linear, nn.Embedding, False
-  model = Transformer(**MODEL_PARAMS[model_size]["args"], linear=linear, embedding=embedding, max_context=max_context, jit=True)
-
-  if not load_weights: return model
+  if quantize == "int8": linear = Int8Linear
+  elif quantize == "nf4": linear = NF4Linear(64)
+  else: linear = nn.Linear
+  model = Transformer(**MODEL_PARAMS[model_size]["args"], linear=linear, max_context=8192, jit=True)

  # load 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)
+  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 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)
+      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
@@ -238,7 +220,7 @@ if __name__ == "__main__":
  parser = argparse.ArgumentParser()
  parser.add_argument("--download_model", action="store_true", help="Download a model")
  parser.add_argument("--model", type=Path, help="Model path")
-  parser.add_argument("--size", choices=["1B", "8B", "70B"], default="1B", help="Model size")
+  parser.add_argument("--size", choices=["1B", "8B", "70B"], default="8B", help="Model size")
  parser.add_argument("--shard", type=int, default=1, help="Shard the model across multiple devices")
  parser.add_argument("--quantize", choices=["int8", "nf4", "float16"], help="Quantization method")
  parser.add_argument("--no_api", action="store_true", help="Disable the api and run a cli test interface")
@@ -252,8 +234,8 @@ if __name__ == "__main__":
  parser.add_argument("--profile", action="store_true", help="Output profile data")
  args = parser.parse_args()

-  # download_model is the default without a model passed in
-  if args.download_model or not args.model:
+  assert (args.model and not args.download_model) or (not args.model and args.download_model), "either download or provide model"
+  if args.download_model:
    if args.size == "1B":
      fetch("https://huggingface.co/bofenghuang/Meta-Llama-3-8B/resolve/main/original/tokenizer.model", "tokenizer.model", subdir="llama3-1b-instruct")
      args.model = fetch("https://huggingface.co/bartowski/Llama-3.2-1B-Instruct-GGUF/resolve/main/Llama-3.2-1B-Instruct-Q6_K.gguf", "Llama-3.2-1B-Instruct-Q6_K.gguf", subdir="llama3-1b-instruct")
@@ -265,11 +247,11 @@ if __name__ == "__main__":
      fetch("https://huggingface.co/TriAiExperiments/SFR-Iterative-DPO-LLaMA-3-8B-R/resolve/main/model-00004-of-00004.safetensors", "model-00004-of-00004.safetensors", subdir="llama3-8b-sfr")
      args.model = fetch("https://huggingface.co/TriAiExperiments/SFR-Iterative-DPO-LLaMA-3-8B-R/raw/main/model.safetensors.index.json", "model.safetensors.index.json", subdir="llama3-8b-sfr")
    elif args.size == "70B":
-      subdir = "DeepSeek-R1-Distill-Llama-70B"
-      args.model = fetch("https://huggingface.co/deepseek-ai/DeepSeek-R1-Distill-Llama-70B/resolve/main/model.safetensors.index.json?download=true", "model.safetensors.index.json", subdir=subdir)
+      subdir = "Llama-3.1-Nemotron-70B-Instruct-HF"
+      args.model = fetch("https://huggingface.co/nvidia/Llama-3.1-Nemotron-70B-Instruct-HF/resolve/main/model.safetensors.index.json?download=true", "model.safetensors.index.json", subdir=subdir)
      fetch("https://huggingface.co/bofenghuang/Meta-Llama-3-8B/resolve/main/original/tokenizer.model", "tokenizer.model", subdir=subdir)
-      for i in range(17):
-        fetch(f"https://huggingface.co/deepseek-ai/DeepSeek-R1-Distill-Llama-70B/resolve/main/model-{i+1:05d}-of-000017.safetensors?download=true", f"model-{i+1:05d}-of-000017.safetensors", subdir=subdir)
+      for i in range(30):
+        fetch(f"https://huggingface.co/nvidia/Llama-3.1-Nemotron-70B-Instruct-HF/resolve/main/model-{i+1:05d}-of-00030.safetensors?download=true", f"model-{i+1:05d}-of-00030.safetensors", subdir=subdir)

  assert args.model is not None, "please provide --model option"

@@ -438,12 +420,11 @@ 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 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()
+          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])
@@ -2,12 +2,13 @@
 import os
 if "NOOPT" not in os.environ: os.environ["NOOPT"] = "1"
 from tinygrad import Device, nn, Tensor, dtypes, Variable
-Device.DEFAULT = "CPU"
+Device.DEFAULT = "CLANG"
 from train_gpt2 import GPT, GPTConfig
 from tinygrad.helpers import dedup, to_function_name, flatten, getenv, GlobalCounters, ansilen, to_function_name
+from tinygrad.engine.schedule import create_schedule
 from tinygrad.engine.realize import get_kernel, run_schedule
 from tinygrad.engine.memory import memory_planner
-from tinygrad.uop.ops import Ops
+from tinygrad.ops import Ops

 TIMING = getenv("TIMING")

@@ -25,7 +26,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 stabilize
+  for i in range(warmup_count):  # TODO: why does it take three and not two to stablize
    GlobalCounters.reset()
    X = Tensor.empty(4, 64, dtype=dtypes.int).reshape(B, T)
    Y = Tensor.empty(4, 64, dtype=dtypes.int).reshape(B, T)
@@ -36,16 +37,16 @@ if __name__ == "__main__":
      tensors = optimizer.schedule_step()
    else:
      tensors = []
-    sched = loss.schedule(*tensors)
+    sched = create_schedule([loss.lazydata] + [x.lazydata for x in tensors])
    print(f"calls {i}:", len(sched))
    #run_schedule(sched[:])
  sched = memory_planner(sched)
  ast_dedup = dedup([si.ast for si in sched if si.ast.op is Ops.SINK])
  srcs = {}
  for ast in ast_dedup:
-    k = get_kernel(Device["CPU"].renderer, ast)
+    k = get_kernel(Device["CLANG"].renderer, ast)
    k.linearize()
-    src = Device["CPU"].renderer.render(to_function_name(k.name), k.uops)
+    src = Device["CLANG"].renderer.render(to_function_name(k.name), k.uops)
    srcs[ast] = (k.name, src)
  print("functions:", len(srcs))
  used_buffers = dedup(flatten([si.bufs for si in sched]))
@@ -99,7 +99,7 @@ class GPT:

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

    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,7 +124,6 @@ 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
@@ -132,10 +131,6 @@ 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|>"})
@@ -170,32 +165,23 @@ 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
-  @Tensor.train()
-  def step(x:Tensor, y:Tensor) -> Tensor:
+  def step(x, y):
    _, loss = model(x, y)
    optimizer.zero_grad()
    loss.backward()
    return loss.realize(*optimizer.schedule_step())

-  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")
+  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")

  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, ...])
@@ -1,157 +0,0 @@
-# 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
-  @Tensor.test()
-  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,7 +3,6 @@ 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):
@@ -31,19 +30,18 @@ if __name__ == "__main__":
                      help="Path to the downloaded weights")
  args = parser.parse_args()

-  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)
+  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
@@ -55,8 +53,7 @@ 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"):
-        with WallTimeEvent(BenchEvent.STEP):
-          tok = model(Tensor([toks[start_pos:]]), 0 if start_pos == 0 else Variable("start_pos", 1, 1024).bind(start_pos), args.temperature).item()
+        tok = model(Tensor([toks[start_pos:]]), 0 if start_pos == 0 else Variable("start_pos", 1, 1024).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, OSX
+from tinygrad.helpers import getenv, prod, Context, round_up, tqdm

 ### ResNet

@@ -129,15 +129,14 @@ 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)
-  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))
+  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))
  procs = []

  try:
    # disk:shm is slower
-    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}")
+    #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")
    Y = [None] * (batch_size*BATCH_COUNT)

    for _ in range(cpu_count()):
@@ -171,13 +170,13 @@ def batch_load_resnet(batch_size=64, val=False, shuffle=True, seed=None, pad_fir

 def process_batch_bert(data: List[dict]) -> dict[str, Tensor]:
  return {
-    "input_ids": Tensor(np.concatenate([s["input_ids"] for s in data], axis=0), dtype=dtypes.int32, device="CPU"),
-    "input_mask": Tensor(np.concatenate([s["input_mask"] for s in data], axis=0), dtype=dtypes.int32, device="CPU"),
-    "segment_ids": Tensor(np.concatenate([s["segment_ids"] for s in data], axis=0), dtype=dtypes.int32, device="CPU"),
-    "masked_lm_positions": Tensor(np.concatenate([s["masked_lm_positions"] for s in data], axis=0), dtype=dtypes.int32, device="CPU"),
-    "masked_lm_ids": Tensor(np.concatenate([s["masked_lm_ids"] for s in data], axis=0), dtype=dtypes.int32, device="CPU"),
-    "masked_lm_weights": Tensor(np.concatenate([s["masked_lm_weights"] for s in data], axis=0), dtype=dtypes.float32, device="CPU"),
-    "next_sentence_labels": Tensor(np.concatenate([s["next_sentence_labels"] for s in data], axis=0), dtype=dtypes.int32, device="CPU"),
+    "input_ids": Tensor(np.concatenate([s["input_ids"] for s in data], axis=0), dtype=dtypes.float32),
+    "input_mask": Tensor(np.concatenate([s["input_mask"] for s in data], axis=0), dtype=dtypes.default_float),
+    "segment_ids": Tensor(np.concatenate([s["segment_ids"] for s in data], axis=0), dtype=dtypes.float32),
+    "masked_lm_positions": Tensor(np.concatenate([s["masked_lm_positions"] for s in data], axis=0), dtype=dtypes.float32),
+    "masked_lm_ids": Tensor(np.concatenate([s["masked_lm_ids"] for s in data], axis=0), dtype=dtypes.float32),
+    "masked_lm_weights": Tensor(np.concatenate([s["masked_lm_weights"] for s in data], axis=0), dtype=dtypes.float32),
+    "next_sentence_labels": Tensor(np.concatenate([s["next_sentence_labels"] for s in data], axis=0), dtype=dtypes.float32),
  }

 def load_file(file: str):
@@ -224,8 +223,14 @@ def batch_load_train_bert(BS:int):
  assert cycle_length > 0, "cycle_length must be greater than 0"

  dataset = InterleavedDataset(train_files, cycle_length)
+  buffer = [dataset.get() for _ in range(1000)]
  while True:
-    yield process_batch_bert([dataset.get() for _ in range(BS)])
+    batch = []
+    for _ in range(BS):
+      index = random.randint(0, 999)
+      batch.append(buffer[index])
+      buffer[index] = dataset.get()
+    yield process_batch_bert(batch)

 # Reference: https://github.com/mlcommons/training/blob/1c8a098ae3e70962a4f7422c0b0bd35ae639e357/language_model/tensorflow/bert/run_pretraining.py, Line 416
 def batch_load_val_bert(BS:int):
@@ -313,7 +318,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):
@@ -349,167 +354,6 @@ def batch_load_unet3d(preprocessed_dataset_dir:Path, batch_size:int=6, val:bool=
      # happens with BENCHMARK set
      pass

-### RetinaNet
-
-def load_retinanet_data(base_dir:Path, val:bool, queue_in:Queue, queue_out:Queue,
-                        imgs:Tensor, boxes:Tensor, labels:Tensor, matches:Tensor|None=None,
-                        anchors:Tensor|None=None, seed:int|None=None):
-  from extra.datasets.openimages import image_load, random_horizontal_flip, resize
-  from examples.mlperf.helpers import box_iou, find_matches, generate_anchors
-  import torch
-
-  while (data:=queue_in.get()) is not None:
-    idx, img, tgt = data
-    img = image_load(base_dir, img["subset"], img["file_name"])
-
-    if val:
-      img = resize(img)[0]
-    else:
-      if seed is not None:
-        np.random.seed(seed)
-        random.seed(seed)
-        torch.manual_seed(seed)
-
-      img, tgt = random_horizontal_flip(img, tgt)
-      img, tgt, _ = resize(img, tgt=tgt)
-      match_quality_matrix = box_iou(tgt["boxes"], (anchor := np.concatenate(generate_anchors((800, 800)))))
-      match_idxs = find_matches(match_quality_matrix, allow_low_quality_matches=True)
-      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()
-
-    imgs[idx].contiguous().realize().lazydata.base.realized.as_buffer(force_zero_copy=True)[:] = img.tobytes()
-
-    queue_out.put(idx)
-  queue_out.put(None)
-
-def batch_load_retinanet(dataset, val:bool, base_dir:Path, batch_size:int=32, shuffle:bool=True, seed:int|None=None):
-  def _enqueue_batch(bc):
-    from extra.datasets.openimages import prepare_target
-    for idx in range(bc * batch_size, (bc+1) * batch_size):
-      img = dataset.loadImgs(next(dataset_iter))[0]
-      ann = dataset.loadAnns(dataset.getAnnIds(img_id:=img["id"]))
-      tgt = prepare_target(ann, img_id, (img["height"], img["width"]))
-
-      if img_ids is not None:
-        img_ids[idx] = img_id
-
-      if img_sizes is not None:
-        img_sizes[idx] = tgt["image_size"]
-
-      queue_in.put((idx, img, tgt))
-
-  def _setup_shared_mem(shm_name:str, size:tuple[int, ...], dtype:dtypes) -> tuple[shared_memory.SharedMemory, Tensor]:
-    if 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(size))
-    shm_tensor = Tensor.empty(*size, dtype=dtype, device=f"disk:/dev/shm/{shm_name}")
-    return shm, shm_tensor
-
-  image_ids = sorted(dataset.imgs.keys())
-  batch_count = min(32, len(image_ids) // batch_size)
-
-  queue_in, queue_out = Queue(), Queue()
-  procs, data_out_count = [], [0] * batch_count
-
-  shm_imgs, imgs = _setup_shared_mem("retinanet_imgs", (batch_size * batch_count, 800, 800, 3), dtypes.uint8)
-
-  if val:
-    boxes, labels, matches, anchors = None, None, None, None
-    img_ids, img_sizes = [None] * (batch_size * batch_count), [None] * (batch_size * batch_count)
-  else:
-    img_ids, img_sizes = None, None
-    shm_boxes, boxes = _setup_shared_mem("retinanet_boxes", (batch_size * batch_count, 120087, 4), dtypes.float32)
-    shm_labels, labels = _setup_shared_mem("retinanet_labels", (batch_size * batch_count, 120087), dtypes.int64)
-    shm_matches, matches = _setup_shared_mem("retinanet_matches", (batch_size * batch_count, 120087), dtypes.int64)
-    shm_anchors, anchors = _setup_shared_mem("retinanet_anchors", (batch_size * batch_count, 120087, 4), dtypes.float64)
-
-  shutdown = False
-  class Cookie:
-    def __init__(self, bc):
-      self.bc = bc
-    def __del__(self):
-      if not shutdown:
-        try: _enqueue_batch(self.bc)
-        except StopIteration: pass
-
-  def shuffle_indices(indices, seed):
-    rng = random.Random(seed)
-    rng.shuffle(indices)
-
-  if shuffle: shuffle_indices(image_ids, seed=seed)
-  dataset_iter = iter(image_ids)
-
-  try:
-    for _ in range(cpu_count()):
-      proc = Process(
-        target=load_retinanet_data,
-        args=(base_dir, val, queue_in, queue_out, imgs, boxes, labels),
-        kwargs={"matches": matches, "anchors": anchors, "seed": seed}
-      )
-      proc.daemon = True
-      proc.start()
-      procs.append(proc)
-
-    for bc in range(batch_count):
-      _enqueue_batch(bc)
-
-    for _ in range(len(image_ids) // batch_size):
-      while True:
-        bc = queue_out.get() // batch_size
-        data_out_count[bc] += 1
-        if data_out_count[bc] == batch_size: break
-
-      data_out_count[bc] = 0
-
-      if val:
-        yield (imgs[bc * batch_size:(bc + 1) * batch_size],
-               img_ids[bc * batch_size:(bc + 1) * batch_size],
-               img_sizes[bc * batch_size:(bc + 1) * batch_size],
-               Cookie(bc))
-      else:
-        yield (imgs[bc * batch_size:(bc + 1) * batch_size],
-               boxes[bc * batch_size:(bc + 1) * batch_size],
-               labels[bc * batch_size:(bc + 1) * batch_size],
-               matches[bc * batch_size:(bc + 1) * batch_size],
-               anchors[bc * batch_size:(bc + 1) * batch_size],
-               Cookie(bc))
-  finally:
-    shutdown = True
-
-    for _ in procs: queue_in.put(None)
-    queue_in.close()
-
-    for _ in procs:
-      while queue_out.get() is not None: pass
-    queue_out.close()
-
-    # shutdown processes
-    for proc in procs: proc.join()
-
-    shm_imgs.close()
-
-    if not val:
-      shm_boxes.close()
-      shm_labels.close()
-      shm_matches.close()
-      shm_anchors.close()
-
-    try:
-      shm_imgs.unlink()
-
-      if not val:
-        shm_boxes.unlink()
-        shm_labels.unlink()
-        shm_matches.unlink()
-        shm_anchors.unlink()
-    except FileNotFoundError:
-      # happens with BENCHMARK set
-      pass
-
 if __name__ == "__main__":
  def load_unet3d(val):
    assert not val, "validation set is not supported due to different sizes on inputs"
@@ -530,14 +374,6 @@ if __name__ == "__main__":
      for x,y,c in batch_load_resnet(val=val):
        pbar.update(x.shape[0])

-  def load_retinanet(val):
-    from extra.datasets.openimages import BASEDIR, download_dataset
-    from pycocotools.coco import COCO
-    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])
-
  load_fn_name = f"load_{getenv('MODEL', 'resnet')}"
  if load_fn_name in globals():
    globals()[load_fn_name](getenv("VAL", 1))
@@ -1,7 +1,6 @@
 from collections import OrderedDict
 import unicodedata
 from typing import Optional
-import math
 import numpy as np
 from tinygrad.nn import state
 from tinygrad.tensor import Tensor, dtypes
@@ -196,18 +195,20 @@ def get_bert_qa_prediction(features, example, start_end_logits):
  return "empty"

 def get_mlperf_bert_config():
-  """benchmark is BERT-large"""
-  ret = {"attention_probs_dropout_prob": 0.1, "hidden_dropout_prob": 0.1, "vocab_size": 30522, "type_vocab_size": 2, "max_position_embeddings": 512}
+  """Config is BERT-large"""
+  return {
+    "attention_probs_dropout_prob": 0.1,
+    "hidden_dropout_prob": 0.1,
+    "hidden_size": 1024,
+    "intermediate_size": 4096,
+    "max_position_embeddings": 512,
+    "num_attention_heads": 16,
+    "num_hidden_layers": 24,
+    "type_vocab_size": 2,
+    "vocab_size": 30522
+  }

-  match (bert_size:=getenv("BERT_SIZE", "large")):
-    case "large": ret.update({"hidden_size": 1024, "intermediate_size": 4096, "num_attention_heads": 16, "num_hidden_layers": 24})
-    case "tiny": ret.update({"hidden_size": 128, "intermediate_size": 512, "num_attention_heads": 2, "num_hidden_layers": 2})
-    case _: raise RuntimeError(f"unhandled {bert_size=}")
-
-  if (bert_layers:=getenv("BERT_LAYERS")): ret["num_hidden_layers"] = bert_layers
-  return ret
-
-def get_mlperf_bert_model():
+def get_mlperf_bert_model(checkpoint_path:Optional[str]=None):
  from extra.models import bert
  from examples.mlperf.initializers import LinearBert, EmbeddingBert, LayerNormBert

@@ -219,138 +220,21 @@ def get_mlperf_bert_model():
  config = get_mlperf_bert_config()
  if getenv("DISABLE_DROPOUT", 0):
    config["hidden_dropout_prob"] = config["attention_probs_dropout_prob"] = 0.0
-  return BertForPretraining(**config)
+  model = BertForPretraining(**config)
+  return model.load_from_pretrained(checkpoint_path) if checkpoint_path else model

-def get_fake_data_bert(BS:int):
+def get_data_bert(GPUS:list[str], it):
+  data: dict[str, Tensor] = next(it)
+  for key in data.keys(): data[key].shard_(GPUS, axis=0)
+  return data
+
+def get_fake_data_bert(GPUS:list[str], BS:int):
  return {
-    "input_ids": Tensor.empty((BS, 512), dtype=dtypes.int32, device="CPU"),
-    "input_mask": Tensor.empty((BS, 512), dtype=dtypes.int32, device="CPU"),
-    "segment_ids": Tensor.empty((BS, 512), dtype=dtypes.int32, device="CPU"),
-    "masked_lm_positions": Tensor.empty((BS, 76), dtype=dtypes.int32, device="CPU"),
-    "masked_lm_ids": Tensor.empty((BS, 76), dtype=dtypes.int32, device="CPU"),
-    "masked_lm_weights": Tensor.empty((BS, 76), dtype=dtypes.float32, device="CPU"),
-    "next_sentence_labels": Tensor.empty((BS, 1), dtype=dtypes.int32, device="CPU"),
+    "input_ids": Tensor.empty((BS, 512), dtype=dtypes.float32).contiguous().shard_(GPUS, axis=0),
+    "input_mask": Tensor.empty((BS, 512), dtype=dtypes.default_float).contiguous().shard_(GPUS, axis=0),
+    "segment_ids": Tensor.empty((BS, 512), dtype=dtypes.float32).contiguous().shard_(GPUS, axis=0),
+    "masked_lm_positions": Tensor.empty((BS, 76), dtype=dtypes.float32).contiguous().shard_(GPUS, axis=0),
+    "masked_lm_ids": Tensor.empty((BS, 76), dtype=dtypes.float32).contiguous().shard_(GPUS, axis=0),
+    "masked_lm_weights": Tensor.empty((BS, 76), dtype=dtypes.float32).contiguous().shard_(GPUS, axis=0),
+    "next_sentence_labels": Tensor.empty((BS, 1), dtype=dtypes.float32).contiguous().shard_(GPUS, axis=0),
  }
-
-def find_matches(match_quality_matrix:np.ndarray, high_threshold:float=0.5, low_threshold:float=0.4, allow_low_quality_matches:bool=False) -> np.ndarray:
-  BELOW_LOW_THRESHOLD, BETWEEN_THRESHOLDS = -1, -2
-
-  def _set_low_quality_matches_(matches:np.ndarray, all_matches:np.ndarray, match_quality_matrix:np.ndarray):
-    highest_quality_foreach_gt = np.max(match_quality_matrix, axis=1)
-    pred_inds_to_update = np.nonzero(match_quality_matrix == highest_quality_foreach_gt[:, None])[1]
-    matches[pred_inds_to_update] = all_matches[pred_inds_to_update]
-
-  assert low_threshold <= high_threshold
-
-  matched_vals, matches = match_quality_matrix.max(axis=0), match_quality_matrix.argmax(axis=0)
-  all_matches = np.copy(matches) if allow_low_quality_matches else None
-  below_low_threshold = matched_vals < low_threshold
-  between_thresholds = (matched_vals >= low_threshold) & (matched_vals < high_threshold)
-  matches[below_low_threshold] = BELOW_LOW_THRESHOLD
-  matches[between_thresholds] = BETWEEN_THRESHOLDS
-
-  if allow_low_quality_matches:
-    assert all_matches is not None
-    _set_low_quality_matches_(matches, all_matches, match_quality_matrix)
-
-  return matches
-
-def box_iou(boxes1:np.ndarray, boxes2:np.ndarray) -> np.ndarray:
-  def _box_area(boxes:np.ndarray) -> np.ndarray: return (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
-
-  def _box_inter_union(boxes1:np.ndarray, boxes2:np.ndarray) -> tuple[np.ndarray, np.ndarray]:
-    area1, area2 = _box_area(boxes1), _box_area(boxes2)
-    lt, rb = np.maximum(boxes1[:, None, :2], boxes2[:, :2]), np.minimum(boxes1[:, None, 2:], boxes2[:, 2:])
-    wh = np.clip(rb - lt, a_min=0, a_max=None)
-    inter = wh[:, :, 0] * wh[:, :, 1]
-    union = area1[:, None] + area2 - inter
-    return inter, union
-
-  inter, union = _box_inter_union(boxes1, boxes2)
-  return inter / union
-
-def generate_anchors(input_size:tuple[int, int], scales:Optional[tuple[Tensor, ...]]=None, aspect_ratios:Optional[tuple[Tensor, ...]]=None) -> list[np.ndarray]:
-  def _compute_grid_sizes(input_size:tuple[int, int]) -> np.ndarray:
-    return np.ceil(np.array(input_size)[None, :] / 2 ** np.arange(3, 8)[:, None])
-
-  scales = tuple((i, int(i * 2 ** (1/3)), int(i * 2 ** (2/3))) for i in 2 ** np.arange(5, 10)) if scales is None else scales
-  aspect_ratios = ((0.5, 1.0, 2.0),) * len(scales) if aspect_ratios is None else aspect_ratios
-  aspect_ratios = tuple(ar for ar in aspect_ratios)
-  grid_sizes = _compute_grid_sizes(input_size)
-
-  assert len(scales) == len(aspect_ratios) == len(grid_sizes), "scales, aspect_ratios, and grid_sizes must have the same length"
-
-  anchors = []
-  for s, ar, gs in zip(scales, aspect_ratios, grid_sizes):
-    s, ar = np.array(s), np.array(ar)
-    h_ratios = np.sqrt(ar)
-    w_ratios = 1 / h_ratios
-    ws = (w_ratios[:, None] * s[None, :]).reshape(-1)
-    hs = (h_ratios[:, None] * s[None, :]).reshape(-1)
-    base_anchors = (np.stack([-ws, -hs, ws, hs], axis=1) / 2).round()
-    stride_h, stride_w = input_size[0] // gs[0], input_size[1] // gs[1]
-    shifts_x, shifts_y = np.meshgrid(np.arange(gs[1]) * stride_w, np.arange(gs[0]) * stride_h)
-    shifts_x, shifts_y = shifts_x.reshape(-1), shifts_y.reshape(-1)
-    shifts = np.stack([shifts_x, shifts_y, shifts_x, shifts_y], axis=1, dtype=np.float32)
-    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
@@ -1,5 +1,5 @@
 import math
-from typing import Union
+from typing import Union, Tuple

 from tinygrad import Tensor, nn, dtypes
 from tinygrad.helpers import prod, argfix
@@ -53,12 +53,10 @@ 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)
-    # 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()
+    return (arange == idx).mul(vals).sum(2, acc_dtype=vals.dtype)

 class LayerNormBert:
-  def __init__(self, normalized_shape:Union[int, tuple[int, ...]], eps:float=1e-12, elementwise_affine:bool=True):
+  def __init__(self, normalized_shape:Union[int, Tuple[int, ...]], eps:float=1e-12, elementwise_affine:bool=True):
    self.normalized_shape = (normalized_shape,) if isinstance(normalized_shape, int) else tuple(normalized_shape)
    self.axis, self.eps, self.elementwise_affine = tuple(-1-i for i in range(len(self.normalized_shape))), eps, elementwise_affine
    self.weight, self.bias = (Tensor.ones(*self.normalized_shape, dtype=dtypes.float32), Tensor.zeros(*self.normalized_shape, dtype=dtypes.float32)) if elementwise_affine else (None, None)
@@ -68,62 +66,3 @@ 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)
@@ -1,29 +1,6 @@
 from examples.mlperf.metrics import dice_score
-from tinygrad import Tensor

 def dice_ce_loss(pred, tgt):
  ce = pred.permute(0, 2, 3, 4, 1).sparse_categorical_crossentropy(tgt.squeeze(1))
  dice = (1.0 - dice_score(pred, tgt, argmax=False, to_one_hot_x=False)).mean()
  return (dice + ce) / 2
-
-def sigmoid_focal_loss(pred:Tensor, tgt:Tensor, alpha:float=0.25, gamma:float=2.0, reduction:str="none") -> Tensor:
-  assert reduction in ["mean", "sum", "none"], f"unsupported reduction {reduction}"
-  p, ce_loss = pred.sigmoid(), pred.binary_crossentropy_logits(tgt, reduction="none")
-  p_t = p * tgt + (1 - p) * (1 - tgt)
-  loss = ce_loss * ((1 - p_t) ** gamma)
-
-  if alpha >= 0:
-    alpha_t = alpha * tgt + (1 - alpha) * (1 - tgt)
-    loss = loss * alpha_t
-
-  if reduction == "mean": loss = loss.mean()
-  elif reduction == "sum": loss = loss.sum()
-  return loss
-
-def l1_loss(pred:Tensor, tgt:Tensor, reduction:str="none") -> Tensor:
-  assert reduction in ["mean", "sum", "none"], f"unsupported reduction {reduction}"
-  loss = (pred - tgt).abs()
-
-  if reduction == "mean": loss = loss.mean()
-  elif reduction == "sum": loss = loss.sum()
-  return loss
@@ -5,63 +5,61 @@ 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
-  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
+  # 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
@@ -83,43 +81,47 @@ 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
-  tlog("imports")
-
-  mdl = RetinaNet(ResNeXt50_32X4D())
-  mdl.load_from_pretrained()
-  tlog("loaded models")
-
-  coco = COCO(download_dataset(base_dir:=getenv("BASEDIR", BASEDIR), 'validation'))
+  coco = COCO(download_dataset(base_dir:=getenv("BASE_DIR", 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")

-  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")
+  from tinygrad.engine.jit import TinyJit
+  mdlrun = TinyJit(lambda x: mdl(input_fixup(x)).realize())
+
+  n, bs = 0, 8
  st = time.perf_counter()
-  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 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}
      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)
@@ -127,18 +129,13 @@ 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)))
-    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
+    st = time.perf_counter()

  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
@@ -1,15 +1,14 @@
-import os, time, math, functools, random
+import os, time, math, functools
 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, Profiling
+from tinygrad.helpers import getenv, BEAM, WINO, round_up, diskcache_clear, FUSE_CONV_BW
 from tinygrad.nn.state import get_parameters, get_state_dict, safe_load, safe_save
-from tinygrad.nn.optim import LAMB, LARS, SGD, OptimizerGroup, Adam
+from tinygrad.nn.optim import LAMB, LARS, SGD, OptimizerGroup

 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

@@ -80,7 +79,7 @@ def train_resnet():
  lr_warmup_epochs  = config["lr_warmup_epochs"]  = getenv("WARMUP_EPOCHS", 2)
  decay             = config["decay"]             = getenv("DECAY", 2e-4)

-  loss_scaler       = config["LOSS_SCALER"]       = getenv("LOSS_SCALER", 256.0 if dtypes.default_float == dtypes.float16 else 1.0)
+  loss_scaler       = config["LOSS_SCALER"]       = getenv("LOSS_SCALER", 128.0 if dtypes.default_float == dtypes.float16 else 1.0)

  target, achieved  = getenv("TARGET", 0.759), False
  eval_start_epoch  = getenv("EVAL_START_EPOCH", 0)
@@ -206,25 +205,24 @@ def train_resnet():
    st = time.perf_counter()
    while proc is not None:
      GlobalCounters.reset()
-      with WallTimeEvent(BenchEvent.STEP):
-        (loss, top_1), y, proc = train_step(proc[0], proc[1]), proc[2], proc[3]
+      (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:
@@ -275,7 +273,7 @@ def train_resnet():
      else:
        it = iter(tqdm(batch_load_resnet(batch_size=EVAL_BS, val=True, shuffle=False, pad_first_batch=True), total=steps_in_val_epoch))
        i, proc = 0, data_get(it)
-
+        
      prev_cookies = []
      while proc is not None:
        GlobalCounters.reset()
@@ -345,349 +343,8 @@ def train_resnet():
        safe_save(get_training_state(model, optimizer_group, scheduler_group), fn)

 def train_retinanet():
-  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), Tensor.test():
-        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
+  # TODO: Retinanet
+  pass

 def train_unet3d():
  """
@@ -789,7 +446,7 @@ def train_unet3d():
    loss.backward()
    optim.step()
    return loss.realize()
-
+  
  @Tensor.train(mode=False)
  @Tensor.test()
  def eval_step(model, x, y):
@@ -798,7 +455,7 @@ def train_unet3d():
    loss = dice_ce_loss(y_hat, y)
    score = dice_score(y_hat, y)
    return loss.realize(), score.realize()
-
+  
  if WANDB: wandb.init(config=config, project=PROJ_NAME)

  step_times, start_epoch = [], 1
@@ -807,7 +464,7 @@ def train_unet3d():
  next_eval_at = start_eval_at

  print(f"Training on {GPUS}")
-
+  
  if BENCHMARK: print("Benchmarking UNet3D")
  else: print(f"Start evaluation at epoch {start_eval_at} and every {evaluate_every} epoch(s) afterwards")

@@ -894,7 +551,7 @@ def train_unet3d():

        if mean_dice >= TARGET_METRIC:
          is_successful = True
-          save_checkpoint(get_state_dict(model), "./ckpts/unet3d.safe")
+          save_checkpoint(get_state_dict(model), f"./ckpts/unet3d.safe")
        elif mean_dice < 1e-6:
          print("Model diverging. Aborting.")
          diverged = True
@@ -915,49 +572,39 @@ def train_rnnt():
  pass

@TinyJit
-def train_step_bert(model, optimizer, scheduler, loss_scaler:float, input_ids:Tensor, segment_ids:Tensor, attention_mask:Tensor,
-                    masked_positions:Tensor, masked_lm_ids:Tensor, masked_lm_weights:Tensor, next_sentence_labels:Tensor, GPUS):
-  for t in [input_ids, segment_ids, attention_mask, masked_positions, masked_lm_ids, masked_lm_weights, next_sentence_labels]:
-    if len(GPUS) > 1: t.shard_(GPUS, axis=0)
-    else: t.to_(GPUS[0])
+def train_step_bert(model, optimizer, scheduler, loss_scaler:float, input_ids:Tensor, segment_ids:Tensor, attention_mask:Tensor, masked_positions:Tensor, masked_lm_ids:Tensor, masked_lm_weights:Tensor, next_sentence_labels:Tensor):
  optimizer.zero_grad()

  lm_logits, seq_relationship_logits = model(input_ids, attention_mask, masked_positions, segment_ids)
  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)
-  for p in optimizer.params:
+  global_norm = Tensor([0.0], dtype=dtypes.float32, device=optimizer[0].device).realize()
+  for p in optimizer.params: 
    p.grad = p.grad / loss_scaler
    global_norm += p.grad.float().square().sum()
-  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)
+  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)

  optimizer.step()
  scheduler.step()
-  # TODO: no to("CPU") here because it blocks and messes the python time
-  Tensor.realize(loss, global_norm, optimizer.optimizers[0].lr)
-  return loss, global_norm, optimizer.optimizers[0].lr
+  return loss.realize()

@TinyJit
-def eval_step_bert(model, input_ids:Tensor, segment_ids:Tensor, attention_mask:Tensor, masked_positions:Tensor, masked_lm_ids:Tensor,
-                   masked_lm_weights:Tensor, next_sentence_labels:Tensor, GPUS):
-  for t in [input_ids, segment_ids, attention_mask, masked_positions, masked_lm_ids, masked_lm_weights, next_sentence_labels]:
-    if len(GPUS) > 1: t.shard_(GPUS, axis=0)
-    else: t.to_(GPUS[0])
+def eval_step_bert(model, input_ids:Tensor, segment_ids:Tensor, attention_mask:Tensor, masked_positions:Tensor, masked_lm_ids:Tensor, masked_lm_weights:Tensor, next_sentence_labels:Tensor):
  lm_logits, seq_relationship_logits = model(input_ids, attention_mask, masked_positions, segment_ids)
-  masked_lm_accuracy, seq_relationship_accuracy, masked_lm_loss, next_sentence_loss = \
-    model.accuracy(lm_logits, seq_relationship_logits, masked_lm_ids, masked_lm_weights, next_sentence_labels)
-  for t in [masked_lm_accuracy, seq_relationship_accuracy, masked_lm_loss, next_sentence_loss]:
-    t.to_("CPU")
-  Tensor.realize(masked_lm_accuracy, seq_relationship_accuracy, masked_lm_loss, next_sentence_loss)
-  return masked_lm_accuracy, seq_relationship_accuracy, masked_lm_loss, next_sentence_loss
+  masked_lm_accuracy, seq_relationship_accuracy, masked_lm_loss, next_sentence_loss = model.accuracy(lm_logits, seq_relationship_logits, masked_lm_ids, masked_lm_weights, next_sentence_labels)
+  return {
+    "masked_lm_accuracy": masked_lm_accuracy.realize(),
+    "next_sentence_accuracy": seq_relationship_accuracy.realize(),
+    "masked_lm_loss": masked_lm_loss.realize(),
+    "next_sentence_loss": next_sentence_loss.realize()
+  }

 def train_bert():
  # NOTE: pip install tensorflow, wandb required
  from examples.mlperf.dataloader import batch_load_train_bert, batch_load_val_bert
-  from examples.mlperf.helpers import get_mlperf_bert_model, get_fake_data_bert
+  from examples.mlperf.helpers import get_mlperf_bert_model, get_data_bert, get_fake_data_bert
  from examples.mlperf.lr_schedulers import PolynomialDecayWithWarmup

  config = {}
@@ -981,7 +628,7 @@ def train_bert():
    MLLOGGER.logger.propagate = False

    if INITMLPERF:
-      assert BENCHMARK, "BENCHMARK must be set for INITMLPERF"
+      assert BENCHMARK, f"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)
@@ -1002,11 +649,9 @@ def train_bert():
  # ** hyperparameters **
  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)
+  max_lr             = config["OPT_BASE_LEARNING_RATE"] = getenv("OPT_BASE_LEARNING_RATE", 0.0001 * math.sqrt(BS/66))

-  train_steps        = config["TRAIN_STEPS"]            = getenv("TRAIN_STEPS", 3600000 // BS)
+  train_steps        = config["TRAIN_STEPS"]            = getenv("TRAIN_STEPS", 3630000 // 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
@@ -1015,7 +660,7 @@ def train_bert():
  save_ckpt_dir      = config["SAVE_CKPT_DIR"]          = getenv("SAVE_CKPT_DIR", "./ckpts")
  init_ckpt          = config["INIT_CKPT_DIR"]          = getenv("INIT_CKPT_DIR", BASEDIR)

-  loss_scaler        = config["LOSS_SCALER"]            = getenv("LOSS_SCALER", 2.0**11 if dtypes.default_float == dtypes.float16 else 1.0)
+  loss_scaler        = config["LOSS_SCALER"]            = getenv("LOSS_SCALER", 2.0**10 if dtypes.default_float == dtypes.float16 else 1.0)
  decay              = config["DECAY"]                  = getenv("DECAY", 0.01)
  epsilon            = config["EPSILON"]                = getenv("EPSILON", 1e-6)
  poly_power         = config["POLY_POWER"]             = getenv("POLY_POWER", 1.0)
@@ -1040,18 +685,11 @@ def train_bert():

  # ** init model **

-  model = get_mlperf_bert_model()
-  if RUNMLPERF:
-    model.load_from_pretrained(init_ckpt)
-  else:
-    # for init, zero out all weights
-    for p in get_parameters(model):
-      p = p.assign(Tensor.zeros_like(p).contiguous()).realize()
-
+  model = get_mlperf_bert_model(init_ckpt if RUNMLPERF else None)
+  
+  for _, x in get_state_dict(model).items():
+    x.realize().to_(GPUS)
  parameters = get_parameters(model)
-  if len(GPUS) > 1:
-    for p in parameters:
-      p.to_(GPUS)

  # ** Log run config **
  for key, value in config.items(): print(f'HParam: "{key}": {value}')
@@ -1059,8 +697,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, 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_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_group = OptimizerGroup(optimizer_wd, optimizer_no_wd)

  # ** LR scheduler **
@@ -1079,8 +717,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=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_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_LR_DECAY_POLY_POWER, value=config["POLY_POWER"])
      MLLOGGER.event(key=mllog_constants.OPT_LAMB_EPSILON, value=config["EPSILON"])

@@ -1097,7 +735,7 @@ def train_bert():
  previous_step = None
  if ckpt:=getenv("RESUME", ""):
    load_training_state(model, optimizer_group, scheduler_group, safe_load(ckpt))
-    start_step = int(scheduler_wd.epoch_counter.item())
+    start_step = int(scheduler_wd.epoch_counter.numpy().item())
    print(f"resuming from {ckpt} at step {start_step}")

  if RUNMLPERF:
@@ -1105,76 +743,70 @@ def train_bert():
    eval_it = iter(batch_load_val_bert(EVAL_BS))
    train_it = iter(tqdm(batch_load_train_bert(BS), total=train_steps, disable=BENCHMARK))
    for _ in range(start_step): next(train_it) # Fast forward
-  else:
-    # repeat fake data
-    def repeat_fake(bs):
-      while True: yield get_fake_data_bert(bs)
-    eval_it = iter(repeat_fake(EVAL_BS))
-    train_it = iter(repeat_fake(BS))
+

  step_times = []
  # ** train loop **
  wc_start = time.perf_counter()
-
-  i, train_data = start_step, next(train_it)
-
  if RUNMLPERF:
+    # only load real data with RUNMLPERF
+    i, train_data = start_step, get_data_bert(GPUS, train_it)
    if MLLOGGER:
      MLLOGGER.start(key=mllog_constants.EPOCH_START, value=i*BS, metadata={"epoch_num": i*BS})
+  else:
+    i, train_data = start_step, get_fake_data_bert(GPUS, BS)

  while train_data is not None and i < train_steps and not achieved:
-    if getenv("TRAIN", 1):
-      Tensor.training = True
-      BEAM.value = TRAIN_BEAM
-      st = time.perf_counter()
-      GlobalCounters.reset()
-      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)
+    Tensor.training = True
+    BEAM.value = TRAIN_BEAM
+    st = time.perf_counter()
+    GlobalCounters.reset()
+    loss = 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"])

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

-        try:
-          next_data = next(train_it)
-        except StopIteration:
-          next_data = None
+    try:
+      if RUNMLPERF:
+        next_data = get_data_bert(GPUS, train_it)
+      else:
+        next_data = get_fake_data_bert(GPUS, BS)
+    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()
-        assert not math.isnan(loss)
-        lr = lr.item()
+    device_str = loss.device if isinstance(loss.device, str) else f"{loss.device[0]} * {len(loss.device)}"
+    loss = loss.numpy().item()

-      cl = time.perf_counter()
-      if BENCHMARK: step_times.append(cl - st)
+    cl = time.perf_counter()
+    if BENCHMARK: step_times.append(cl - st)

-      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, {lr:.6f} LR, "
-        f"{GlobalCounters.mem_used / 1e9:.2f} GB used, {GlobalCounters.global_ops * 1e-9 / (cl - st):9.2f} GFLOPS")
-      if WANDB:
-        wandb.log({"lr": lr, "train/loss": loss, "train/global_norm": global_norm.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": (i+1)*BS})
+    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, {optimizer_wd.lr.numpy()[0]:.6f} LR, "
+      f"{GlobalCounters.mem_used / 1e9:.2f} GB used, {GlobalCounters.global_ops * 1e-9 / (cl - st):9.2f} GFLOPS")
+    if WANDB:
+      wandb.log({"lr": optimizer_wd.lr.numpy(), "train/loss": loss, "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": (i+1)*BS})

-      train_data, next_data = next_data, None
-      i += 1
+    train_data, next_data = next_data, None
+    i += 1

-      if i == BENCHMARK:
-        median_step_time = sorted(step_times)[(BENCHMARK + 1) // 2]  # in seconds
-        estimated_total_minutes = int(median_step_time * train_steps / 60)
-        print(f"Estimated training time: {estimated_total_minutes // 60}h{estimated_total_minutes % 60}m")
-        print(f"epoch global_ops: {train_steps * GlobalCounters.global_ops:_}, "
-              f"epoch global_mem: {train_steps * GlobalCounters.global_mem:_}")
+    if i == BENCHMARK:
+      median_step_time = sorted(step_times)[(BENCHMARK + 1) // 2]  # in seconds
+      estimated_total_minutes = int(median_step_time * train_steps / 60)
+      print(f"Estimated training time: {estimated_total_minutes // 60}h{estimated_total_minutes % 60}m")
+      print(f"epoch global_ops: {train_steps * GlobalCounters.global_ops:_}, "
+            f"epoch global_mem: {train_steps * GlobalCounters.global_mem:_}")

    # ** eval loop **
-    if i % eval_step_freq == 0 or (BENCHMARK and i == BENCHMARK) or i == train_steps:
+    if i % eval_step_freq == 0 or (BENCHMARK and i == BENCHMARK):
      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"): train_step_bert.reset()
-      elif getenv("FREE_INTERMEDIATE", 1) and train_step_bert.captured is not None: train_step_bert.captured.free_intermediates()
+      if getenv("RESET_STEP", 1): train_step_bert.reset()
      eval_lm_losses = []
      eval_clsf_losses = []
      eval_lm_accs = []
@@ -1184,14 +816,19 @@ def train_bert():
      BEAM.value = EVAL_BEAM

      for j in tqdm(range(max_eval_steps), desc="Evaluating", total=max_eval_steps, disable=BENCHMARK):
-        eval_data = next(eval_it)
+        if RUNMLPERF:
+          eval_data = get_data_bert(GPUS, eval_it)
+        else:
+          eval_data = get_fake_data_bert(GPUS, EVAL_BS)
        GlobalCounters.reset()
        st = time.time()

-        lm_acc, clsf_acc, lm_loss, clsf_loss = eval_step_bert(model,
+        eval_result: dict[str, Tensor] = eval_step_bert(model,
          eval_data["input_ids"], eval_data["segment_ids"], eval_data["input_mask"], eval_data["masked_lm_positions"],
-          eval_data["masked_lm_ids"], eval_data["masked_lm_weights"], eval_data["next_sentence_labels"], GPUS)
-        lm_acc, clsf_acc, lm_loss, clsf_loss = lm_acc.item(), clsf_acc.item(), lm_loss.item(), clsf_loss.item()
+          eval_data["masked_lm_ids"], eval_data["masked_lm_weights"], eval_data["next_sentence_labels"])
+
+        lm_loss, clsf_loss  = eval_result["masked_lm_loss"].item(), eval_result["next_sentence_loss"].item()
+        lm_acc, clsf_acc = eval_result["masked_lm_accuracy"].item(), eval_result["next_sentence_accuracy"].item()

        eval_lm_losses.append(lm_loss)
        eval_clsf_losses.append(clsf_loss)
@@ -1201,16 +838,14 @@ def train_bert():
        et = time.time()
        eval_times.append(et - st)

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

-      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
+      if getenv("RESET_STEP", 1): eval_step_bert.reset()
+      del eval_data, eval_result
      avg_lm_loss = sum(eval_lm_losses) / len(eval_lm_losses)
      avg_clsf_loss = sum(eval_clsf_losses) / len(eval_clsf_losses)
      avg_lm_acc = sum(eval_lm_accs) / len(eval_lm_accs)
@@ -1222,7 +857,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, "epoch": (i+1)*BS})
+                    "eval/clsf_accuracy": avg_clsf_acc, "eval/forward_time": avg_fw_time})

      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"]})
@@ -1249,14 +884,11 @@ 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"
@@ -1286,4 +918,4 @@ if __name__ == "__main__":
      nm = f"train_{m}"
      if nm in globals():
        print(f"training {m}")
-        with Profiling(enabled=getenv("PYPROFILE")): globals()[nm]()
+        globals()[nm]()
@@ -1,15 +0,0 @@
-#!/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
@@ -1,69 +0,0 @@
-# 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
-```
@@ -1,14 +0,0 @@
-#!/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
@@ -1,17 +0,0 @@
-#!/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
@@ -1,27 +0,0 @@
-#!/bin/bash
-
-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,20 +4,24 @@ This problem uses BERT for NLP.

 ## Requirements

-Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0.
+Install tinygrad and mlperf-logging from master.
 ```
 git clone https://github.com/tinygrad/tinygrad.git
 python3 -m pip install -e ".[mlperf]"
 ```
-Also install gdown (for dataset), numpy, tqdm and tensorflow.
+Also install tqdm and tensorflow.
 ```
-pip install gdown numpy tqdm tensorflow
+pip install 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
@@ -52,18 +56,18 @@ BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-eval

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

 ### tinybox_red

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

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

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

-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 BEAM=4 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=512
 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
+export BENCHMARK=10 DEBUG=2

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

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

-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 BEAM=4 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=512
 export IGNORE_JIT_FIRST_BEAM=1
 export BASEDIR="/raid/datasets/wiki"

@@ -1,13 +1,11 @@
 #!/bin/bash

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

-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 BEAM=4 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=512
 export IGNORE_JIT_FIRST_BEAM=1
 export BASEDIR="/raid/datasets/wiki"

@@ -19,7 +17,7 @@ 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
+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
@@ -4,20 +4,24 @@ This problem uses BERT for NLP.

 ## Requirements

-Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0.
+Install tinygrad and mlperf-logging from master.
 ```
 git clone https://github.com/tinygrad/tinygrad.git
 python3 -m pip install -e ".[mlperf]"
 ```
-Also install gdown (for dataset), numpy, tqdm and tensorflow.
+Also install tqdm and tensorflow.
 ```
-pip install gdown numpy tqdm tensorflow
+pip install 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
@@ -52,18 +56,18 @@ BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-eval

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

 ### tinybox_red

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

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

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

-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 BEAM=3
 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
+export BENCHMARK=10 DEBUG=2

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

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

-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 BEAM=3
 export IGNORE_JIT_FIRST_BEAM=1
 export BASEDIR="/raid/datasets/wiki"

@@ -1,13 +1,11 @@
 #!/bin/bash

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

-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 BEAM=3
 export IGNORE_JIT_FIRST_BEAM=1
 export BASEDIR="/raid/datasets/wiki"

@@ -18,13 +16,8 @@ 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
+BENCHMARK=10 INITMLPERF=1 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,8 @@
+#!/bin/bash
+
+rocm-smi --setprofile compute
+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
@@ -1,10 +1,10 @@
 #!/bin/bash

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

-export RESET_STEP=0
+export LAZYCACHE=0 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

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

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

-export RESET_STEP=0
+export LAZYCACHE=0 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

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

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

-export RESET_STEP=0
+export LAZYCACHE=0 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 LOGMLPERF=1

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

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

-export RESET_STEP=0
+export LAZYCACHE=0 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}
+export BENCHMARK=10 DEBUG=2

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

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

-export RESET_STEP=0
+export LAZYCACHE=0 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

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

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

-export RESET_STEP=0
+export LAZYCACHE=0 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

@@ -1,38 +0,0 @@
-# 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
-```
@@ -1,14 +0,0 @@
-#!/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
@@ -1,15 +0,0 @@
-#!/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
@@ -1,23 +0,0 @@
-#!/bin/bash
-
-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
@@ -1,14 +0,0 @@
-#!/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
@@ -1,15 +0,0 @@
-#!/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
@@ -1,39 +0,0 @@
-{
-    "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 32-Core Processor",
-    "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",
-    "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": ""
-  }
-  
@@ -28,7 +28,7 @@
  "accelerator_interconnect_topology": "",
  "cooling": "air",
  "hw_notes": "",
-  "framework": "tinygrad, branch mlperf_training_v5.0",
+  "framework": "tinygrad, commit b5546912e24e0a864b35924da4efa5d71cfe368b",
  "other_software_stack": {
    "python": "3.10.12",
    "CUDA": "12.4"
@@ -28,9 +28,10 @@
  "accelerator_interconnect_topology": "",
  "cooling": "air",
  "hw_notes": "",
-  "framework": "tinygrad, branch mlperf_training_v5.0",
+  "framework": "tinygrad, commit b5546912e24e0a864b35924da4efa5d71cfe368b",
  "other_software_stack": {
-    "python": "3.10.12"
+    "python": "3.10.12",
+    "ROCm": "6.1.3"
  },
  "operating_system": "Ubuntu 22.04.4",
  "sw_notes": ""
@@ -16,9 +16,9 @@ class LinearGen:
    self.l4 = Tensor.scaled_uniform(1024, 784)

  def forward(self, x):
-    x = x.dot(self.l1).leaky_relu(0.2)
-    x = x.dot(self.l2).leaky_relu(0.2)
-    x = x.dot(self.l3).leaky_relu(0.2)
+    x = x.dot(self.l1).leakyrelu(0.2)
+    x = x.dot(self.l2).leakyrelu(0.2)
+    x = x.dot(self.l3).leakyrelu(0.2)
    x = x.dot(self.l4).tanh()
    return x

@@ -31,9 +31,9 @@ class LinearDisc:

  def forward(self, x):
    # balance the discriminator inputs with const bias (.add(1))
-    x = x.dot(self.l1).add(1).leaky_relu(0.2).dropout(0.3)
-    x = x.dot(self.l2).leaky_relu(0.2).dropout(0.3)
-    x = x.dot(self.l3).leaky_relu(0.2).dropout(0.3)
+    x = x.dot(self.l1).add(1).leakyrelu(0.2).dropout(0.3)
+    x = x.dot(self.l2).leakyrelu(0.2).dropout(0.3)
+    x = x.dot(self.l3).leakyrelu(0.2).dropout(0.3)
    x = x.dot(self.l4).log_softmax()
    return x

@@ -1,94 +0,0 @@
-# https://arxiv.org/pdf/2409.02060
-import time
-import numpy as np
-np.set_printoptions(suppress=True, linewidth=1000)
-import functools
-from tinygrad import Tensor, nn, Device, GlobalCounters
-from tinygrad.helpers import Timing, getenv
-from extra.models.llama import Transformer, convert_from_huggingface
-
-class MixtureFeedForward:
-  def __init__(self, num_experts:int, activated_experts:int, dim:int, hidden_dim:int, linear=nn.Linear):
-    self.activated_experts = activated_experts
-    self.gate = nn.Linear(dim, num_experts, bias=False)
-    self.up_proj = Tensor.zeros(num_experts, hidden_dim, dim, dtype='bfloat16')
-    self.down_proj = Tensor.zeros(num_experts, dim, hidden_dim, dtype='bfloat16')
-    self.gate_proj = Tensor.zeros(num_experts, hidden_dim, dim, dtype='bfloat16')
-  def __call__(self, x:Tensor) -> Tensor:
-    assert x.shape[0] == 1, "only BS=1"
-    assert x.shape[1] == 1, "only length=1"
-    g = self.gate(x).float().softmax(-1)
-
-    g = g.squeeze() # (BS, length, num_experts) -> (num_experts,)
-    probs, sel = g.topk(self.activated_experts)
-
-    # run MoE
-    x_up_gate = x.dot(self.gate_proj[sel].permute(0,2,1)).silu() * x.dot(self.up_proj[sel].permute(0,2,1))
-    x_down = x_up_gate.dot(self.down_proj[sel].permute(0,2,1))
-    return (x_down.float() * probs.reshape(self.activated_experts, 1, 1)).sum(axis=0)
-
-# model is bf16, 1.3B active, 6.9B total
-# M3 Max is 400 GB/s, so 400/2.6 = ~154 tok/s
-
-def fetch_weights() -> dict[str, Tensor]:
-  # TODO: make this lazy so the 3 fetches can happen in parallel
-  m1 = Tensor.from_url("https://huggingface.co/allenai/OLMoE-1B-7B-0924/resolve/main/model-00001-of-00003.safetensors").to(Device.DEFAULT)
-  m2 = Tensor.from_url("https://huggingface.co/allenai/OLMoE-1B-7B-0924/resolve/main/model-00002-of-00003.safetensors").to(Device.DEFAULT)
-  m3 = Tensor.from_url("https://huggingface.co/allenai/OLMoE-1B-7B-0924/resolve/main/model-00003-of-00003.safetensors").to(Device.DEFAULT)
-  return {**nn.state.safe_load(m1), **nn.state.safe_load(m2), **nn.state.safe_load(m3)}
-
-if __name__ == "__main__":
-  if getenv("TORCH"):
-    from transformers import OlmoeForCausalLM, AutoTokenizer
-    model = OlmoeForCausalLM.from_pretrained("allenai/OLMoE-1B-7B-0924")
-    tokenizer = AutoTokenizer.from_pretrained("allenai/OLMoE-1B-7B-0924")
-    inputs = tokenizer("Hello", return_tensors="pt")
-    generate_ids = model.generate(inputs.input_ids, max_length=30)
-    out = tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
-    print(out)
-    exit(0)
-
-  with Timing("create model: "):
-    model = Transformer(n_layers=16, dim=2048, hidden_dim=1024, n_heads=16, norm_eps=1e-5, qk_norm=1e-5, max_context=1024,
-                        vocab_size=50304, feed_forward=functools.partial(MixtureFeedForward, 64, 8))
-    model_state_dict = nn.state.get_state_dict(model)
-    del model_state_dict['freqs_cis']
-
-  with Timing("load weights to GPU: "):
-    nhf_state = convert_from_huggingface(fetch_weights(), 16, 16, 16)
-    # NOTE: i'm not sure this actually needs float32, it may just change the type of things downstream from it. but doesn't match torch w/o this
-    for needs_float32 in ['tok_embeddings.weight']: nhf_state[needs_float32] = nhf_state[needs_float32].float()
-  print(f"ram used: {GlobalCounters.mem_used/1e9:.2f} GB")
-
-  with Timing("unpack weights: "):
-    nn.state.load_state_dict(model, nhf_state, verbose=False, strict=False, consume=True, realize=False)
-    assert len(nhf_state) == 0
-    Tensor.realize(*list(nn.state.get_state_dict(model).values()))
-  print(f"ram used: {GlobalCounters.mem_used/1e9:.2f} GB")
-
-  count = 30
-  temperature = 0
-
-  with Timing("load tokenizer: "):
-    from transformers import AutoTokenizer
-    tokenizer = AutoTokenizer.from_pretrained("allenai/OLMoE-1B-7B-0924")
-
-  toks = [12092]
-  start_pos = 0
-  timings = []
-  for i in range(count):
-    GlobalCounters.reset()
-    st = time.perf_counter()
-    tok = model(Tensor([toks[start_pos:]]), start_pos, temperature).item()
-    timings.append(time.perf_counter()-st)
-    toks.append(tok)
-    start_pos += 1
-    print(toks)
-    print(tokenizer.decode(toks))
-  print(f"fastest token {min(timings)*1e3:.2f} ms, {1/min(timings):.1f} tok/s")
-
-  if temperature == 0:
-    # Hello, I am a newbie to this forum and I am trying to get a better understanding of the different types of data that can be stored in a
-    assert toks == [12092, 13, 309, 717, 247, 747, 17782, 281, 436, 12209, 285, 309, 717, 2820, 281, 755,
-                    247, 1805, 4685, 273, 253, 1027, 3510, 273, 941, 326, 476, 320, 7141, 275, 247], "BAD OUTPUT!"
-
@@ -0,0 +1,211 @@
+#!/usr/bin/env python3
+import os, sys, io, pathlib, json, struct
+import numpy as np
+sys.path.insert(0, str(pathlib.Path(__file__).parents[1]))
+
+if "FLOAT16" not in os.environ: os.environ["FLOAT16"] = "1"
+if "IMAGE" not in os.environ: os.environ["IMAGE"] = "2"
+if "NOLOCALS" not in os.environ: os.environ["NOLOCALS"] = "1"
+
+OPENPILOT_MODEL = "https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx"
+
+import onnx
+from typing import Tuple, List, Optional, Dict, cast
+from extra.onnx import get_run_onnx
+from tinygrad import Tensor, Device, GlobalCounters, dtypes
+from tinygrad.dtype import ImageDType
+from tinygrad.device import Buffer
+from tinygrad.helpers import partition, Context, fetch, getenv, DEBUG, tqdm
+from tinygrad.engine.realize import run_schedule, lower_schedule, ExecItem, CompiledRunner
+from tinygrad.engine.memory import memory_planner
+from tinygrad.engine.schedule import ScheduleItem, create_schedule
+from tinygrad.ops import Ops
+from tinygrad.tensor import _to_np_dtype
+Device.DEFAULT = "GPU"
+
+def get_schedule(onnx_data) -> Tuple[List[ScheduleItem], List[ScheduleItem]]:
+  Tensor.no_grad = True
+  Tensor.training = False
+
+  # load the model
+  onnx_model = onnx.load(io.BytesIO(onnx_data))
+  run_onnx = get_run_onnx(onnx_model)
+  input_shapes = {inp.name:tuple(x.dim_value for x in inp.type.tensor_type.shape.dim) for inp in onnx_model.graph.input}
+
+  # run the model
+  inputs = {k:Tensor.empty(*shp) for k,shp in input_shapes.items()}
+  ret: Tensor = next(iter(run_onnx(inputs).values())).cast(dtypes.float32).contiguous()
+  schedule = create_schedule([ret.lazydata])
+
+  # filter schedule that don't depend on the inputs
+  input_lb = [x.lazydata.base.buffer for x in inputs.values()]
+  depends = set(input_lb)
+  for si in schedule:
+    if any(b in depends for b in si.inputs):
+      for out in si.outputs: depends.add(out)
+
+  # run all kernels that don't depend on the inputs
+  # NOTE: there's two extra kernels due to fusions that now happen since the weights aren't realized
+  schedule, schedule_independent = partition(schedule, lambda si: any(out in depends for out in si.outputs))
+  print(f"{len(schedule)} schedule items depend on the input, {len(schedule_independent)} don't")
+
+  # confirm no non-sink metaop in the (non independent) schedule except for the ones that load the input buffers
+  assert all(si.ast.op is Ops.SINK or out in input_lb for si in schedule for out in si.outputs), "has non SINK ops, can't compile to Thneed"
+  return schedule, schedule_independent, inputs
+
+def test_vs_onnx(onnx_data, eis:Optional[List[ExecItem]], inputs:Dict[str, Tensor]):
+  import onnx
+  #import pyopencl as cl
+  #from extra.thneed import Thneed
+  import numpy as np
+  onnx_model = onnx.load(io.BytesIO(onnx_data))
+
+  input_shapes = {inp.name:tuple(x.dim_value for x in inp.type.tensor_type.shape.dim) for inp in onnx_model.graph.input}
+  Tensor.manual_seed(1337)
+  new_inputs = {k:Tensor.randn(*shp, requires_grad=False)*8 for k,shp in input_shapes.items()}
+  new_np_inputs = {k:v.realize().numpy() for k,v in new_inputs.items()}
+
+  if getenv("ORT"):
+    # test with onnxruntime
+    import onnxruntime as ort
+    onnx_session = ort.InferenceSession(onnx_data)
+    onnx_output = onnx_session.run([onnx_model.graph.output[0].name], {k:v.astype(np.float16) for k,v in new_np_inputs.items()})
+    new_torch_out = onnx_output[0]
+    print("got ort outputs")
+  else:
+    # test with torch
+    from test.models.test_onnx import run_onnx_torch
+    new_torch_out = run_onnx_torch(onnx_model, new_np_inputs).numpy()
+    print("got torch outputs")
+
+  # if you don't have a schedule
+  if eis is None:
+    run_onnx = get_run_onnx(onnx_model)
+    new_tinygrad_out = next(iter(run_onnx(new_inputs).values())).cast(dtypes.float32).numpy()
+    np.testing.assert_allclose(new_torch_out, new_tinygrad_out, atol=1e-4, rtol=1e-2)
+    print("classic self-test passed!")
+    return
+
+  # set inputs
+  for k,v in inputs.items(): v.lazydata.base.realized.copyin(new_np_inputs[k].data)
+
+  # run code (all buffers have been allocated)
+  GlobalCounters.reset()
+  output = eis[-1].bufs[0]
+  for ei in eis: ei.run()
+
+  new_tinygrad_out = np.frombuffer(output.as_buffer(), dtype=_to_np_dtype(output.dtype))
+  np.testing.assert_allclose(new_torch_out.reshape(new_tinygrad_out.shape), new_tinygrad_out, atol=1e-4, rtol=1e-2)
+  print("semi-thneed self-test passed!")
+
+if __name__ == "__main__":
+  onnx_data = fetch(sys.argv[1] if len(sys.argv) > 1 else OPENPILOT_MODEL).read_bytes()
+
+  # quick test for ONNX issues
+  #thneed_test_onnx(onnx_data, None)
+  #exit(0)
+
+  schedule, schedule_independent, inputs = get_schedule(onnx_data)
+  schedule, schedule_input = partition(schedule, lambda x: x.ast.op is Ops.SINK)
+  print(f"{len(schedule_input)} inputs")
+
+  run_schedule(schedule_independent)
+  run_schedule(schedule_input)
+  with Context(DEBUG=max(DEBUG.value, 2), BEAM=getenv("LATEBEAM")):
+    schedule = memory_planner(schedule)
+    for si in schedule:
+      for b in si.outputs:
+        assert not b.is_allocated(), "output should not be allocated"
+    image_count = sum(isinstance(out.dtype, ImageDType) for si in schedule for out in si.outputs)
+    print(f"**** compiling real kernels {image_count}/{len(schedule)} images ****")
+    eis = list(tqdm(lower_schedule(schedule), total=len(schedule)))
+
+  print("kernel count:", len(eis))
+  assert len(eis) <= getenv("ALLOWED_KERNEL_COUNT", 0) or getenv("ALLOWED_KERNEL_COUNT", 0) == 0, "too many kernels!"
+
+  # new simple thneed
+  def to_ref(b:Buffer): return struct.pack("Q", id(b)).decode("latin_1")
+
+  seen_buffers = set()
+  input_buffers = [x.lazydata.buffer for x in inputs.values()]
+  jdat = {"binaries": [], "programs": {}, "kernels": [], "objects": []}
+  jdat["inputs"] = {k:to_ref(v.lazydata.buffer) for k,v in inputs.items()}
+  jdat["outputs"] = [to_ref(eis[-1].bufs[0])]
+  weights = []
+  for i,ei in enumerate(eis):
+    #print("***", i)
+    for b in ei.bufs:
+      needs_load = b.is_allocated() and b not in input_buffers
+      #print(b, needs_load)
+      if b in seen_buffers: continue
+      seen_buffers.add(b)
+      if isinstance(b.dtype, ImageDType):
+        base_dtype = dtypes.float16 if b.dtype.fmt == 'e' else dtypes.float32
+        row_pitch = (b.dtype.shape[0]*4*base_dtype.itemsize + 63)//64 * 64
+        size = row_pitch * b.dtype.shape[1]
+        jdat['objects'].append({
+          "id": to_ref(b), "needs_load": needs_load, "size": size, "arg_type": "image2d_t",
+          "width": b.dtype.shape[0], "height": b.dtype.shape[1], "row_pitch": row_pitch, "float32": b.dtype.base == dtypes.float32,
+        })
+        if needs_load:
+          t = Tensor.empty(b.dtype.shape, dtype=b.dtype)
+          t.lazydata.buffer = b
+          data = t.cast(dtypes.float32).pad(((0, row_pitch//(4*base_dtype.itemsize)-b.dtype.shape[0]), (0,0), (0,0))).contiguous().numpy()
+          # NOTE: this cast must be done in numpy for platforms that don't support half
+          if base_dtype == dtypes.float16: data = data.astype(np.float16)
+          weights.append(data.tobytes())
+          assert len(weights[-1]) == size, "wrong size buffer"
+      else:
+        jdat['objects'].append({
+          "id": to_ref(b), "arg_type": b.dtype.name + "*", "needs_load": needs_load, "size": b.nbytes,
+        })
+        if needs_load:
+          weights.append(b.as_buffer())
+          assert len(weights[-1]) == b.nbytes, "wrong size buffer"
+
+  saved_binaries = set()
+  binaries = []
+  gated_read_image_count = 0
+  GlobalCounters.reset()
+  with Context(DEBUG=max(DEBUG.value, 2)):
+    for ei in eis:
+      prg = cast(CompiledRunner, ei.prg)
+      assert len(prg.p.vars) == 0
+      if prg.p.function_name not in saved_binaries:
+        jdat['binaries'].append({"name":prg.p.function_name, "length":len(prg.lib)})
+        binaries.append(prg.lib)
+        saved_binaries.add(prg.p.function_name)
+      gated_read_image_count += prg.p.src.count("?read_image")
+      ei.run()
+      jdat['kernels'].append({
+        "name": prg.p.function_name,
+        "work_dim": len(prg.p.global_size),
+        "global_work_size": prg.p.global_size,
+        "local_work_size": prg.p.local_size,
+        "num_args": len(ei.bufs),
+        "args": [to_ref(b) for b in ei.bufs],
+        "arg_size": [8]*len(ei.bufs),
+      })
+
+  if (allowed_gated_read_image:=getenv("ALLOWED_GATED_READ_IMAGE", -1)) != -1:
+    assert gated_read_image_count <= allowed_gated_read_image, \
+      f"too many gated read_image! {gated_read_image_count=}, {allowed_gated_read_image=}"
+
+  output_fn = sys.argv[2] if len(sys.argv) >= 3 else "/tmp/output.thneed"
+  print(f"saving thneed to {output_fn} with {len(weights)} buffers and {len(binaries)} binaries")
+  with open(output_fn, "wb") as f:
+    j = json.dumps(jdat, ensure_ascii=False).encode('latin_1')
+    f.write(struct.pack("I", len(j)))
+    f.write(j)
+    for w in weights: f.write(w)
+    for b in binaries: f.write(b)
+    print("saved", f.tell(), "bytes")
+
+  FLOAT16 = getenv("FLOAT16", 0)
+  if FLOAT16 == 0:
+    try:
+      test_vs_onnx(onnx_data, eis, inputs)
+    except ModuleNotFoundError as e:
+      print(f"TEST NOT HAPPENING {e}")
+
+
@@ -5,150 +5,76 @@ if "IMAGE" not in os.environ: os.environ["IMAGE"] = "2"
 if "NOLOCALS" not in os.environ: os.environ["NOLOCALS"] = "1"
 if "JIT_BATCH_SIZE" not in os.environ: os.environ["JIT_BATCH_SIZE"] = "0"

-from tinygrad import fetch, Tensor, TinyJit, Context, GlobalCounters, Device
+from tinygrad import fetch, Tensor, TinyJit, Context, GlobalCounters
 from tinygrad.helpers import DEBUG, getenv
 from tinygrad.tensor import _from_np_dtype
-from tinygrad.engine.realize import CompiledRunner

 import onnx
 from onnx.helper import tensor_dtype_to_np_dtype
-from tinygrad.frontend.onnx import OnnxRunner
+from extra.onnx import get_run_onnx   # TODO: port to main tinygrad

 OPENPILOT_MODEL = sys.argv[1] if len(sys.argv) > 1 else "https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx"
-OUTPUT = sys.argv[2] if len(sys.argv) > 2 else "/tmp/openpilot.pkl"
+OUTPUT = "/tmp/openpilot.pkl"

-def compile(onnx_file):
-  onnx_model = onnx.load(onnx_file)
-  run_onnx = OnnxRunner(onnx_model)
+def compile():
+  Tensor.no_grad = True
+  Tensor.training = False
+
+  onnx_bytes = fetch(OPENPILOT_MODEL)
+  onnx_model = onnx.load(onnx_bytes)
+  run_onnx = get_run_onnx(onnx_model)
  print("loaded model")

  input_shapes = {inp.name:tuple(x.dim_value for x in inp.type.tensor_type.shape.dim) for inp in onnx_model.graph.input}
  input_types = {inp.name: tensor_dtype_to_np_dtype(inp.type.tensor_type.elem_type) for inp in onnx_model.graph.input}
-  # Float inputs and outputs to tinyjits for openpilot are always float32
-  input_types = {k:(np.float32 if v==np.float16 else v) for k,v in input_types.items()}
+  if getenv("FLOAT16", 0) == 0: input_types = {k:(np.float32 if v==np.float16 else v) for k,v in input_types.items()}
  Tensor.manual_seed(100)
  new_inputs = {k:Tensor.randn(*shp, dtype=_from_np_dtype(input_types[k])).mul(8).realize() for k,shp in sorted(input_shapes.items())}
-  new_inputs_numpy = {k:v.numpy() for k,v in new_inputs.items()}
  print("created tensors")

-  run_onnx_jit = TinyJit(lambda **kwargs:
-                         next(iter(run_onnx({k:v.to(Device.DEFAULT) for k,v in kwargs.items()}).values())).cast('float32'), prune=True)
+  run_onnx_jit = TinyJit(lambda **kwargs: run_onnx(kwargs), prune=True)
  for i in range(3):
    GlobalCounters.reset()
    print(f"run {i}")
-    inputs = {**{k:v.clone() for k,v in new_inputs.items() if 'img' in k},
-              **{k:Tensor(v, device="NPY").realize() for k,v in new_inputs_numpy.items() if 'img' not in k}}
    with Context(DEBUG=max(DEBUG.value, 2 if i == 2 else 1)):
-      ret = run_onnx_jit(**inputs).numpy()
+      ret = next(iter(run_onnx_jit(**new_inputs).values())).cast('float32').numpy()
    # copy i == 1 so use of JITBEAM is okay
    if i == 1: test_val = np.copy(ret)
  print(f"captured {len(run_onnx_jit.captured.jit_cache)} kernels")
-  np.testing.assert_equal(test_val, ret, "JIT run failed")
+  np.testing.assert_equal(test_val, ret)
  print("jit run validated")

-  # checks from compile2
-  kernel_count = 0
-  read_image_count = 0
-  gated_read_image_count = 0
-  for ei in run_onnx_jit.captured.jit_cache:
-    if isinstance(ei.prg, CompiledRunner):
-      kernel_count += 1
-      read_image_count += ei.prg.p.src.count("read_image")
-      gated_read_image_count += ei.prg.p.src.count("?read_image")
-  print(f"{kernel_count=},  {read_image_count=}, {gated_read_image_count=}")
-  if (allowed_kernel_count:=getenv("ALLOWED_KERNEL_COUNT", -1)) != -1:
-    assert kernel_count <= allowed_kernel_count, f"too many kernels! {kernel_count=}, {allowed_kernel_count=}"
-  if (allowed_read_image:=getenv("ALLOWED_READ_IMAGE", -1)) != -1:
-    assert read_image_count == allowed_read_image, f"different read_image! {read_image_count=}, {allowed_read_image=}"
-  if (allowed_gated_read_image:=getenv("ALLOWED_GATED_READ_IMAGE", -1)) != -1:
-    assert gated_read_image_count <= allowed_gated_read_image, f"too many gated read_image! {gated_read_image_count=}, {allowed_gated_read_image=}"
-
  with open(OUTPUT, "wb") as f:
    pickle.dump(run_onnx_jit, f)
-  mdl_sz = os.path.getsize(onnx_file)
+  mdl_sz = os.path.getsize(onnx_bytes)
  pkl_sz = os.path.getsize(OUTPUT)
  print(f"mdl size is {mdl_sz/1e6:.2f}M")
  print(f"pkl size is {pkl_sz/1e6:.2f}M")
  print("**** compile done ****")
  return test_val

-def test_vs_compile(run, new_inputs, test_val=None):
+def test(test_val=None):
+  with open(OUTPUT, "rb") as f:
+    run = pickle.load(f)
+  Tensor.manual_seed(100)
+  new_inputs = {nm:Tensor.randn(*st.shape, dtype=dtype).mul(8).realize() for nm, (st, _, dtype, _) in
+                sorted(zip(run.captured.expected_names, run.captured.expected_st_vars_dtype_device))}
  new_inputs_numpy = {k:v.numpy() for k,v in new_inputs.items()}
-
-  # create fake "from_blob" tensors for the inputs, and wrapped NPY tensors for the numpy inputs (these have the same underlying memory)
-  inputs = {**{k:v for k,v in new_inputs.items() if 'img' in k},
-            **{k:Tensor(v, device="NPY").realize() for k,v in new_inputs_numpy.items() if 'img' not in k}}
-
-  # run 20 times
  for _ in range(20):
    st = time.perf_counter()
+    # Need to cast non-image inputs from numpy, this is only realistic way to run it
+    inputs = {**{k:v for k,v in new_inputs.items() if 'img' in k},
+              **{k:Tensor(v) for k,v in new_inputs_numpy.items() if 'img' not in k}}
    out = run(**inputs)
    mt = time.perf_counter()
-    val = out.numpy()
+    val = out['outputs'].numpy()
    et = time.perf_counter()
    print(f"enqueue {(mt-st)*1e3:6.2f} ms -- total run {(et-st)*1e3:6.2f} ms")
  print(out, val.shape, val.dtype)
  if test_val is not None: np.testing.assert_equal(test_val, val)
  print("**** test done ****")

-  # test that changing the numpy changes the model outputs
-  if any([x.device == 'NPY' for x in inputs.values()]):
-    for v in new_inputs_numpy.values(): v *= 2
-    out = run(**inputs)
-    changed_val = out.numpy()
-    np.testing.assert_raises(AssertionError, np.testing.assert_array_equal, val, changed_val)
-  return val
-
-def test_vs_onnx(new_inputs, test_val, onnx_file, ort=False):
-  new_inputs_numpy = {k:v.numpy() for k,v in new_inputs.items()}
-  onnx_model = onnx.load(onnx_file)
-
-  timings = []
-  if ort:
-    # test with onnxruntime
-    import onnxruntime as ort
-    onnx_session = ort.InferenceSession(onnx_file)
-    for _ in range(1 if test_val is not None else 5):
-      st = time.perf_counter()
-      onnx_output = onnx_session.run([onnx_model.graph.output[0].name], {k:v.astype(np.float16) for k,v in new_inputs_numpy.items()})
-      timings.append(time.perf_counter() - st)
-    new_torch_out = onnx_output[0]
-  else:
-    # test with torch
-    import torch
-    from onnx2torch import convert
-    inputs = {k.name:new_inputs_numpy[k.name] for k in onnx_model.graph.input}
-    torch_model = convert(onnx_model).float()
-    with torch.no_grad():
-      for _ in range(1 if test_val is not None else 5):
-        st = time.perf_counter()
-        torch_out = torch_model(*[torch.tensor(x) for x in inputs.values()])
-        timings.append(time.perf_counter() - st)
-      new_torch_out = torch_out.numpy()
-
-  if test_val is not None:
-    np.testing.assert_allclose(new_torch_out.reshape(test_val.shape), test_val, atol=1e-4, rtol=1e-2)
-    print("test vs onnx passed")
-  return timings
-
 if __name__ == "__main__":
-  onnx_file = fetch(OPENPILOT_MODEL)
-  test_val = compile(onnx_file) if not getenv("RUN") else None
-
-  with open(OUTPUT, "rb") as f: pickle_loaded = pickle.load(f)
-
-  # same randomness as compile
-  Tensor.manual_seed(100)
-  new_inputs = {nm:Tensor.randn(*st.shape, dtype=dtype).mul(8).realize() for nm, (st, _, dtype, _) in
-                sorted(zip(pickle_loaded.captured.expected_names, pickle_loaded.captured.expected_st_vars_dtype_device))}
-
-  test_val = test_vs_compile(pickle_loaded, new_inputs, test_val)
-  if getenv("BENCHMARK"):
-    for be in ["torch", "ort"]:
-      try:
-        timings = test_vs_onnx(new_inputs, None, onnx_file, be=="ort")
-        print(f"timing {be}: {min(timings)*1000:.2f} ms")
-      except Exception as e:
-        print(f"{be} fail with {e}")
-  if not getenv("FLOAT16"): test_vs_onnx(new_inputs, test_val, onnx_file, getenv("ORT"))
+  test_val = compile() if not getenv("RUN") else None
+  test(test_val)

@@ -0,0 +1,2 @@
+#!/bin/bash
+NOLOCALS=1 FLOAT16=1 DEBUGCL=1 IMAGE=2 GPU=1 python3 examples/openpilot/compile2.py
@@ -1,5 +1,5 @@
-from tinygrad import dtypes, getenv, Device
-from tinygrad.helpers import trange, colored, DEBUG, temp
+from tinygrad import dtypes
+from tinygrad.helpers import trange
 from tinygrad.nn.datasets import mnist
 import torch
 from torch import nn, optim
@@ -26,20 +26,14 @@ class Model(nn.Module):
    return self.lin(torch.flatten(x, 1))

 if __name__ == "__main__":
-  if getenv("TINY_BACKEND"):
-    import tinygrad.frontend.torch
-    device = torch.device("tiny")
-  else:
-    device = torch.device({"METAL":"mps","NV":"cuda"}.get(Device.DEFAULT, "cpu"))
-  if DEBUG >= 1: print(f"using torch backend {device}")
+  mps_device = torch.device("mps")
  X_train, Y_train, X_test, Y_test = mnist()
-  X_train = torch.tensor(X_train.float().numpy(), device=device)
-  Y_train = torch.tensor(Y_train.cast(dtypes.int64).numpy(), device=device)
-  X_test = torch.tensor(X_test.float().numpy(), device=device)
-  Y_test = torch.tensor(Y_test.cast(dtypes.int64).numpy(), device=device)
+  X_train = torch.tensor(X_train.float().numpy(), device=mps_device)
+  Y_train = torch.tensor(Y_train.cast(dtypes.int64).numpy(), device=mps_device)
+  X_test = torch.tensor(X_test.float().numpy(), device=mps_device)
+  Y_test = torch.tensor(Y_test.cast(dtypes.int64).numpy(), device=mps_device)

-  if getenv("TORCHVIZ"): torch.cuda.memory._record_memory_history()
-  model = Model().to(device)
+  model = Model().to(mps_device)
  optimizer = optim.Adam(model.parameters(), 1e-3)

  loss_fn = nn.CrossEntropyLoss()
@@ -54,16 +48,8 @@ if __name__ == "__main__":
    return loss

  test_acc = float('nan')
-  for i in (t:=trange(getenv("STEPS", 70))):
+  for i in (t:=trange(70)):
    samples = torch.randint(0, X_train.shape[0], (512,))  # putting this in JIT didn't work well
    loss = step(samples)
    if i%10 == 9: test_acc = ((model(X_test).argmax(axis=-1) == Y_test).sum() * 100 / X_test.shape[0]).item()
    t.set_description(f"loss: {loss.item():6.2f} test_accuracy: {test_acc:5.2f}%")
-
-  # verify eval acc
-  if target := getenv("TARGET_EVAL_ACC_PCT", 0.0):
-    if test_acc >= target and test_acc != 100.0: print(colored(f"{test_acc=} >= {target}", "green"))
-    else: raise ValueError(colored(f"{test_acc=} < {target}", "red"))
-  if getenv("TORCHVIZ"):
-    torch.cuda.memory._dump_snapshot(fp:=temp("torchviz.pkl", append_user=True))
-    print(f"saved torch memory snapshot to {fp}, view in https://pytorch.org/memory_viz")
@@ -44,7 +44,7 @@ def load_model(model_path:Path, model_params:Dict[str, Union[int, float]]) -> Tr
  model.layers = updated_layers

  # load weights
-  weights = fix_bf16(convert_from_huggingface(load(str(model_path / "model.safetensors.index.json")), model_params["n_layers"], model_params["n_heads"], model_params["n_kv_heads"], permute_layers=False))
+  weights = fix_bf16(convert_from_huggingface(load(str(model_path / "model.safetensors.index.json")), model, model_params["n_heads"], model_params["n_kv_heads"], permute_layers=False))

  # replace weights in model
  load_state_dict(model, weights, strict=False, consume=True)
@@ -5,7 +5,6 @@ from examples.stable_diffusion import AutoencoderKL, get_alphas_cumprod
 from examples.sdxl import DPMPP2MSampler, append_dims, LegacyDDPMDiscretization
 from extra.models.unet import UNetModel
 from extra.models.clip import FrozenOpenClipEmbedder
-from extra.bench_log import BenchEvent, WallTimeEvent

 from typing import Dict
 import argparse, tempfile, os
@@ -118,14 +117,12 @@ if __name__ == "__main__":
  if not weights_fn:
    weights_url = args.weights_url if args.weights_url else default_weights_url
    weights_fn  = fetch(weights_url, os.path.basename(str(weights_url)))
+  load_state_dict(model, safe_load(weights_fn), strict=False)

-  with WallTimeEvent(BenchEvent.LOAD_WEIGHTS):
-    load_state_dict(model, safe_load(weights_fn), strict=False)
-
-    if args.fp16:
-      for k,v in get_state_dict(model).items():
-        if k.startswith("model"):
-          v.replace(v.cast(dtypes.float16).realize())
+  if args.fp16:
+    for k,v in get_state_dict(model).items():
+      if k.startswith("model"):
+        v.replace(v.cast(dtypes.float16).realize())

  c  = { "crossattn": model.cond_stage_model(args.prompt) }
  uc = { "crossattn": model.cond_stage_model("") }
@@ -3,13 +3,12 @@
 # Stability-AI/generative-models | MIT     | https://github.com/Stability-AI/generative-models/blob/fbdc58cab9f4ee2be7a5e1f2e2787ecd9311942f/LICENSE-CODE
 # mlfoundations/open_clip        | MIT     | https://github.com/mlfoundations/open_clip/blob/58e4e39aaabc6040839b0d2a7e8bf20979e4558a/LICENSE

-from tinygrad import Tensor, TinyJit, dtypes, GlobalCounters
+from tinygrad import Tensor, TinyJit, dtypes
 from tinygrad.nn import Conv2d, GroupNorm
-from tinygrad.nn.state import safe_load, load_state_dict, get_state_dict
-from tinygrad.helpers import fetch, trange, colored, Timing
+from tinygrad.nn.state import safe_load, load_state_dict
+from tinygrad.helpers import fetch, trange, colored, Timing, GlobalCounters
 from extra.models.clip import Embedder, FrozenClosedClipEmbedder, FrozenOpenClipEmbedder
 from extra.models.unet import UNetModel, Upsample, Downsample, timestep_embedding
-from extra.bench_log import BenchEvent, WallTimeEvent
 from examples.stable_diffusion import ResnetBlock, Mid
 import numpy as np

@@ -346,19 +345,18 @@ class DPMPP2MSampler:
    old_denoised = None
    for i in trange(num_sigmas - 1):
      with Timing("step in ", enabled=timing, on_exit=lambda _: f", using {GlobalCounters.mem_used/1e9:.2f} GB"):
-        GlobalCounters.reset()
-        with WallTimeEvent(BenchEvent.STEP):
-          x, old_denoised = self.sampler_step(
-            old_denoised=old_denoised,
-            prev_sigma=(None if i==0 else sigmas[i-1].expand(x.shape[0])),
-            sigma=sigmas[i].expand(x.shape[0]),
-            next_sigma=sigmas[i+1].expand(x.shape[0]),
-            denoiser=denoiser,
-            x=x,
-            c=c,
-            uc=uc,
-          )
-          x.realize(old_denoised)
+        x, old_denoised = self.sampler_step(
+          old_denoised=old_denoised,
+          prev_sigma=(None if i==0 else sigmas[i-1].expand(x.shape[0])),
+          sigma=sigmas[i].expand(x.shape[0]),
+          next_sigma=sigmas[i+1].expand(x.shape[0]),
+          denoiser=denoiser,
+          x=x,
+          c=c,
+          uc=uc,
+        )
+        x.realize()
+        old_denoised.realize()

    return x

@@ -386,13 +384,7 @@ if __name__ == "__main__":

  default_weight_url = 'https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/resolve/main/sd_xl_base_1.0.safetensors'
  weights = args.weights if args.weights else fetch(default_weight_url, 'sd_xl_base_1.0.safetensors')
-  loaded_weights = load_state_dict(model, safe_load(weights), strict=False, verbose=False, realize=False)
-
-  start_mem_used = GlobalCounters.mem_used
-  with Timing("loaded weights in ", lambda et_ns: f", {(B:=(GlobalCounters.mem_used-start_mem_used))/1e9:.2f} GB loaded at {B/et_ns:.2f} GB/s"):
-    with WallTimeEvent(BenchEvent.LOAD_WEIGHTS):
-      Tensor.realize(*loaded_weights)
-    del loaded_weights
+  load_state_dict(model, safe_load(weights), strict=False)

  N = 1
  C = 4
@@ -403,7 +395,8 @@ if __name__ == "__main__":

  c, uc = model.create_conditioning([args.prompt], args.width, args.height)
  del model.conditioner
-  Tensor.realize(*c.values(), *uc.values())
+  for v in c .values(): v.realize()
+  for v in uc.values(): v.realize()
  print("created batch")

  # https://github.com/Stability-AI/generative-models/blob/fbdc58cab9f4ee2be7a5e1f2e2787ecd9311942f/sgm/inference/helpers.py#L101
@@ -1,48 +0,0 @@
-import os, pathlib, argparse
-from examples.llama3 import Tokenizer
-from tabulate import tabulate
-from tinygrad import fetch
-from tinygrad.helpers import flatten
-
-# llama 3 tokenizer
-tokenizer = Tokenizer(fetch("https://huggingface.co/bofenghuang/Meta-Llama-3-8B/resolve/main/original/tokenizer.model").as_posix())
-
-def read_code(base_path):
-  ret = []
-  for path, _, files in os.walk(os.path.join(base_path, "tinygrad")):
-    for name in files:
-      if not name.endswith(".py"): continue
-      if 'tinygrad/runtime/autogen' in path.replace('\\', '/'): continue
-      fullpath = os.path.join(path, name)
-      code = pathlib.Path(fullpath).read_text()
-      ret.append(("### " + fullpath.split("tinygrad/", 1)[1], code))
-  return ret
-
-def write_code_to_file(filename, code_list):
-  """Writes the combined code to a specified file."""
-  with open(filename, 'w') as f:
-    f.write('\n'.join(flatten(code_list)))
-
-if __name__ == "__main__":
-  parser = argparse.ArgumentParser(description="Analyze and optionally save tinygrad code.")
-  parser.add_argument("--output", help="Output file to write the combined code to.")
-  args = parser.parse_args()
-
-  ret = read_code(".")
-
-  table = []
-  for name,code in ret:
-    table.append([name, len(tokenizer.encode(name+"\x00"+code))])
-  print(tabulate([["name", "llm tokens"]]+sorted(table, key=lambda x: -x[1]), headers="firstrow"))
-
-  code_str = '\x00'.join(flatten(ret))
-  print(f"code has {len(code_str)} chars")
-  newline_count = code_str.count('\n')
-  print(f"code has {newline_count} newlines")
-
-  encoded = tokenizer.encode(code_str)
-  print(f"code has {len(encoded)} tokens")
-
-  if args.output:
-    write_code_to_file(args.output, ret)
-    print(f"Combined code written to {args.output}")
@@ -437,14 +437,14 @@ class Generator:
    x = self.conv_pre(x)
    if g is not None:  x = x + self.cond(g)
    for i in range(self.num_upsamples):
-      x, xs = self.ups[i](x.leaky_relu(LRELU_SLOPE)), None
+      x, xs = self.ups[i](x.leakyrelu(LRELU_SLOPE)), None
      x_source = self.noise_convs[i](har_source)
      x = x + x_source
      for j in range(self.num_kernels):
        if xs is None: xs = self.resblocks[i * self.num_kernels + j].forward(x)
        else: xs += self.resblocks[i * self.num_kernels + j].forward(x)
      x = xs / self.num_kernels
-    return self.conv_post(x.leaky_relu()).tanh()
+    return self.conv_post(x.leakyrelu()).tanh()

 # **** helpers ****

@@ -504,7 +504,7 @@ def load_checkpoint_enc(checkpoint_path, model: ContentVec, optimizer=None, skip
        obj, v = getattr(parent, "weight"), weight_norm(weight_v, weight_g, 0)
        weight_g, weight_v, parent, skip = None, None, None, False
      if not skip and obj.shape == v.shape:
-        if "feature_extractor" in key and (isinstance(parent, (nn.GroupNorm, nn.LayerNorm))):  # cast
+        if "feature_extractor" in key and (isinstance(parent, nn.GroupNorm) or isinstance(parent, nn.LayerNorm)):  # cast
          obj.assign(v.to(obj.device).float())
        else:
          obj.assign(v.to(obj.device))
@@ -14,7 +14,6 @@ from tinygrad.nn import Conv2d, GroupNorm
 from tinygrad.nn.state import torch_load, load_state_dict, get_state_dict
 from extra.models.clip import Closed, Tokenizer
 from extra.models.unet import UNetModel
-from extra.bench_log import BenchEvent, WallTimeEvent

 class AttnBlock:
  def __init__(self, in_channels):
@@ -233,13 +232,12 @@ if __name__ == "__main__":
  model = StableDiffusion()

  # load in weights
-  with WallTimeEvent(BenchEvent.LOAD_WEIGHTS):
-    load_state_dict(model, torch_load(fetch('https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt', 'sd-v1-4.ckpt'))['state_dict'], strict=False)
+  load_state_dict(model, torch_load(fetch('https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt', 'sd-v1-4.ckpt'))['state_dict'], strict=False)

-    if args.fp16:
-      for k,v in get_state_dict(model).items():
-        if k.startswith("model"):
-          v.replace(v.cast(dtypes.float16).realize())
+  if args.fp16:
+    for k,v in get_state_dict(model).items():
+      if k.startswith("model"):
+        v.replace(v.cast(dtypes.float16).realize())

  # run through CLIP to get context
  tokenizer = Tokenizer.ClipTokenizer()
@@ -272,10 +270,9 @@ if __name__ == "__main__":
      GlobalCounters.reset()
      t.set_description("%3d %3d" % (index, timestep))
      with Timing("step in ", enabled=args.timing, on_exit=lambda _: f", using {GlobalCounters.mem_used/1e9:.2f} GB"):
-        with WallTimeEvent(BenchEvent.STEP):
-          tid = Tensor([index])
-          latent = run(model, unconditional_context, context, latent, Tensor([timestep]), alphas[tid], alphas_prev[tid], Tensor([args.guidance]))
-          if args.timing: Device[Device.DEFAULT].synchronize()
+        tid = Tensor([index])
+        latent = run(model, unconditional_context, context, latent, Tensor([timestep]), alphas[tid], alphas_prev[tid], Tensor([args.guidance]))
+        if args.timing: Device[Device.DEFAULT].synchronize()
    del run

  # upsample latent space to image with autoencoder
@@ -5,13 +5,10 @@
 #  - symbolic removal

 from examples.beautiful_mnist import Model
-from tinygrad import Tensor, nn, getenv, GlobalCounters, Variable
+from tinygrad import Tensor, nn, getenv, GlobalCounters
 from tinygrad.nn.datasets import mnist
 from tinygrad.helpers import trange, DEBUG

-# STEPS=70 python3 examples/stunning_mnist.py
-# NOTE: it's broken with STACK=1, why?
-
 if __name__ == "__main__":
  X_train, Y_train, X_test, Y_test = mnist()
  print("*** got data")
@@ -27,21 +24,19 @@ if __name__ == "__main__":
  print("*** got samples")

  with Tensor.train():
+    # TODO: this shouldn't be a for loop. something like: (contract is still up in the air)
    """
    i = UOp.range(samples.shape[0])  # TODO: fix range function on UOp
    losses = model(X_samp[i]).sparse_categorical_crossentropy(Y_samp[i]).backward().contract(i)
    opt.schedule_steps(i)
    """
-    # TODO: this shouldn't be a for loop. something like: (contract is still up in the air)
-    vi = Variable('i', 0, samples.shape[0]-1)
    losses = []
    for i in range(samples.shape[0]):
-      vib = vi.bind(i)
      opt.zero_grad()
-      losses.append(model(X_samp[vib]).sparse_categorical_crossentropy(Y_samp[vib]).backward())
+      losses.append(model(X_samp[i]).sparse_categorical_crossentropy(Y_samp[i]).backward())
      opt.schedule_step()
    # TODO: this stack currently breaks the "generator" aspect of losses. it probably shouldn't
-    if getenv("STACK", 0): losses = Tensor.stack(*losses)
+    #losses = Tensor.stack(*losses)
  print("*** scheduled training")

  # evaluate the model
@@ -54,8 +49,5 @@ if __name__ == "__main__":

  # only actually do anything at the end
  if getenv("LOSS", 1):
-    for i in (t:=trange(len(losses))):
-      GlobalCounters.reset()
-      t.set_description(f"loss: {losses[i].item():6.2f}")
-  if getenv("TEST", 1):
-    print(f"test_accuracy: {test_acc.item():5.2f}%")
+    for i in (t:=trange(len(losses))): t.set_description(f"loss: {losses[i].item():6.2f}")
+  print(f"test_accuracy: {test_acc.item():5.2f}%")
@@ -0,0 +1,5 @@
+from tinygrad import Tensor, dtypes
+
+if __name__ == "__main__":
+    a = Tensor([1.0,2.0,3.0], dtype=dtypes.half)
+    print((a*2.0).numpy())
@@ -1,81 +0,0 @@
-import random, sys
-import numpy as np
-from extra.datasets.imagenet import get_imagenet_categories, get_val_files, center_crop
-from examples.benchmark_onnx import load_onnx_model
-from PIL import Image
-from tinygrad import Tensor, dtypes, GlobalCounters
-from tinygrad.helpers import fetch, getenv
-
-# works:
-#  ~70% - https://github.com/onnx/models/raw/refs/heads/main/validated/vision/classification/resnet/model/resnet50-v2-7.onnx
-#  ~43% - https://github.com/onnx/models/raw/refs/heads/main/Computer_Vision/alexnet_Opset16_torch_hub/alexnet_Opset16.onnx
-#  ~72% - https://github.com/xamcat/mobcat-samples/raw/refs/heads/master/onnx_runtime/InferencingSample/InferencingSample/mobilenetv2-7.onnx
-#  ~71% - https://github.com/axinc-ai/onnx-quantization/raw/refs/heads/main/models/mobilenetv2_1.0.opt.onnx
-#  ~67% - https://github.com/xamcat/mobcat-samples/raw/refs/heads/master/onnx_runtime/InferencingSample/InferencingSample/mobilenetv2-7-quantized.onnx
-# broken:
-#  https://github.com/MTlab/onnx2caffe/raw/refs/heads/master/model/MobileNetV2.onnx
-#  https://huggingface.co/qualcomm/MobileNet-v2-Quantized/resolve/main/MobileNet-v2-Quantized.onnx
-#  ~35% - https://github.com/axinc-ai/onnx-quantization/raw/refs/heads/main/models/mobilenev2_quantized.onnx
-
-# QUANT=1 python3 examples/test_onnx_imagenet.py
-#   https://github.com/xamcat/mobcat-samples/raw/refs/heads/master/onnx_runtime/InferencingSample/InferencingSample/mobilenetv2-7.onnx
-# DONT_REALIZE_EXPAND=1 python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
-# VIZ=1 DONT_REALIZE_EXPAND=1 python3 examples/benchmark_onnx.py /tmp/model.quant.onnx
-
-def imagenet_dataloader(cnt=0):
-  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)
-  files = get_val_files()
-  random.shuffle(files)
-  if cnt != 0: files = files[:cnt]
-  cir = get_imagenet_categories()
-  for fn in files:
-    img = Image.open(fn)
-    img = img.convert('RGB') if img.mode != "RGB" else img
-    img = center_crop(img)
-    img = np.array(img)
-    img = Tensor(img).permute(2,0,1).reshape(1,3,224,224)
-    img = ((img.cast(dtypes.float32)/255.0) - input_mean) / input_std
-    y = cir[fn.split("/")[-2]]
-    yield img,y
-
-if __name__ == "__main__":
-  fn = sys.argv[1]
-  if getenv("QUANT"):
-    from onnxruntime.quantization import quantize_dynamic, quantize_static, QuantFormat, QuantType, CalibrationDataReader
-    model_fp32 = fetch(fn)
-    fn = '/tmp/model.quant.onnx'
-    if getenv("DYNAMIC"):
-      quantize_dynamic(model_fp32, fn)
-    else:
-      class ImagenetReader(CalibrationDataReader):
-        def __init__(self):
-          self.iter = imagenet_dataloader(cnt=1000)
-        def get_next(self) -> dict:
-          try:
-            img,y = next(self.iter)
-          except StopIteration:
-            return None
-          return {"input": img.numpy()}
-      quantize_static(model_fp32, fn, ImagenetReader(), quant_format=QuantFormat.QDQ, per_channel=False,
-                      activation_type=QuantType.QUInt8, weight_type=QuantType.QUInt8,
-                      extra_options={"ActivationSymmetric": False})
-
-  run_onnx_jit, input_specs = load_onnx_model(fetch(fn))
-  t_name, t_spec = list(input_specs.items())[0]
-  assert t_spec.shape[1:] == (3,224,224), f"shape is {t_spec.shape}"
-
-  hit = 0
-  for i,(img,y) in enumerate(imagenet_dataloader(cnt=getenv("CNT", 100))):
-    GlobalCounters.reset()
-    p = run_onnx_jit(**{t_name:img})
-    assert p.shape == (1,1000)
-    t = p.to('cpu').argmax().item()
-    hit += y==t
-    print(f"target: {y:3d}  pred: {t:3d}  acc: {hit/(i+1)*100:.2f}%")
-
-  MS_TARGET = 13.4
-  print(f"need {GlobalCounters.global_ops/1e9*(1000/MS_TARGET):.2f} GFLOPS for {MS_TARGET:.2f} ms")
-
-  import pickle
-  with open("/tmp/im.pkl", "wb") as f: pickle.dump(run_onnx_jit, f)
@@ -1,19 +0,0 @@
-import sys, pickle
-from tinygrad import GlobalCounters
-from tinygrad.helpers import fetch, getenv
-from examples.test_onnx_imagenet import imagenet_dataloader
-
-if __name__ == "__main__":
-  with open(fetch(sys.argv[1]), "rb") as f:
-    run_onnx_jit = pickle.load(f)
-  input_name = run_onnx_jit.captured.expected_names[0]
-  device = run_onnx_jit.captured.expected_st_vars_dtype_device[0][-1]
-  print(f"input goes into {input_name=} on {device=}")
-  hit = 0
-  for i,(img,y) in enumerate(imagenet_dataloader(cnt=getenv("CNT", 100))):
-    GlobalCounters.reset()
-    p = run_onnx_jit(**{input_name:img.to(device)})
-    assert p.shape == (1,1000)
-    t = p.to('cpu').argmax().item()
-    hit += y==t
-    print(f"target: {y:3d}  pred: {t:3d}  acc: {hit/(i+1)*100:.2f}%")
@@ -1,11 +0,0 @@
-/*!
-Pure v3.0.0
-Copyright 2013 Yahoo!
-Licensed under the BSD License.
-https://github.com/pure-css/pure/blob/master/LICENSE
-*/
-/*!
-normalize.css v | MIT License | https://necolas.github.io/normalize.css/
-Copyright (c) Nicolas Gallagher and Jonathan Neal
-*/
-/*! normalize.css v8.0.1 | MIT License | github.com/necolas/normalize.css */html{line-height:1.15;-webkit-text-size-adjust:100%}body{margin:0}main{display:block}h1{font-size:2em;margin:.67em 0}hr{box-sizing:content-box;height:0;overflow:visible}pre{font-family:monospace,monospace;font-size:1em}a{background-color:transparent}abbr[title]{border-bottom:none;text-decoration:underline;-webkit-text-decoration:underline dotted;text-decoration:underline dotted}b,strong{font-weight:bolder}code,kbd,samp{font-family:monospace,monospace;font-size:1em}small{font-size:80%}sub,sup{font-size:75%;line-height:0;position:relative;vertical-align:baseline}sub{bottom:-.25em}sup{top:-.5em}img{border-style:none}button,input,optgroup,select,textarea{font-family:inherit;font-size:100%;line-height:1.15;margin:0}button,input{overflow:visible}button,select{text-transform:none}[type=button],[type=reset],[type=submit],button{-webkit-appearance:button}[type=button]::-moz-focus-inner,[type=reset]::-moz-focus-inner,[type=submit]::-moz-focus-inner,button::-moz-focus-inner{border-style:none;padding:0}[type=button]:-moz-focusring,[type=reset]:-moz-focusring,[type=submit]:-moz-focusring,button:-moz-focusring{outline:1px dotted ButtonText}fieldset{padding:.35em .75em .625em}legend{box-sizing:border-box;color:inherit;display:table;max-width:100%;padding:0;white-space:normal}progress{vertical-align:baseline}textarea{overflow:auto}[type=checkbox],[type=radio]{box-sizing:border-box;padding:0}[type=number]::-webkit-inner-spin-button,[type=number]::-webkit-outer-spin-button{height:auto}[type=search]{-webkit-appearance:textfield;outline-offset:-2px}[type=search]::-webkit-search-decoration{-webkit-appearance:none}::-webkit-file-upload-button{-webkit-appearance:button;font:inherit}details{display:block}summary{display:list-item}template{display:none}[hidden]{display:none}html{font-family:sans-serif}.hidden,[hidden]{display:none!important}.pure-img{max-width:100%;height:auto;display:block}
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
Ahmed Harmouche	fc9ae53a05	Progress with pack/unpack wgsl built-ins	2024-12-05 16:12:21 +01:00
Ahmed Harmouche	4669e467ef	f16 comp/decomp WIP	2024-12-05 14:14:23 +01:00