mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-08-14 19:38:27 +00:00
Compare commits
2
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 George HotzandGitHub
|
a07c9da26b | ||
 geohot
|
816a359a3c |
@@ -11,5 +11,5 @@ runs:
|
||||
git fetch origin $CURRENT_SHA
|
||||
export COMMIT_MESSAGE=$(git show -s --format=%B "$CURRENT_SHA")
|
||||
export CURRENT_HEAD=$(git rev-parse HEAD)
|
||||
cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && CHECK_OOB=0 PYTHONPATH=. python3 process_replay.py
|
||||
cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && IGNORE_OOB=1 PYTHONPATH=. python3 process_replay.py
|
||||
git checkout $CURRENT_HEAD # restore to branch
|
||||
|
||||
@@ -45,55 +45,43 @@ inputs:
|
||||
description: "Install mesa"
|
||||
required: false
|
||||
default: 'false'
|
||||
tinydreno:
|
||||
description: "Install tinydreno"
|
||||
required: false
|
||||
default: 'false'
|
||||
runs:
|
||||
using: "composite"
|
||||
steps:
|
||||
- name: Set up Python ${{ inputs.python-version }}
|
||||
id: setup-python
|
||||
uses: actions/setup-python@v6
|
||||
uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: ${{ inputs.python-version }}
|
||||
|
||||
# **** Caching packages ****
|
||||
|
||||
- name: Cache Python packages (PR)
|
||||
if: github.event_name == 'pull_request'
|
||||
id: restore-venv-pr
|
||||
uses: actions/cache/restore@v4
|
||||
with:
|
||||
path: ${{ github.workspace }}/.venv
|
||||
key: venv-${{ runner.os }}-${{ runner.arch }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ env.CACHE_VERSION }}
|
||||
- name: Cache Python packages
|
||||
if: github.event_name != 'pull_request'
|
||||
id: restore-venv
|
||||
uses: actions/cache@v5
|
||||
uses: actions/cache@v4
|
||||
with:
|
||||
path: ${{ github.workspace }}/.venv
|
||||
key: venv-${{ runner.os }}-${{ runner.arch }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ env.CACHE_VERSION }}
|
||||
key: venv-${{ runner.os }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ env.CACHE_VERSION }}
|
||||
|
||||
# **** Caching downloads ****
|
||||
|
||||
- name: Cache downloads (PR)
|
||||
if: inputs.key != '' && github.event_name == 'pull_request'
|
||||
uses: actions/cache/restore@v4
|
||||
- name: Cache downloads (Linux)
|
||||
if: inputs.key != '' && runner.os == 'Linux'
|
||||
uses: actions/cache@v4
|
||||
with:
|
||||
path: ${{ runner.os == 'Linux' && '~/.cache/tinygrad/downloads/' || '~/Library/Caches/tinygrad/downloads/' }}
|
||||
path: ~/.cache/tinygrad/downloads/
|
||||
key: downloads-${{ github.job }}-${{ inputs.key }}-${{ env.CACHE_VERSION }}
|
||||
- name: Cache downloads
|
||||
if: inputs.key != '' && github.event_name != 'pull_request'
|
||||
uses: actions/cache@v5
|
||||
- name: Cache downloads (macOS)
|
||||
if: inputs.key != '' && runner.os == 'macOS'
|
||||
uses: actions/cache@v4
|
||||
with:
|
||||
path: ${{ runner.os == 'Linux' && '~/.cache/tinygrad/downloads/' || '~/Library/Caches/tinygrad/downloads/' }}
|
||||
path: ~/Library/Caches/tinygrad/downloads/
|
||||
key: downloads-${{ github.job }}-${{ inputs.key }}-${{ env.CACHE_VERSION }}
|
||||
|
||||
# **** Python deps ****
|
||||
|
||||
- name: Install dependencies in venv (with extra)
|
||||
if: inputs.deps != '' && steps.restore-venv-pr.outputs.cache-hit != 'true' && steps.restore-venv.outputs.cache-hit != 'true'
|
||||
if: inputs.deps != '' && steps.restore-venv.outputs.cache-hit != 'true'
|
||||
shell: bash
|
||||
run: |
|
||||
python -m venv .venv
|
||||
@@ -104,7 +92,7 @@ runs:
|
||||
fi
|
||||
python -m pip install -e ".[${{ inputs.deps }}]" ${{ inputs.pydeps }} --extra-index-url https://download.pytorch.org/whl/cpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/Triton-Nightly/pypi/simple/
|
||||
- name: Install dependencies in venv (without extra)
|
||||
if: inputs.deps == '' && steps.restore-venv-pr.outputs.cache-hit != 'true' && steps.restore-venv.outputs.cache-hit != 'true'
|
||||
if: inputs.deps == '' && steps.restore-venv.outputs.cache-hit != 'true'
|
||||
shell: bash
|
||||
run: |
|
||||
python -m venv .venv
|
||||
@@ -149,7 +137,7 @@ runs:
|
||||
run: |
|
||||
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/7.1 $(lsb_release -cs) main
|
||||
deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/6.2 $(lsb_release -cs) main
|
||||
EOF
|
||||
echo -e 'Package: *\nPin: release o=repo.radeon.com\nPin-Priority: 600' | sudo tee /etc/apt/preferences.d/rocm-pin-600
|
||||
|
||||
@@ -194,18 +182,12 @@ runs:
|
||||
echo "pkgs=$pkgs" >> "$GITHUB_OUTPUT"
|
||||
echo "hash=$(echo -n "$pkgs" | sha256sum | cut -d' ' -f1)" >> "$GITHUB_OUTPUT"
|
||||
|
||||
- name: Cache apt (PR)
|
||||
if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true') && github.event_name == 'pull_request'
|
||||
uses: actions/cache/restore@v4
|
||||
with:
|
||||
path: /var/cache/apt/archives/
|
||||
key: ${{ runner.os }}-${{ runner.arch }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
|
||||
- name: Cache apt
|
||||
if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true') && github.event_name != 'pull_request'
|
||||
uses: actions/cache@v5
|
||||
if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true')
|
||||
uses: actions/cache@v4
|
||||
with:
|
||||
path: /var/cache/apt/archives/
|
||||
key: ${{ runner.os }}-${{ runner.arch }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
|
||||
key: ${{ runner.os }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
|
||||
|
||||
- name: Run apt Update + Install
|
||||
if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true')
|
||||
@@ -237,7 +219,7 @@ runs:
|
||||
shell: bash
|
||||
run: |
|
||||
sudo mkdir -p /usr/local/lib
|
||||
curl -s -H "Authorization: token $GH_TOKEN" curl -s https://api.github.com/repos/tinygrad/amdcomgr_dylib/releases/latest | \
|
||||
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 -fL -o /usr/local/lib/libamd_comgr.dylib
|
||||
cargo build --release --manifest-path ./extra/remu/Cargo.toml
|
||||
@@ -257,26 +239,17 @@ runs:
|
||||
ln -s /opt/homebrew/opt/[email protected] /opt/homebrew/opt/boost || true
|
||||
ln -s /opt/homebrew/opt/boost/lib/libboost_atomic-mt.dylib /opt/homebrew/opt/boost/lib/libboost_atomic.dylib || true
|
||||
ln -s /opt/homebrew/opt/boost/lib/libboost_thread-mt.dylib /opt/homebrew/opt/boost/lib/libboost_thread.dylib || true
|
||||
- name: Cache gpuocelot (PR)
|
||||
if: inputs.ocelot == 'true' && github.event_name == 'pull_request'
|
||||
id: cache-build-pr
|
||||
uses: actions/cache/restore@v4
|
||||
env:
|
||||
cache-name: cache-gpuocelot-build-1
|
||||
with:
|
||||
path: ${{ github.workspace }}/gpuocelot/ocelot
|
||||
key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-${{ env.CACHE_VERSION }}
|
||||
- name: Cache gpuocelot
|
||||
if: inputs.ocelot == 'true' && github.event_name != 'pull_request'
|
||||
if: inputs.ocelot == 'true'
|
||||
id: cache-build
|
||||
uses: actions/cache@v5
|
||||
uses: actions/cache@v4
|
||||
env:
|
||||
cache-name: cache-gpuocelot-build-1
|
||||
with:
|
||||
path: ${{ github.workspace }}/gpuocelot/ocelot
|
||||
key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-${{ env.CACHE_VERSION }}
|
||||
- name: Clone/compile gpuocelot
|
||||
if: inputs.ocelot == 'true' && steps.cache-build-pr.outputs.cache-hit != 'true' && steps.cache-build.outputs.cache-hit != 'true'
|
||||
if: inputs.ocelot == 'true' && steps.cache-build.outputs.cache-hit != 'true'
|
||||
shell: bash
|
||||
run: |
|
||||
git clone --recurse-submodules https://github.com/gpuocelot/gpuocelot.git ${{ github.workspace }}/gpuocelot
|
||||
@@ -287,7 +260,6 @@ runs:
|
||||
|
||||
CMAKE_ARGS="-Wno-dev -G Ninja -DOCELOT_BUILD_TOOLS=OFF -DCMAKE_BUILD_ALWAYS=0 -DBUILD_TESTS_CUDA=OFF -DCMAKE_POLICY_VERSION_MINIMUM=3.5"
|
||||
if [[ "${{ runner.os }}" == "macOS" ]]; then
|
||||
sudo xcode-select -s /Applications/Xcode_16.2.app/Contents/Developer
|
||||
CMAKE_ARGS="$CMAKE_ARGS -DBoost_INCLUDE_DIR=$(brew --prefix boost)/include -DBoost_LIBRARY_DIR=$(brew --prefix boost)/lib"
|
||||
fi
|
||||
|
||||
@@ -331,9 +303,3 @@ runs:
|
||||
if: inputs.mesa == 'true' && runner.os == 'macOS'
|
||||
shell: bash
|
||||
run: brew install sirhcm/tinymesa/tinymesa_cpu
|
||||
|
||||
# *** tinydreno ***
|
||||
- name: Install tinydreno (linux)
|
||||
if: inputs.tinydreno == 'true' && runner.os == 'Linux'
|
||||
shell: bash
|
||||
run: sudo curl -fL https://github.com/sirhcm/tinydreno/raw/refs/heads/master/libllvm-qcom.so -o /usr/lib/libllvm-qcom.so
|
||||
|
||||
+108
-76
@@ -14,12 +14,10 @@ on:
|
||||
paths:
|
||||
- 'tinygrad/runtime/autogen/**/*'
|
||||
- 'tinygrad/runtime/support/autogen.py'
|
||||
- '.github/workflows/autogen.yml'
|
||||
workflow_dispatch:
|
||||
paths:
|
||||
- 'tinygrad/runtime/autogen/**/*'
|
||||
- 'tinygrad/runtime/support/autogen.py'
|
||||
- '.github/workflows/autogen.yml'
|
||||
|
||||
jobs:
|
||||
autogen:
|
||||
@@ -28,11 +26,10 @@ jobs:
|
||||
timeout-minutes: 15
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
uses: actions/checkout@v6
|
||||
uses: actions/checkout@v4
|
||||
- name: Setup Environment
|
||||
uses: ./.github/actions/setup-tinygrad
|
||||
with:
|
||||
key: 'autogen'
|
||||
opencl: 'true'
|
||||
amd: 'true'
|
||||
cuda: 'true'
|
||||
@@ -41,104 +38,139 @@ jobs:
|
||||
mesa: 'true'
|
||||
pydeps: 'pyyaml mako'
|
||||
- name: Install autogen support packages
|
||||
run: sudo apt-get install -y --no-install-recommends libclang-20-dev llvm-20-dev hip-dev libusb-1.0-0-dev libdrm-dev
|
||||
- name: Regenerate autogen files
|
||||
run: sudo apt-get install -y --no-install-recommends libclang-20-dev llvm-20-dev hip-dev libusb-1.0-0-dev
|
||||
- name: Verify OpenCL autogen
|
||||
run: |
|
||||
find tinygrad/runtime/autogen -type f -name "*.py" -not -path "*/amd/*" -not -name "__init__.py" -not -name "comgr.py" -not -name "metal.py" -not -name "iokit.py" -not -name "corefoundation.py" -not -name "libclang.py" -delete
|
||||
mv tinygrad/runtime/autogen/opencl.py /tmp/opencl.py.bak
|
||||
python3 -c "from tinygrad.runtime.autogen import opencl"
|
||||
python3 -c "from tinygrad.runtime.autogen import cuda, nvrtc, nvjitlink, nv_570, nv_580, nv"
|
||||
python3 -c "from tinygrad.runtime.autogen import comgr_3, hsa, hip, amd_gpu, sqtt, rocprof, amdgpu_kd, amdgpu_drm"
|
||||
python3 -c "from tinygrad.runtime.autogen.am import am, pm4_soc15, pm4_nv, sdma_4_0_0, sdma_5_0_0, sdma_6_0_0, smu_v13_0_0, smu_v13_0_6, smu_v13_0_12, smu_v14_0_2"
|
||||
python3 -c "from tinygrad.runtime.autogen import libc, kfd, io_uring, ib, pci, vfio"
|
||||
python3 -c "from tinygrad.runtime.autogen import llvm"
|
||||
python3 -c "from tinygrad.runtime.autogen import webgpu"
|
||||
python3 -c "from tinygrad.runtime.autogen import kgsl, qcom_dsp"
|
||||
python3 -c "from tinygrad.runtime.autogen import libusb"
|
||||
python3 -c "from tinygrad.runtime.autogen import mesa"
|
||||
python3 -c "from tinygrad.runtime.autogen import avcodec"
|
||||
python3 -c "from tinygrad.runtime.autogen import llvm_qcom"
|
||||
REGEN=1 python3 -c "from tinygrad.runtime.autogen import libclang"
|
||||
- name: Check for differences
|
||||
diff /tmp/opencl.py.bak tinygrad/runtime/autogen/opencl.py
|
||||
- name: Verify CUDA autogen
|
||||
run: |
|
||||
if ! git diff --quiet; then
|
||||
git diff
|
||||
git diff > autogen-ubuntu.patch
|
||||
echo "Autogen mismatch detected. Patch available at: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}#artifacts"
|
||||
exit 1
|
||||
fi
|
||||
- name: Upload patch artifact
|
||||
if: failure()
|
||||
uses: actions/upload-artifact@v7
|
||||
with:
|
||||
name: autogen-ubuntu-patch
|
||||
path: autogen-ubuntu.patch
|
||||
|
||||
mv tinygrad/runtime/autogen/cuda.py /tmp/cuda.py.bak
|
||||
mv tinygrad/runtime/autogen/nvrtc.py /tmp/nvrtc.py.bak
|
||||
mv tinygrad/runtime/autogen/nvjitlink.py /tmp/nvjitlink.py.bak
|
||||
mv tinygrad/runtime/autogen/nv_570.py /tmp/nv_570.py.bak
|
||||
mv tinygrad/runtime/autogen/nv.py /tmp/nv.py.bak
|
||||
python3 -c "from tinygrad.runtime.autogen import cuda, nvrtc, nvjitlink, nv_570, nv"
|
||||
diff /tmp/cuda.py.bak tinygrad/runtime/autogen/cuda.py
|
||||
diff /tmp/nvrtc.py.bak tinygrad/runtime/autogen/nvrtc.py
|
||||
diff /tmp/nvjitlink.py.bak tinygrad/runtime/autogen/nvjitlink.py
|
||||
diff /tmp/nv_570.py.bak tinygrad/runtime/autogen/nv_570.py
|
||||
diff /tmp/nv.py.bak tinygrad/runtime/autogen/nv.py
|
||||
- name: Verify AMD autogen
|
||||
run: |
|
||||
mv tinygrad/runtime/autogen/comgr.py /tmp/comgr.py.bak
|
||||
mv tinygrad/runtime/autogen/hsa.py /tmp/hsa.py.bak
|
||||
mv tinygrad/runtime/autogen/hip.py /tmp/hip.py.bak
|
||||
mv tinygrad/runtime/autogen/amd_gpu.py /tmp/amd_gpu.py.bak
|
||||
mv tinygrad/runtime/autogen/sqtt.py /tmp/sqtt.py.bak
|
||||
mv tinygrad/runtime/autogen/rocprof.py /tmp/rocprof.py.bak
|
||||
mv tinygrad/runtime/autogen/am/am.py /tmp/am_am.py.bak
|
||||
mv tinygrad/runtime/autogen/am/pm4_soc15.py /tmp/am_pm4_soc15.py.bak
|
||||
mv tinygrad/runtime/autogen/am/pm4_nv.py /tmp/am_pm4_nv.py.bak
|
||||
mv tinygrad/runtime/autogen/am/sdma_4_0_0.py /tmp/am_sdma_4_0_0.py.bak
|
||||
mv tinygrad/runtime/autogen/am/sdma_5_0_0.py /tmp/am_sdma_5_0_0.py.bak
|
||||
mv tinygrad/runtime/autogen/am/sdma_6_0_0.py /tmp/am_sdma_6_0_0.py.bak
|
||||
mv tinygrad/runtime/autogen/am/smu_v13_0_0.py /tmp/am_smu_v13_0_0.py.bak
|
||||
mv tinygrad/runtime/autogen/am/smu_v14_0_2.py /tmp/am_smu_v14_0_2.py.bak
|
||||
python3 -c "from tinygrad.runtime.autogen import comgr, hsa, hip, amd_gpu, sqtt, rocprof; from tinygrad.runtime.autogen.am import am, pm4_soc15, pm4_nv, sdma_4_0_0, sdma_5_0_0, sdma_6_0_0, smu_v13_0_0, smu_v14_0_2"
|
||||
diff /tmp/comgr.py.bak tinygrad/runtime/autogen/comgr.py
|
||||
diff /tmp/hsa.py.bak tinygrad/runtime/autogen/hsa.py
|
||||
diff /tmp/hip.py.bak tinygrad/runtime/autogen/hip.py
|
||||
diff /tmp/amd_gpu.py.bak tinygrad/runtime/autogen/amd_gpu.py
|
||||
diff /tmp/sqtt.py.bak tinygrad/runtime/autogen/sqtt.py
|
||||
diff /tmp/rocprof.py.bak tinygrad/runtime/autogen/rocprof.py
|
||||
diff /tmp/am_am.py.bak tinygrad/runtime/autogen/am/am.py
|
||||
diff /tmp/am_pm4_soc15.py.bak tinygrad/runtime/autogen/am/pm4_soc15.py
|
||||
diff /tmp/am_pm4_nv.py.bak tinygrad/runtime/autogen/am/pm4_nv.py
|
||||
diff /tmp/am_sdma_4_0_0.py.bak tinygrad/runtime/autogen/am/sdma_4_0_0.py
|
||||
diff /tmp/am_sdma_5_0_0.py.bak tinygrad/runtime/autogen/am/sdma_5_0_0.py
|
||||
diff /tmp/am_sdma_6_0_0.py.bak tinygrad/runtime/autogen/am/sdma_6_0_0.py
|
||||
diff /tmp/am_smu_v13_0_0.py.bak tinygrad/runtime/autogen/am/smu_v13_0_0.py
|
||||
diff /tmp/am_smu_v14_0_2.py.bak tinygrad/runtime/autogen/am/smu_v14_0_2.py
|
||||
- name: Verify Linux autogen
|
||||
run: |
|
||||
mv tinygrad/runtime/autogen/libc.py /tmp/libc.py.bak
|
||||
mv tinygrad/runtime/autogen/kfd.py /tmp/kfd.py.bak
|
||||
mv tinygrad/runtime/autogen/io_uring.py /tmp/io_uring.py.bak
|
||||
mv tinygrad/runtime/autogen/ib.py /tmp/ib.py.bak
|
||||
mv tinygrad/runtime/autogen/pci.py /tmp/pci.py.bak
|
||||
mv tinygrad/runtime/autogen/vfio.py /tmp/vfio.py.bak
|
||||
python3 -c "from tinygrad.runtime.autogen import libc, kfd, io_uring, ib, pci, vfio"
|
||||
diff /tmp/libc.py.bak tinygrad/runtime/autogen/libc.py
|
||||
diff /tmp/kfd.py.bak tinygrad/runtime/autogen/kfd.py
|
||||
diff /tmp/io_uring.py.bak tinygrad/runtime/autogen/io_uring.py
|
||||
diff /tmp/ib.py.bak tinygrad/runtime/autogen/ib.py
|
||||
diff /tmp/pci.py.bak tinygrad/runtime/autogen/pci.py
|
||||
diff /tmp/vfio.py.bak tinygrad/runtime/autogen/vfio.py
|
||||
- name: Verify LLVM autogen
|
||||
run: |
|
||||
mv tinygrad/runtime/autogen/llvm.py /tmp/llvm.py.bak
|
||||
python3 -c "from tinygrad.runtime.autogen import llvm"
|
||||
diff /tmp/llvm.py.bak tinygrad/runtime/autogen/llvm.py
|
||||
- name: Verify WebGPU autogen
|
||||
run: |
|
||||
mv tinygrad/runtime/autogen/webgpu.py /tmp/webgpu.py.bak
|
||||
python3 -c "from tinygrad.runtime.autogen import webgpu"
|
||||
diff /tmp/webgpu.py.bak tinygrad/runtime/autogen/webgpu.py
|
||||
- name: Verify Qualcomm autogen
|
||||
run: |
|
||||
mv tinygrad/runtime/autogen/kgsl.py /tmp/kgsl.py.bak
|
||||
mv tinygrad/runtime/autogen/qcom_dsp.py /tmp/qcom_dsp.py.bak
|
||||
python3 -c "from tinygrad.runtime.autogen import kgsl, qcom_dsp"
|
||||
diff /tmp/kgsl.py.bak tinygrad/runtime/autogen/kgsl.py
|
||||
diff /tmp/qcom_dsp.py.bak tinygrad/runtime/autogen/qcom_dsp.py
|
||||
- name: Verify libusb autogen
|
||||
run: |
|
||||
mv tinygrad/runtime/autogen/libusb.py /tmp/libusb.py.bak
|
||||
python3 -c "from tinygrad.runtime.autogen import libusb"
|
||||
diff /tmp/libusb.py.bak tinygrad/runtime/autogen/libusb.py
|
||||
- name: Verify mesa autogen
|
||||
run: |
|
||||
mv tinygrad/runtime/autogen/mesa.py /tmp/mesa.py.bak
|
||||
python3 -c "from tinygrad.runtime.autogen import mesa"
|
||||
diff /tmp/mesa.py.bak tinygrad/runtime/autogen/mesa.py
|
||||
- name: Verify libclang autogen
|
||||
run: |
|
||||
cp tinygrad/runtime/autogen/libclang.py /tmp/libclang.py.bak
|
||||
REGEN=1 python3 -c "from tinygrad.runtime.autogen import libclang"
|
||||
diff /tmp/libclang.py.bak tinygrad/runtime/autogen/libclang.py
|
||||
autogen-mac:
|
||||
name: In-tree Autogen (macos)
|
||||
runs-on: macos-14
|
||||
timeout-minutes: 15
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
uses: actions/checkout@v6
|
||||
uses: actions/checkout@v4
|
||||
- name: Setup Environment
|
||||
uses: ./.github/actions/setup-tinygrad
|
||||
with:
|
||||
key: 'autogen-mac'
|
||||
llvm: 'true'
|
||||
- name: Regenerate autogen files
|
||||
- name: Verify macos autogen
|
||||
run: |
|
||||
rm tinygrad/runtime/autogen/metal.py tinygrad/runtime/autogen/iokit.py tinygrad/runtime/autogen/corefoundation.py
|
||||
python3 -c "from tinygrad.runtime.autogen import metal, iokit, corefoundation"
|
||||
- name: Check for differences
|
||||
run: |
|
||||
if ! git diff --quiet; then
|
||||
git diff
|
||||
git diff > autogen-macos.patch
|
||||
echo "Autogen mismatch detected. Patch available at: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}#artifacts"
|
||||
exit 1
|
||||
fi
|
||||
- name: Upload patch artifact
|
||||
if: failure()
|
||||
uses: actions/upload-artifact@v7
|
||||
with:
|
||||
name: autogen-macos-patch
|
||||
path: autogen-macos.patch
|
||||
|
||||
autogen-comgr-2:
|
||||
name: In-tree Autogen (comgr 2)
|
||||
mv tinygrad/runtime/autogen/metal.py /tmp/metal.py.bak
|
||||
LIBCLANG_PATH=/opt/homebrew/opt/llvm@20/lib/libclang.dylib python3 -c "from tinygrad.runtime.autogen import metal"
|
||||
diff /tmp/metal.py.bak tinygrad/runtime/autogen/metal.py
|
||||
autogen-comgr-3:
|
||||
name: In-tree Autogen (comgr 3)
|
||||
runs-on: ubuntu-24.04
|
||||
timeout-minutes: 15
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
uses: actions/checkout@v6
|
||||
uses: actions/checkout@v4
|
||||
- name: Setup Environment
|
||||
uses: ./.github/actions/setup-tinygrad
|
||||
with:
|
||||
key: 'autogen-comgr'
|
||||
- name: Install autogen support packages
|
||||
run: |
|
||||
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.2 $(lsb_release -cs) main
|
||||
deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/6.4 $(lsb_release -cs) main
|
||||
EOF
|
||||
echo -e 'Package: *\nPin: release o=repo.radeon.com\nPin-Priority: 600' | sudo tee /etc/apt/preferences.d/rocm-pin-600
|
||||
sudo apt -qq update || true
|
||||
sudo apt-get install -y --no-install-recommends libclang-20-dev comgr
|
||||
- name: Regenerate autogen files
|
||||
- name: Verify comgr (3) autogen
|
||||
run: |
|
||||
rm tinygrad/runtime/autogen/comgr.py
|
||||
python3 -c "from tinygrad.runtime.autogen import comgr"
|
||||
- name: Check for differences
|
||||
run: |
|
||||
if ! git diff --quiet; then
|
||||
git diff
|
||||
git diff > autogen-comgr2.patch
|
||||
echo "Autogen mismatch detected. Patch available at: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}#artifacts"
|
||||
exit 1
|
||||
fi
|
||||
- name: Upload patch artifact
|
||||
if: failure()
|
||||
uses: actions/upload-artifact@v7
|
||||
with:
|
||||
name: autogen-comgr2-patch
|
||||
path: autogen-comgr2.patch
|
||||
mv tinygrad/runtime/autogen/comgr_3.py /tmp/comgr_3.py.bak
|
||||
python3 -c "from tinygrad.runtime.autogen import comgr_3"
|
||||
diff /tmp/comgr_3.py.bak tinygrad/runtime/autogen/comgr_3.py
|
||||
|
||||
+302
-278
@@ -16,41 +16,6 @@ on:
|
||||
workflow_dispatch:
|
||||
|
||||
jobs:
|
||||
# the goal of this test is to replicate a normal person on a laptop running the test
|
||||
# no process replay, no benchmarks, no CI, just a normal laptop person
|
||||
# the 3 minute timeout should not be raised
|
||||
testmacpytest:
|
||||
name: Mac pytest
|
||||
env:
|
||||
CI: ""
|
||||
CAPTURE_PROCESS_REPLAY: "0"
|
||||
runs-on: [self-hosted, macOS]
|
||||
timeout-minutes: 3
|
||||
defaults:
|
||||
run:
|
||||
shell: bash -e -o pipefail {0}
|
||||
if: github.repository_owner == 'tinygrad'
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
uses: actions/checkout@v6
|
||||
# brew install uv
|
||||
- name: setup python environment
|
||||
run: |
|
||||
rm -rf /tmp/tinygrad_pytest_ci
|
||||
uv venv /tmp/tinygrad_pytest_ci
|
||||
source /tmp/tinygrad_pytest_ci/bin/activate
|
||||
uv pip install .[testing]
|
||||
- name: setup staging db
|
||||
run: |
|
||||
echo "CACHEDB=/tmp/pytest-db-ci.db" >> $GITHUB_ENV
|
||||
rm -f /tmp/pytest-db-ci*
|
||||
- name: Run pytest -nauto
|
||||
run: |
|
||||
source /tmp/tinygrad_pytest_ci/bin/activate
|
||||
pytest -nauto --durations=20
|
||||
- name: openpilot compile3 0.10.1 driving_vision
|
||||
run: FLOAT16=1 DEV=CL IMAGE=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
|
||||
|
||||
testmacbenchmark:
|
||||
name: Mac Benchmark
|
||||
env:
|
||||
@@ -64,7 +29,7 @@ jobs:
|
||||
if: github.repository_owner == 'tinygrad'
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
uses: actions/checkout@v6
|
||||
uses: actions/checkout@v4
|
||||
- name: Symlink models and datasets
|
||||
run: |
|
||||
mkdir -p weights
|
||||
@@ -84,81 +49,106 @@ jobs:
|
||||
- name: Print macOS version
|
||||
run: sw_vers
|
||||
- name: Run Stable Diffusion
|
||||
run: BENCHMARK_LOG=stable_diffusion JIT=1 ASSERT_MIN_STEP_TIME=720 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing
|
||||
run: BENCHMARK_LOG=stable_diffusion JIT=1 ASSERT_MIN_STEP_TIME=720 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 ASSERT_MIN_STEP_TIME=720 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing
|
||||
run: BENCHMARK_LOG=stable_diffusion_fp32 JIT=1 ASSERT_MIN_STEP_TIME=720 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd_no_fp16.txt
|
||||
- name: Run Stable Diffusion v2
|
||||
# TODO: very slow step time
|
||||
run: BENCHMARK_LOG=stable_diffusion_v2 JIT=1 ASSERT_MIN_STEP_TIME=4500 python3.11 examples/sdv2.py --fp16 --seed 0 --noshow --timing
|
||||
run: BENCHMARK_LOG=stable_diffusion_v2 JIT=1 ASSERT_MIN_STEP_TIME=4500 python3.11 examples/sdv2.py --fp16 --seed 0 --noshow --timing | tee sdv2.txt
|
||||
# process replay can't capture this, the graph is too large
|
||||
- name: Run SDXL
|
||||
run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=5000 CAPTURE_PROCESS_REPLAY=0 JIT=1 python3.11 examples/sdxl.py --seed 0 --noshow --timing
|
||||
run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=5000 CAPTURE_PROCESS_REPLAY=0 JIT=1 python3.11 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
|
||||
- name: Run model inference benchmark
|
||||
run: DEV=METAL NOCLANG=1 python3.11 test/external/external_model_benchmark.py
|
||||
run: METAL=1 NOCLANG=1 python3.11 test/external/external_model_benchmark.py
|
||||
- name: Test speed vs torch
|
||||
run: BIG=2 MPS=1 python3.11 test/speed/external_test_speed_v_torch.py
|
||||
run: BIG=2 MPS=1 python3.11 test/speed/external_test_speed_v_torch.py | tee torch_speed.txt
|
||||
- name: Test tensor cores
|
||||
run: DEV=METAL python3.11 test/opt/test_tensor_cores.py
|
||||
run: METAL=1 python3.11 test/opt/test_tensor_cores.py
|
||||
- name: Test AMX tensor cores
|
||||
run: |
|
||||
DEBUG=2 DEV=CPU CPU_LLVM=0 AMX=1 python3.11 test/opt/test_tensor_cores.py
|
||||
DEBUG=2 DEV=CPU CPU_LLVM=1 AMX=1 python3.11 test/opt/test_tensor_cores.py
|
||||
DEBUG=2 DEV=CPU CPU_LLVM=0 AMX=1 python3.11 test/opt/test_gen_float4.py TestFloat4.test_float4_multidim_amx TestFloat4.test_float4_multidim_unaligned_load_amx
|
||||
DEBUG=2 DEV=CPU CPU_LLVM=1 AMX=1 python3.11 test/opt/test_gen_float4.py TestFloat4.test_float4_multidim_amx TestFloat4.test_float4_multidim_unaligned_load_amx
|
||||
DEBUG=2 CPU=1 CPU_LLVM=0 AMX=1 python3.11 test/opt/test_tensor_cores.py
|
||||
DEBUG=2 CPU=1 CPU_LLVM=1 AMX=1 python3.11 test/opt/test_tensor_cores.py
|
||||
DEBUG=2 CPU=1 CPU_LLVM=0 AMX=1 python3.11 test/opt/test_gen_float4.py TestFloat4.test_float4_multidim_amx TestFloat4.test_float4_multidim_unaligned_load_amx
|
||||
DEBUG=2 CPU=1 CPU_LLVM=1 AMX=1 python3.11 test/opt/test_gen_float4.py TestFloat4.test_float4_multidim_amx TestFloat4.test_float4_multidim_unaligned_load_amx
|
||||
- name: Run Tensor Core GEMM (float)
|
||||
run: DEBUG=2 SHOULD_USE_TC=1 python3.11 extra/gemm/simple_matmul.py
|
||||
run: DEBUG=2 SHOULD_USE_TC=1 python3.11 extra/gemm/simple_matmul.py | tee matmul.txt
|
||||
- name: Run Tensor Core GEMM (half)
|
||||
run: DEBUG=2 SHOULD_USE_TC=1 HALF=1 python3.11 extra/gemm/simple_matmul.py
|
||||
run: DEBUG=2 SHOULD_USE_TC=1 HALF=1 python3.11 extra/gemm/simple_matmul.py | tee matmul_half.txt
|
||||
- name: Run Tensor Core GEMM (bfloat16)
|
||||
run: DEBUG=2 SHOULD_USE_TC=1 BFLOAT16=1 python3.11 extra/gemm/simple_matmul.py
|
||||
run: DEBUG=2 SHOULD_USE_TC=1 BFLOAT16=1 python3.11 extra/gemm/simple_matmul.py | tee matmul_bfloat16.txt
|
||||
- name: Fuzz Padded Tensor Core GEMM
|
||||
run: DEV=METAL 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
|
||||
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
|
||||
BENCHMARK_LOG=llama JIT=1 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
|
||||
BENCHMARK_LOG=llama_nojit JIT=0 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_unjitted.txt
|
||||
BENCHMARK_LOG=llama JIT=1 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_jitted.txt
|
||||
- name: Run LLaMA with BEAM
|
||||
run: BENCHMARK_LOG=llama_beam JITBEAM=2 IGNORE_BEAM_CACHE=1 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
|
||||
run: BENCHMARK_LOG=llama_beam JITBEAM=2 IGNORE_BEAM_CACHE=1 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_beam.txt
|
||||
- name: Run quantized LLaMA
|
||||
run: |
|
||||
BENCHMARK_LOG=llama_int8 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing --quantize int8
|
||||
BENCHMARK_LOG=llama_nf4 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing --quantize nf4
|
||||
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
|
||||
BENCHMARK_LOG=llama3_nf4 python3.11 examples/llama3.py --size 8B --temperature 0 --benchmark --quantize nf4
|
||||
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
|
||||
# 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
|
||||
BENCHMARK_LOG=gpt2 JIT=1 ASSERT_MIN_STEP_TIME=13 python3.11 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing
|
||||
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 ASSERT_MIN_STEP_TIME=13 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
|
||||
run: BENCHMARK_LOG=gpt2_half HALF=1 python3.11 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
|
||||
- name: Run GPT2 w HALF/BEAM
|
||||
run: BENCHMARK_LOG=gpt2_half_beam HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3.11 examples/gpt2.py --count 10 --temperature 0 --timing
|
||||
run: BENCHMARK_LOG=gpt2_half_beam HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3.11 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half_beam.txt
|
||||
- name: Run OLMoE
|
||||
run: BENCHMARK_LOG=olmoe python3.11 examples/olmoe.py
|
||||
- name: Train MNIST
|
||||
run: time PYTHONPATH=. TARGET_EVAL_ACC_PCT=96.0 python3.11 examples/beautiful_mnist.py
|
||||
run: time PYTHONPATH=. TARGET_EVAL_ACC_PCT=96.0 python3.11 examples/beautiful_mnist.py | tee beautiful_mnist.txt
|
||||
|
||||
# NOTE: this is failing in CI. it is not failing on my machine and I don't really have a way to debug it
|
||||
# the error is "RuntimeError: Internal Error (0000000e:Internal Error)"
|
||||
#- name: Run 10 CIFAR training steps
|
||||
# run: BENCHMARK_LOG=cifar_10steps JIT=1 ASSERT_MIN_STEP_TIME=3000 STEPS=10 python3.11 examples/hlb_cifar10.py
|
||||
# run: BENCHMARK_LOG=cifar_10steps JIT=1 ASSERT_MIN_STEP_TIME=3000 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 ASSERT_MIN_STEP_TIME=3000 STEPS=10 DEFAULT_FLOAT=HALF python3.11 examples/hlb_cifar10.py
|
||||
# run: BENCHMARK_LOG=cifar_10steps_half JIT=2 ASSERT_MIN_STEP_TIME=3000 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
|
||||
# run: STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3.11 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
|
||||
# TODO: too slow
|
||||
# - name: Run 10 CIFAR training steps w winograd
|
||||
# run: BENCHMARK_LOG=cifar_10steps_wino JIT=1 ASSERT_MIN_STEP_TIME=150 WINO=1 STEPS=10 python3.11 examples/hlb_cifar10.py
|
||||
- uses: actions/upload-artifact@v7
|
||||
# run: BENCHMARK_LOG=cifar_10steps_wino JIT=1 ASSERT_MIN_STEP_TIME=150 WINO=1 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar_wino.txt
|
||||
- uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: Speed (Mac)
|
||||
path: |
|
||||
onnx_inference_speed.csv
|
||||
torch_speed.txt
|
||||
llama_unjitted.txt
|
||||
llama_jitted.txt
|
||||
llama_beam.txt
|
||||
llama_int8.txt
|
||||
llama_nf4.txt
|
||||
llama3_int8.txt
|
||||
llama3_nf4.txt
|
||||
llama_four_gpu.txt
|
||||
gpt2_unjitted.txt
|
||||
gpt2_jitted.txt
|
||||
gpt2_half.txt
|
||||
gpt2_half_beam.txt
|
||||
matmul.txt
|
||||
matmul_half.txt
|
||||
matmul_bfloat16.txt
|
||||
sd.txt
|
||||
sd_no_fp16.txt
|
||||
sdv2.txt
|
||||
sdxl.txt
|
||||
beautiful_mnist.txt
|
||||
train_cifar.txt
|
||||
train_cifar_half.txt
|
||||
train_cifar_bf16.txt
|
||||
train_cifar_wino.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.11 process_replay.py
|
||||
|
||||
@@ -174,28 +164,24 @@ jobs:
|
||||
if: github.repository_owner == 'tinygrad'
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
uses: actions/checkout@v6
|
||||
uses: actions/checkout@v4
|
||||
- name: setup staging db
|
||||
if: github.ref == 'refs/heads/update_benchmark_staging'
|
||||
run: |
|
||||
echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
|
||||
rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
|
||||
- name: Kill stale pids
|
||||
run: |
|
||||
PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
|
||||
PYTHONPATH=. ./extra/hcq/hcq_smi.py nv kill_pids
|
||||
- name: UsbGPU boot time
|
||||
run: sudo -E PYTHONPATH=. GMMU=0 DEBUG=2 AM_RESET=1 DEV=AMD AMD_IFACE=USB time python3.11 test/test_tiny.py TestTiny.test_plus
|
||||
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=. GMMU=0 DEV=AMD AMD_IFACE=USB python3.11 test/test_tiny.py
|
||||
run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB python3.11 test/test_tiny.py
|
||||
- name: UsbGPU copy speeds
|
||||
run: sudo -E PYTHONPATH=. GMMU=0 DEV=AMD AMD_IFACE=USB python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
|
||||
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=. GMMU=0 DEV=AMD AMD_IFACE=USB GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
|
||||
# run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
|
||||
- name: UsbGPU (USB4/TB) boot time
|
||||
run: PYTHONPATH=. DEBUG=3 DEV=NV NV_IFACE=PCI NV_NAK=1 time python3.11 test/test_tiny.py TestTiny.test_plus
|
||||
run: PYTHONPATH=. DEBUG=3 NV=1 NV_IFACE=PCI NV_NAK=1 time python3.11 test/test_tiny.py TestTiny.test_plus
|
||||
- name: UsbGPU (USB4/TB) tiny tests
|
||||
run: PYTHONPATH=. DEV=NV NV_IFACE=PCI NV_NAK=1 python3.11 test/test_tiny.py
|
||||
run: PYTHONPATH=. NV=1 NV_IFACE=PCI NV_NAK=1 python3.11 test/test_tiny.py
|
||||
|
||||
testnvidiabenchmark:
|
||||
name: tinybox green Benchmark
|
||||
@@ -207,7 +193,7 @@ jobs:
|
||||
if: github.repository_owner == 'tinygrad'
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
uses: actions/checkout@v6
|
||||
uses: actions/checkout@v4
|
||||
- name: Print nvidia-smi
|
||||
run: nvidia-smi
|
||||
- name: Symlink models and datasets
|
||||
@@ -227,71 +213,92 @@ jobs:
|
||||
- name: reset process replay
|
||||
run: test/external/process_replay/reset.py
|
||||
- name: Run model inference benchmark
|
||||
run: DEV=NV CAPTURE_PROCESS_REPLAY=0 NOCLANG=1 python3 test/external/external_model_benchmark.py
|
||||
run: NV=1 CAPTURE_PROCESS_REPLAY=0 NOCLANG=1 python3 test/external/external_model_benchmark.py
|
||||
- name: Test speed vs torch
|
||||
run: DEV=NV CAPTURE_PROCESS_REPLAY=0 HALF=1 BIG=2 TORCHCUDA=1 python3 test/speed/external_test_speed_v_torch.py
|
||||
run: NV=1 CAPTURE_PROCESS_REPLAY=0 HALF=1 BIG=2 TORCHCUDA=1 python3 test/speed/external_test_speed_v_torch.py | tee torch_speed.txt
|
||||
- name: Test speed vs theoretical
|
||||
run: DEV=NV IGNORE_BEAM_CACHE=1 CCACHE=0 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
|
||||
run: NV=1 IGNORE_BEAM_CACHE=1 CCACHE=0 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
|
||||
- name: Test benchmark allreduce
|
||||
run: DEV=NV python test/external/external_benchmark_multitensor_allreduce.py
|
||||
run: NV=1 python test/external/external_benchmark_multitensor_allreduce.py
|
||||
- name: Test tensor cores
|
||||
run: |
|
||||
DEV=NV ALLOW_TF32=1 python3 test/opt/test_tensor_cores.py
|
||||
DEV=NV NV_PTX=1 ALLOW_TF32=1 python3 test/opt/test_tensor_cores.py
|
||||
NV=1 ALLOW_TF32=1 python3 test/opt/test_tensor_cores.py
|
||||
NV=1 NV_PTX=1 ALLOW_TF32=1 python3 test/opt/test_tensor_cores.py
|
||||
- name: Run Tensor Core GEMM (CUDA)
|
||||
run: |
|
||||
DEV=CUDA SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
|
||||
DEV=CUDA SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
|
||||
DEV=CUDA SHOULD_USE_TC=1 ALLOW_TF32=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py
|
||||
DEV=CUDA SHOULD_USE_TC=1 FP8E4M3=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
|
||||
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 SHOULD_USE_TC=1 FP8E4M3=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_fp8.txt
|
||||
- name: Run Tensor Core GEMM (PTX)
|
||||
run: DEV=NV NV_PTX=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
|
||||
run: NV=1 NV_PTX=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_ptx.txt
|
||||
- name: Run Tensor Core GEMM (NV)
|
||||
run: DEV=NV SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
|
||||
- name: Test DEV=NV
|
||||
run: DEBUG=2 DEV=NV python -m pytest -rA test/test_tiny.py
|
||||
- name: Test DEV=CUDA
|
||||
run: DEBUG=2 DEV=CUDA python -m pytest -rA test/test_tiny.py
|
||||
run: NV=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_nv.txt
|
||||
- name: Test NV=1
|
||||
run: DEBUG=2 NV=1 python -m pytest -rA test/test_tiny.py
|
||||
- name: Test CUDA=1
|
||||
run: DEBUG=2 CUDA=1 python -m pytest -rA test/test_tiny.py
|
||||
- name: Run Stable Diffusion
|
||||
run: BENCHMARK_LOG=stable_diffusion DEV=NV python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing
|
||||
run: BENCHMARK_LOG=stable_diffusion NV=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
|
||||
# TODO: too slow
|
||||
# - name: Run SDXL
|
||||
# run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=2000 CAPTURE_PROCESS_REPLAY=0 DEV=NV CAPTURE_PROCESS_REPLAY=0 python3 examples/sdxl.py --seed 0 --noshow --timing
|
||||
# run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=2000 CAPTURE_PROCESS_REPLAY=0 NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
|
||||
- name: Run LLaMA
|
||||
run: |
|
||||
BENCHMARK_LOG=llama_nojit DEV=NV JIT=0 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
|
||||
BENCHMARK_LOG=llama DEV=NV JIT=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
|
||||
BENCHMARK_LOG=llama_nojit NV=1 JIT=0 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_unjitted.txt
|
||||
BENCHMARK_LOG=llama NV=1 JIT=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_jitted.txt
|
||||
- name: Run LLaMA with BEAM
|
||||
run: BENCHMARK_LOG=llama_beam DEV=NV JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
|
||||
run: BENCHMARK_LOG=llama_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_beam.txt
|
||||
# - name: Run LLaMA 7B on 4 GPUs
|
||||
# run: DEV=NV CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0 --timing
|
||||
# run: NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_four_gpu.txt
|
||||
# - name: Run LLaMA 7B on 6 GPUs
|
||||
# run: DEV=NV CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 6 --prompt "Hello." --count 10 --temperature 0 --timing
|
||||
# run: NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 6 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_six_gpu.txt
|
||||
- name: Run LLaMA-3 8B BEAM
|
||||
run: BENCHMARK_LOG=llama3_beam DEV=NV JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
|
||||
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 DEV=NV 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
|
||||
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 quantized LLaMA3
|
||||
run: BENCHMARK_LOG=llama3_fp8 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --temperature 0 --benchmark --quantize fp8
|
||||
run: BENCHMARK_LOG=llama3_fp8 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --temperature 0 --benchmark --quantize fp8 | tee llama3_fp8.txt
|
||||
# - name: Run LLaMA-3 8B on 6 GPUs
|
||||
# run: DEV=NV CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 6 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
|
||||
# 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: DEV=NV CAPTURE_PROCESS_REPLAY=0 MAX_CONTEXT=256 python3 examples/llama.py --gen 2 --size 70B --shard 6 --prompt "Hello." --count 10 --temperature 0 --timing
|
||||
# run: NV=1 CAPTURE_PROCESS_REPLAY=0 MAX_CONTEXT=256 python3 examples/llama.py --gen 2 --size 70B --shard 6 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_2_70B.txt
|
||||
- name: Run Mixtral 8x7B
|
||||
run: time BENCHMARK_LOG=mixtral DEV=NV CAPTURE_PROCESS_REPLAY=0 python3 examples/mixtral.py --temperature 0 --count 10 --timing
|
||||
run: time BENCHMARK_LOG=mixtral NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/mixtral.py --temperature 0 --count 10 --timing | tee mixtral.txt
|
||||
- name: Run GPT2
|
||||
run: |
|
||||
BENCHMARK_LOG=gpt2_nojit DEV=NV JIT=0 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing
|
||||
BENCHMARK_LOG=gpt2 DEV=NV JIT=1 ASSERT_MIN_STEP_TIME=4 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing
|
||||
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 ASSERT_MIN_STEP_TIME=4 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 DEV=NV HALF=1 ASSERT_MIN_STEP_TIME=6 python3 examples/gpt2.py --count 10 --temperature 0 --timing
|
||||
run: BENCHMARK_LOG=gpt2_half NV=1 HALF=1 ASSERT_MIN_STEP_TIME=6 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 DEV=NV HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing
|
||||
- uses: actions/upload-artifact@v7
|
||||
run: BENCHMARK_LOG=gpt2_half_beam NV=1 HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half_beam.txt
|
||||
- uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: Speed (NVIDIA)
|
||||
path: |
|
||||
onnx_inference_speed.csv
|
||||
torch_speed.txt
|
||||
matmul.txt
|
||||
matmul_bfloat16.txt
|
||||
matmul_tf32.txt
|
||||
matmul_ptx.txt
|
||||
matmul_nv.txt
|
||||
sd.txt
|
||||
sdxl.txt
|
||||
llama_unjitted.txt
|
||||
llama_jitted.txt
|
||||
llama_beam.txt
|
||||
llama3_beam.txt
|
||||
llama3_four_gpu.txt
|
||||
llama3_six_gpu.txt
|
||||
llama3_fp8.txt
|
||||
llama_2_70B.txt
|
||||
mixtral.txt
|
||||
gpt2_unjitted.txt
|
||||
gpt2_jitted.txt
|
||||
gpt2_half.txt
|
||||
gpt2_half_beam.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
|
||||
|
||||
@@ -305,7 +312,7 @@ jobs:
|
||||
if: github.repository_owner == 'tinygrad'
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
uses: actions/checkout@v6
|
||||
uses: actions/checkout@v4
|
||||
- name: Symlink models and datasets
|
||||
run: |
|
||||
mkdir -p weights
|
||||
@@ -325,35 +332,49 @@ jobs:
|
||||
run: test/external/process_replay/reset.py
|
||||
# TODO: too slow
|
||||
# - name: Fuzz Padded Tensor Core GEMM (NV)
|
||||
# run: DEV=NV M_START=12 M_STOP=20 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 python3 ./extra/gemm/fuzz_matmul.py
|
||||
# run: NV=1 M_START=12 M_STOP=20 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 python3 ./extra/gemm/fuzz_matmul.py
|
||||
# TODO: too slow
|
||||
# - name: Fuzz Padded Tensor Core GEMM (PTX)
|
||||
# run: DEV=NV NV_PTX=1 M_START=12 M_STOP=20 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 python3 ./extra/gemm/fuzz_matmul.py
|
||||
# run: NV=1 NV_PTX=1 M_START=12 M_STOP=20 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 python3 ./extra/gemm/fuzz_matmul.py
|
||||
- name: HEVC Decode Benchmark
|
||||
run: VALIDATE=1 MAX_FRAMES=100 ASSERT_FPS=1400 JITBEAM=1 DEV=NV PYTHONPATH=. python3 extra/hevc/decode.py
|
||||
run: VALIDATE=1 MAX_FRAMES=100 NV=1 PYTHONPATH=. python3 extra/hevc/decode.py
|
||||
- name: Train MNIST
|
||||
run: time PYTHONPATH=. DEV=NV TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py
|
||||
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 ASSERT_MIN_STEP_TIME=120 DEV=NV STEPS=10 python3 examples/hlb_cifar10.py
|
||||
run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=120 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 ASSERT_MIN_STEP_TIME=120 DEV=NV STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
|
||||
run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=110 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 ASSERT_MIN_STEP_TIME=120 DEV=NV STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py
|
||||
run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=120 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 ASSERT_MIN_STEP_TIME=350 DEV=NV WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
|
||||
# run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=350 NV=1 CAPTURE_PROCESS_REPLAY=0 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
|
||||
- name: Run full CIFAR training w 1 GPU
|
||||
run: time BENCHMARK_LOG=cifar DEV=NV DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
|
||||
run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 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 DEV=NV DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
|
||||
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.0 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 DEV=NV MODEL=resnet python3 examples/mlperf/model_eval.py
|
||||
run: time BENCHMARK_LOG=resnet_eval NV=1 MODEL=resnet python3 examples/mlperf/model_eval.py
|
||||
- name: Run 10 MLPerf ResNet50 training steps (1 gpu)
|
||||
run: BENCHMARK_LOG=resnet_10steps DEV=NV DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
|
||||
run: BENCHMARK_LOG=resnet_10steps NV=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet_one_gpu.txt
|
||||
- name: Run 10 MLPerf ResNet50 training steps (6 gpu)
|
||||
run: BENCHMARK_LOG=resnet_10steps_6gpu DEV=NV CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py
|
||||
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 DEV=NV CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
|
||||
run: BENCHMARK_LOG=bert_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
|
||||
- uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: Speed (NVIDIA Training)
|
||||
path: |
|
||||
beautiful_mnist.txt
|
||||
train_cifar.txt
|
||||
train_cifar_half.txt
|
||||
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
|
||||
- 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
|
||||
|
||||
@@ -367,13 +388,11 @@ jobs:
|
||||
if: github.repository_owner == 'tinygrad'
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
uses: actions/checkout@v6
|
||||
- name: Setcap to python
|
||||
run: ./extra/amdpci/setup_python_cap.sh
|
||||
- name: Remove amd modules
|
||||
run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd rmmod
|
||||
- name: Kill stale pids
|
||||
run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
|
||||
uses: actions/checkout@v4
|
||||
- name: Remove amdgpu
|
||||
run: sudo rmmod amdgpu || true
|
||||
- name: Cleanup running AM processes
|
||||
run: python extra/amdpci/am_smi.py --pids --kill
|
||||
#- name: Insert amdgpu
|
||||
# run: sudo modprobe amdgpu
|
||||
- name: Symlink models and datasets
|
||||
@@ -407,20 +426,18 @@ jobs:
|
||||
#- name: Test speed vs torch
|
||||
# run: |
|
||||
# python3 -c "import torch; print(torch.__version__)"
|
||||
# LD_PRELOAD="/opt/rocm/lib/libhsa-runtime64.so" HSA=1 BIG=2 TORCHCUDA=1 python3 test/speed/external_test_speed_v_torch.py
|
||||
# LD_PRELOAD="/opt/rocm/lib/libhsa-runtime64.so" HSA=1 BIG=2 TORCHCUDA=1 python3 test/speed/external_test_speed_v_torch.py | tee torch_speed.txt
|
||||
- name: Test speed vs theoretical
|
||||
run: DEV=AMD IGNORE_BEAM_CACHE=1 CCACHE=0 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
|
||||
- name: Test tensor cores AMD_LLVM=0
|
||||
run: DEV=AMD AMD_LLVM=0 python3 test/opt/test_tensor_cores.py
|
||||
# TODO: this is flaky
|
||||
# - name: Test tensor cores AMD_LLVM=1
|
||||
# run: DEV=AMD AMD_LLVM=1 python3 test/opt/test_tensor_cores.py
|
||||
- name: Run Tensor Core GEMM (AMD)
|
||||
run: AMD=1 IGNORE_BEAM_CACHE=1 CCACHE=0 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
|
||||
- name: Test tensor cores
|
||||
run: |
|
||||
DEV=AMD SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
|
||||
DEV=AMD SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py
|
||||
- name: Test DEV=AMD
|
||||
run: DEBUG=2 DEV=AMD python -m pytest -rA test/test_tiny.py
|
||||
AMD=1 AMD_LLVM=0 python3 test/opt/test_tensor_cores.py
|
||||
AMD=1 AMD_LLVM=1 python3 test/opt/test_tensor_cores.py
|
||||
AMD=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
|
||||
- name: Run Tensor Core GEMM (AMD)
|
||||
run: AMD=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py | tee matmul_amd.txt
|
||||
- name: Test AMD=1
|
||||
run: DEBUG=2 AMD=1 python -m pytest -rA test/test_tiny.py
|
||||
#- name: Test HIP=1
|
||||
# run: DEBUG=2 HIP=1 python -m pytest -rA test/test_tiny.py
|
||||
# TODO: AMD compiler bug causes this to fail
|
||||
@@ -429,43 +446,65 @@ jobs:
|
||||
#- 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 DEV=AMD AM_RESET=1 python3 test/test_tiny.py TestTiny.test_plus
|
||||
run: time AMD=1 AM_RESET=1 python3 test/test_tiny.py TestTiny.test_plus
|
||||
- name: Test AM warm start time
|
||||
run: time DEV=AMD python3 test/test_tiny.py TestTiny.test_plus
|
||||
run: time AMD=1 python3 test/test_tiny.py TestTiny.test_plus
|
||||
- name: Run Stable Diffusion
|
||||
run: BENCHMARK_LOG=stable_diffusion ASSERT_MIN_STEP_TIME=550 DEV=AMD python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing
|
||||
run: BENCHMARK_LOG=stable_diffusion ASSERT_MIN_STEP_TIME=550 AMD=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
|
||||
- name: Run SDXL
|
||||
run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=3200 CAPTURE_PROCESS_REPLAY=0 DEV=AMD python3 examples/sdxl.py --seed 0 --noshow --timing
|
||||
run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=3200 CAPTURE_PROCESS_REPLAY=0 AMD=1 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
|
||||
- name: Run LLaMA 7B
|
||||
run: |
|
||||
BENCHMARK_LOG=llama_nojit DEV=AMD JIT=0 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
|
||||
BENCHMARK_LOG=llama DEV=AMD JIT=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
|
||||
BENCHMARK_LOG=llama_nojit AMD=1 JIT=0 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_unjitted.txt
|
||||
BENCHMARK_LOG=llama AMD=1 JIT=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_jitted.txt
|
||||
- name: Run LLaMA 7B with BEAM
|
||||
run: BENCHMARK_LOG=llama_beam DEV=AMD JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
|
||||
run: BENCHMARK_LOG=llama_beam AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_beam.txt
|
||||
# - name: Run LLaMA 7B on 4 GPUs
|
||||
# run: DEV=AMD CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0 --timing
|
||||
# run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_four_gpu.txt
|
||||
# - name: Run LLaMA 7B on 6 GPUs
|
||||
# run: DEV=AMD CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 6 --prompt "Hello." --count 10 --temperature 0 --timing
|
||||
# run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 6 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_six_gpu.txt
|
||||
- name: Run LLaMA-3 8B BEAM
|
||||
run: BENCHMARK_LOG=llama3_beam DEV=AMD JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
|
||||
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 DEV=AMD 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
|
||||
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: DEV=AMD CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 6 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
|
||||
# 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: DEV=AMD CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 2 --size 70B --shard 6 --prompt "Hello." --count 10 --temperature 0 --timing
|
||||
# run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 2 --size 70B --shard 6 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_2_70B.txt
|
||||
- name: Run Mixtral 8x7B
|
||||
run: time BENCHMARK_LOG=mixtral DEV=AMD python3 examples/mixtral.py --temperature 0 --count 10 --timing
|
||||
run: time BENCHMARK_LOG=mixtral AMD=1 python3 examples/mixtral.py --temperature 0 --count 10 --timing | tee mixtral.txt
|
||||
- name: Run GPT2
|
||||
run: |
|
||||
BENCHMARK_LOG=gpt2_nojit DEV=AMD JIT=0 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing
|
||||
BENCHMARK_LOG=gpt2 DEV=AMD JIT=1 ASSERT_MIN_STEP_TIME=5 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing
|
||||
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 ASSERT_MIN_STEP_TIME=5 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 DEV=AMD HALF=1 ASSERT_MIN_STEP_TIME=5 python3 examples/gpt2.py --count 10 --temperature 0 --timing
|
||||
run: BENCHMARK_LOG=gpt2_half AMD=1 HALF=1 ASSERT_MIN_STEP_TIME=5 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 DEV=AMD HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing
|
||||
run: BENCHMARK_LOG=gpt2_half_beam AMD=1 HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half_beam.txt
|
||||
- uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: Speed (AMD)
|
||||
path: |
|
||||
onnx_inference_speed.csv
|
||||
torch_speed.txt
|
||||
llama_unjitted.txt
|
||||
llama_jitted.txt
|
||||
llama_beam.txt
|
||||
llama3_beam.txt
|
||||
llama3_four_gpu.txt
|
||||
llama3_six_gpu.txt
|
||||
llama_2_70B.txt
|
||||
gpt2_unjitted.txt
|
||||
gpt2_jitted.txt
|
||||
gpt2_half.txt
|
||||
gpt2_half_beam.txt
|
||||
matmul.txt
|
||||
matmul_amd.txt
|
||||
sd.txt
|
||||
sdxl.txt
|
||||
mixtral.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
|
||||
|
||||
@@ -479,13 +518,11 @@ jobs:
|
||||
if: github.repository_owner == 'tinygrad'
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
uses: actions/checkout@v6
|
||||
- name: Setcap to python
|
||||
run: ./extra/amdpci/setup_python_cap.sh
|
||||
- name: Remove amd modules
|
||||
run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd rmmod
|
||||
- name: Kill stale pids
|
||||
run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
|
||||
uses: actions/checkout@v4
|
||||
- name: Remove amdgpu
|
||||
run: sudo rmmod amdgpu || true
|
||||
- name: Cleanup running AM processes
|
||||
run: python extra/amdpci/am_smi.py --pids --kill
|
||||
- name: Symlink models and datasets
|
||||
run: |
|
||||
mkdir -p weights
|
||||
@@ -504,23 +541,31 @@ jobs:
|
||||
- name: reset process replay
|
||||
run: test/external/process_replay/reset.py
|
||||
- name: Train MNIST
|
||||
run: time PYTHONPATH=. DEV=AMD TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py
|
||||
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 ASSERT_MIN_STEP_TIME=200 DEV=AMD STEPS=10 python3 examples/hlb_cifar10.py
|
||||
run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=200 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 ASSERT_MIN_STEP_TIME=230 DEV=AMD STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
|
||||
run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=200 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 ASSERT_MIN_STEP_TIME=288 DEV=AMD STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py
|
||||
# run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=288 AMD=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
|
||||
# TODO: too slow
|
||||
# - name: Run 10 CIFAR training steps w winograd
|
||||
# run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=66 DEV=AMD WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
|
||||
# run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=66 AMD=1 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
|
||||
- name: Run full CIFAR training w 1 GPU
|
||||
run: time BENCHMARK_LOG=cifar DEV=AMD DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
|
||||
run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 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 DEV=AMD DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
|
||||
# TODO: broken on some of the machines
|
||||
#- name: Test full tinyfs load
|
||||
# run: TINYFS_ENDPOINT=10.0.52.11:6767 PYTHONPATH=. python extra/tinyfs/fetch_file.py --hash d734f5e3be9f1e9d863bfaa4fc6c1ef2 --len 175866113 --dest mapping.json --check
|
||||
run: time BENCHMARK_LOG=cifar_6gpu AMD=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
|
||||
- uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: Speed (AMD Training)
|
||||
path: |
|
||||
beautiful_mnist.txt
|
||||
train_cifar.txt
|
||||
train_cifar_half.txt
|
||||
train_cifar_bf16.txt
|
||||
train_cifar_wino.txt
|
||||
train_cifar_one_gpu.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
|
||||
|
||||
@@ -534,13 +579,11 @@ jobs:
|
||||
if: github.repository_owner == 'tinygrad'
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
uses: actions/checkout@v6
|
||||
- name: Setcap to python
|
||||
run: ./extra/amdpci/setup_python_cap.sh
|
||||
- name: Remove amd modules
|
||||
run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd rmmod
|
||||
- name: Kill stale pids
|
||||
run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
|
||||
uses: actions/checkout@v4
|
||||
- name: Remove amdgpu
|
||||
run: sudo rmmod amdgpu || true
|
||||
- name: Cleanup running AM processes
|
||||
run: python extra/amdpci/am_smi.py --pids --kill
|
||||
- name: Symlink models and datasets
|
||||
run: |
|
||||
mkdir -p weights
|
||||
@@ -559,14 +602,21 @@ jobs:
|
||||
- name: reset process replay
|
||||
run: test/external/process_replay/reset.py
|
||||
- name: Run MLPerf resnet eval
|
||||
run: time BENCHMARK_LOG=resnet_eval DEV=AMD MODEL=resnet python3 examples/mlperf/model_eval.py
|
||||
run: time BENCHMARK_LOG=resnet_eval AMD=1 MODEL=resnet python3 examples/mlperf/model_eval.py
|
||||
- name: Run 10 MLPerf ResNet50 training steps (1 gpu)
|
||||
run: BENCHMARK_LOG=resnet_10steps DEV=AMD DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
|
||||
run: BENCHMARK_LOG=resnet_10steps AMD=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet_one_gpu.txt
|
||||
- name: Run 10 MLPerf ResNet50 training steps (6 gpu)
|
||||
run: BENCHMARK_LOG=resnet_10steps_6gpu DEV=AMD CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py
|
||||
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 DEV=AMD CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
|
||||
run: BENCHMARK_LOG=bert_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
|
||||
- uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: Speed (AMD MLPerf)
|
||||
path: |
|
||||
train_resnet.txt
|
||||
train_resnet_one_gpu.txt
|
||||
train_bert.txt
|
||||
- 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
|
||||
|
||||
@@ -580,7 +630,7 @@ jobs:
|
||||
if: github.repository_owner == 'tinygrad'
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
uses: actions/checkout@v6
|
||||
uses: actions/checkout@v4
|
||||
- name: setup staging db
|
||||
if: github.ref == 'refs/heads/update_benchmark_staging'
|
||||
run: |
|
||||
@@ -588,54 +638,31 @@ 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 compile3 0.11.0 driving_vision
|
||||
run: BENCHMARK_LOG=openpilot_0_11_0_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_vision.onnx
|
||||
- name: IR3 openpilot compile3 0.11.0 driving_vision
|
||||
run: BENCHMARK_LOG=ir3_openpilot_0_11_0_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM QCOM_IR3=1 FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_vision.onnx
|
||||
- name: openpilot compile3 0.11.0 driving_policy
|
||||
run: BENCHMARK_LOG=openpilot_0_11_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_policy.onnx
|
||||
- name: openpilot compile3 0.11.0 dmonitoring
|
||||
run: BENCHMARK_LOG=openpilot_0_11_0_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/dmonitoring_model.onnx
|
||||
- name: openpilot compile3 0.10.0 driving_policy
|
||||
run: BENCHMARK_LOG=openpilot_0_10_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=4 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.10.0/selfdrive/modeld/models/driving_policy.onnx
|
||||
- name: openpilot compile3 0.10.0 dmonitoring
|
||||
run: BENCHMARK_LOG=openpilot_0_10_0_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.10.0/selfdrive/modeld/models/dmonitoring_model.onnx
|
||||
- name: DEBUG=2 openpilot compile3 0.10.1 driving_vision
|
||||
run: PYTHONPATH="." DEBUG=2 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
|
||||
- name: DEBUG=2 IMAGE=1 openpilot compile3 0.10.1 driving_vision
|
||||
run: PYTHONPATH="." DEBUG=2 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
|
||||
- name: openpilot compile3 0.10.1 driving_vision
|
||||
run: BENCHMARK_LOG=openpilot_0_10_1_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
|
||||
run: BENCHMARK_LOG=openpilot_0_10_1_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
|
||||
- name: openpilot compile3 0.10.1 driving_policy
|
||||
run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_policy.onnx
|
||||
run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=4 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_policy.onnx
|
||||
- name: openpilot compile3 0.10.1 dmonitoring
|
||||
run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/dmonitoring_model.onnx
|
||||
run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=10 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/dmonitoring_model.onnx
|
||||
- name: benchmark MobileNetV2 on DSP
|
||||
run: |
|
||||
# generate quantized weights
|
||||
ln -s /data/home/tiny/tinygrad/extra/datasets/imagenet extra/datasets/imagenet
|
||||
ln -s /data/home/tiny/tinygrad/testsig-*.so .
|
||||
PYTHONPATH=. CC=clang-19 DEV=CPU CPU_LLVM=0 QUANT=1 CNT=0 python3 examples/test_onnx_imagenet.py https://github.com/xamcat/mobcat-samples/raw/refs/heads/master/onnx_runtime/InferencingSample/InferencingSample/mobilenetv2-7.onnx /tmp/model.quant.onnx
|
||||
PYTHONPATH=. CC=clang-19 CPU=1 CPU_LLVM=0 QUANT=1 CNT=0 python3 examples/test_onnx_imagenet.py https://github.com/xamcat/mobcat-samples/raw/refs/heads/master/onnx_runtime/InferencingSample/InferencingSample/mobilenetv2-7.onnx /tmp/model.quant.onnx
|
||||
# benchmark on DSP with NOOPT=1, the devectorizer has issues
|
||||
PYTHONPATH=. CC=clang-19 DEV=DSP NOOPT=1 CNT=2 DEBUG=2 python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
|
||||
PYTHONPATH=. CC=clang-19 DSP=1 NOOPT=1 CNT=2 DEBUG=2 python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
|
||||
- name: Run process replay tests
|
||||
run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
|
||||
|
||||
testcommausbgpubenchmark:
|
||||
name: UsbGPU Benchmark (comma)
|
||||
runs-on: [self-hosted, Linux, comma4]
|
||||
timeout-minutes: 20
|
||||
defaults:
|
||||
run:
|
||||
shell: bash -e -o pipefail {0}
|
||||
if: github.repository_owner == 'tinygrad'
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
uses: actions/checkout@v6
|
||||
- name: setup staging db
|
||||
if: github.ref == 'refs/heads/update_benchmark_staging'
|
||||
run: |
|
||||
echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
|
||||
rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
|
||||
- name: openpilot compile3 0.10.1 driving_vision
|
||||
run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision PYTHONPATH="." GMMU=0 DEV=AMD AMD_LLVM=1 AMD_IFACE=USB ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
|
||||
- name: openpilot load_pickle 0.10.1 driving_vision
|
||||
run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision_load_pickle PYTHONPATH="." GMMU=0 DEV=AMD AMD_IFACE=USB ASSERT_MIN_LOAD_TIME=15 python3 examples/openpilot/load_pickle.py
|
||||
|
||||
testreddriverbenchmark:
|
||||
name: AM Benchmark
|
||||
runs-on: [self-hosted, Linux, tinyboxrandom]
|
||||
@@ -646,13 +673,11 @@ jobs:
|
||||
if: github.repository_owner == 'tinygrad'
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
uses: actions/checkout@v6
|
||||
- name: Setcap to python
|
||||
run: ./extra/amdpci/setup_python_cap.sh
|
||||
uses: actions/checkout@v4
|
||||
- name: Remove amd modules
|
||||
run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd rmmod
|
||||
run: ./extra/hcq/hcq_smi.py amd rmmod
|
||||
- name: Kill stale pids
|
||||
run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
|
||||
run: ./extra/hcq/hcq_smi.py amd kill_pids
|
||||
- name: Symlink models and datasets
|
||||
run: |
|
||||
mkdir -p weights
|
||||
@@ -671,40 +696,40 @@ jobs:
|
||||
- name: reset process replay
|
||||
run: test/external/process_replay/reset.py
|
||||
- name: Test driver cold start time
|
||||
run: time DEBUG=3 DEV=AMD AM_RESET=1 python3 test/test_tiny.py TestTiny.test_plus
|
||||
run: time DEBUG=3 AMD=1 AM_RESET=1 python3 test/test_tiny.py TestTiny.test_plus
|
||||
- name: Test driver warm start time
|
||||
run: time DEBUG=3 DEV=AMD python3 test/test_tiny.py TestTiny.test_plus
|
||||
run: time DEBUG=3 AMD=1 python3 test/test_tiny.py TestTiny.test_plus
|
||||
# Fails on 9070
|
||||
# - name: Test tensor cores
|
||||
# run: |
|
||||
# DEV=AMD AMD_LLVM=0 python3 test/test_linearizer.py test/opt/test_tensor_cores.py
|
||||
# DEV=AMD AMD_LLVM=1 python3 test/test_linearizer.py test/opt/test_tensor_cores.py
|
||||
# DEV=AMD SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
|
||||
# AMD=1 AMD_LLVM=0 python3 test/test_linearizer.py test/opt/test_tensor_cores.py
|
||||
# AMD=1 AMD_LLVM=1 python3 test/test_linearizer.py test/opt/test_tensor_cores.py
|
||||
# AMD=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
|
||||
- name: Run Tensor Core GEMM (AMD)
|
||||
run: DEV=AMD SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py
|
||||
- name: Test DEV=AMD
|
||||
run: DEBUG=2 DEV=AMD python -m pytest -rA test/test_tiny.py
|
||||
run: AMD=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py | tee am_matmul_amd.txt
|
||||
- name: Test AMD=1
|
||||
run: DEBUG=2 AMD=1 python -m pytest -rA test/test_tiny.py
|
||||
- name: Test DISK copy time
|
||||
run: DEV=AMD TESTFILE=/raid/downloads/llama3-8b-sfr/model-00001-of-00004.safetensors python3 test/external/external_benchmark_disk_raw.py
|
||||
run: AMD=1 TESTFILE=/raid/downloads/llama3-8b-sfr/model-00001-of-00004.safetensors python3 test/external/external_benchmark_disk_raw.py
|
||||
- name: Test CPU copy time
|
||||
run: |
|
||||
DEV=AMD GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyDefaulttoCPUJit
|
||||
DEV=AMD GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
|
||||
AMD=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyDefaulttoCPUJit
|
||||
AMD=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
|
||||
- name: Run full CIFAR training w 1 GPU
|
||||
run: time BENCHMARK_LOG=cifar DEV=AMD DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
|
||||
run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee am_train_cifar_one_gpu.txt
|
||||
# - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
|
||||
# run: BENCHMARK_LOG=resnet_10steps DEV=AMD MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
|
||||
# run: BENCHMARK_LOG=resnet_10steps AMD=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee am_train_resnet_one_gpu.txt
|
||||
- name: Run 10 MLPerf Bert training steps (1 gpu)
|
||||
# TODO: remove BERT_LAYERS once scheduler is fast
|
||||
run: BENCHMARK_LOG=bert_10steps DEV=AMD CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
|
||||
- name: Remote
|
||||
run: |
|
||||
pkill -f 'extra/remote/serve.py' || true
|
||||
PYTHONPATH=. python3 extra/remote/serve.py 6482 &
|
||||
sleep 1
|
||||
DEBUG=2 PYTHONPATH=. REMOTE=127.0.0.1:6482 AM_RESET=1 DEV=AMD AMD_IFACE=PCI python3 test/test_tiny.py
|
||||
DEBUG=2 PYTHONPATH=. REMOTE=127.0.0.1:6482 AM_RESET=1 DEV=AMD AMD_AQL=1 AMD_IFACE=PCI python3 test/test_tiny.py
|
||||
pkill -f 'extra/remote/serve.py' || true
|
||||
run: BENCHMARK_LOG=bert_10steps AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee am_train_bert_one_gpu.txt
|
||||
- uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: Speed (AM Driver)
|
||||
path: |
|
||||
am_matmul_amd.txt
|
||||
am_train_cifar_one_gpu.txt
|
||||
am_train_resnet_one_gpu.txt
|
||||
am_train_bert_one_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
|
||||
|
||||
@@ -718,13 +743,11 @@ jobs:
|
||||
if: github.repository_owner == 'tinygrad'
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
uses: actions/checkout@v6
|
||||
- name: Setcap to python
|
||||
run: ./extra/amdpci/setup_python_cap.sh
|
||||
uses: actions/checkout@v4
|
||||
- name: Remove nv modules
|
||||
run: PYTHONPATH=. ./extra/hcq/hcq_smi.py nv rmmod
|
||||
run: ./extra/hcq/hcq_smi.py nv rmmod
|
||||
- name: Kill stale pids
|
||||
run: PYTHONPATH=. ./extra/hcq/hcq_smi.py nv kill_pids
|
||||
run: ./extra/hcq/hcq_smi.py nv kill_pids
|
||||
- name: Symlink models and datasets
|
||||
run: |
|
||||
mkdir -p weights
|
||||
@@ -743,30 +766,31 @@ jobs:
|
||||
- name: reset process replay
|
||||
run: test/external/process_replay/reset.py
|
||||
- name: Test driver start time
|
||||
run: time DEBUG=3 DEV=NV python3 test/test_tiny.py TestTiny.test_plus
|
||||
run: time DEBUG=3 NV=1 python3 test/test_tiny.py TestTiny.test_plus
|
||||
- name: Test tensor cores
|
||||
run: DEV=NV ALLOW_TF32=1 python3 test/opt/test_tensor_cores.py
|
||||
run: NV=1 ALLOW_TF32=1 python3 test/opt/test_tensor_cores.py
|
||||
- name: Test DISK copy time
|
||||
run: DEV=NV TESTFILE=/raid/downloads/llama3-8b-sfr/model-00001-of-00004.safetensors python3 test/external/external_benchmark_disk_raw.py
|
||||
run: NV=1 TESTFILE=/raid/downloads/llama3-8b-sfr/model-00001-of-00004.safetensors python3 test/external/external_benchmark_disk_raw.py
|
||||
- name: Test CPU copy time
|
||||
run: |
|
||||
DEV=NV GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyDefaulttoCPUJit
|
||||
DEV=NV GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
|
||||
NV=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyDefaulttoCPUJit
|
||||
NV=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
|
||||
- name: Test LLAMA-3
|
||||
run: BENCHMARK_LOG=llama3_beam DEV=NV JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --benchmark --temperature 0
|
||||
run: BENCHMARK_LOG=llama3_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --benchmark --temperature 0 | tee nv_llama3_beam.txt
|
||||
- name: Run full CIFAR training w 1 GPU
|
||||
run: time BENCHMARK_LOG=cifar DEV=NV DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
|
||||
run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee nv_train_cifar_one_gpu.txt
|
||||
- name: Run 10 MLPerf ResNet50 training steps (1 gpu)
|
||||
run: BENCHMARK_LOG=resnet_10steps DEV=NV MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
|
||||
run: BENCHMARK_LOG=resnet_10steps NV=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee nv_train_resnet_one_gpu.txt
|
||||
- name: Run 10 MLPerf Bert training steps (1 gpu)
|
||||
# TODO: remove BERT_LAYERS once scheduler is fast
|
||||
run: BENCHMARK_LOG=bert_10steps DEV=NV CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
|
||||
- name: Remote
|
||||
run: |
|
||||
pkill -f 'extra/remote/serve.py' || true
|
||||
PYTHONPATH=. python3 extra/remote/serve.py 6483 &
|
||||
sleep 1
|
||||
DEBUG=2 PYTHONPATH=. REMOTE=127.0.0.1:6483 DEV=NV python3 test/test_tiny.py
|
||||
pkill -f 'extra/remote/serve.py' || true
|
||||
run: BENCHMARK_LOG=bert_10steps NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee nv_train_bert_one_gpu.txt
|
||||
- uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: Speed (NV Driver)
|
||||
path: |
|
||||
nv_llama3_beam.txt
|
||||
nv_train_cifar_one_gpu.txt
|
||||
nv_train_resnet_one_gpu.txt
|
||||
nv_train_bert_one_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
|
||||
|
||||
@@ -14,7 +14,7 @@ jobs:
|
||||
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
uses: actions/checkout@v6
|
||||
uses: actions/checkout@v4
|
||||
- name: Remove amdgpu
|
||||
run: sudo rmmod amdgpu || true
|
||||
- name: Cleanup running AM processes
|
||||
@@ -22,10 +22,10 @@ jobs:
|
||||
- name: Run SDXL with new search
|
||||
# TODO: GCVM_L2_PROTECTION_FAULT_STATUS with llvm19
|
||||
run: |
|
||||
BENCHMARK_LOG=search_sdxl PYTHONPATH=. DEV=AMD JITBEAM=2 IGNORE_BEAM_CACHE=1 CCACHE=0 python examples/sdxl.py --noshow --timing --seed 0
|
||||
BENCHMARK_LOG=search_sdxl PYTHONPATH=. AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 CCACHE=0 python examples/sdxl.py --noshow --timing --seed 0
|
||||
- name: Run SDXL with cached search
|
||||
run: |
|
||||
BENCHMARK_LOG=search_sdxl_cached PYTHONPATH=. DEV=AMD JITBEAM=2 python examples/sdxl.py --noshow --timing --seed 0
|
||||
BENCHMARK_LOG=search_sdxl_cached PYTHONPATH=. AMD=1 JITBEAM=2 python examples/sdxl.py --noshow --timing --seed 0
|
||||
- name: Run winograd cifar with new search
|
||||
run: |
|
||||
BENCHMARK_LOG=search_wino_cifar WINO=1 DEFAULT_FLOAT=HALF JITBEAM=4 IGNORE_BEAM_CACHE=1 CCACHE=0 BS=1024 STEPS=500 python examples/hlb_cifar10.py
|
||||
|
||||
@@ -10,16 +10,16 @@ jobs:
|
||||
deploy:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
- name: Configure Git Credentials
|
||||
run: |
|
||||
git config user.name github-actions[bot]
|
||||
git config user.email 41898282+github-actions[bot]@users.noreply.github.com
|
||||
- uses: actions/setup-python@v6
|
||||
- uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: 3.x
|
||||
- run: echo "cache_id=$(date --utc '+%V')" >> $GITHUB_ENV
|
||||
- uses: actions/cache@v5
|
||||
- uses: actions/cache@v4
|
||||
with:
|
||||
key: mkdocs-material-${{ env.cache_id }}
|
||||
path: .cache
|
||||
|
||||
@@ -16,7 +16,7 @@ jobs:
|
||||
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
uses: actions/checkout@v6
|
||||
uses: actions/checkout@v4
|
||||
- name: Cleanup running AM processes
|
||||
run: python extra/amdpci/am_smi.py --pids --kill
|
||||
- name: Symlink datasets
|
||||
|
||||
@@ -12,9 +12,9 @@ jobs:
|
||||
deploy:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v6
|
||||
uses: actions/setup-python@v2
|
||||
with:
|
||||
python-version: '3.x'
|
||||
- name: Install dependencies
|
||||
|
||||
@@ -15,7 +15,7 @@ jobs:
|
||||
branchstat: ${{ steps.brstat.outputs.stat}}
|
||||
steps:
|
||||
- name: Check code from PR branch
|
||||
uses: actions/checkout@v6
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
repository: ${{ github.event.pull_request.head.repo.full_name }}
|
||||
ref: ${{ github.event.pull_request.head.sha }}
|
||||
@@ -46,18 +46,18 @@ jobs:
|
||||
if: needs.checkbranch.outputs.branchstat == 'false'
|
||||
steps:
|
||||
- name: Checkout code from PR branch
|
||||
uses: actions/checkout@v6
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
repository: ${{ github.event.pull_request.head.repo.full_name }}
|
||||
ref: ${{ github.event.pull_request.head.sha }}
|
||||
path: pr
|
||||
# the base default to tinygrad master and cannot be other fork branch for security purpose
|
||||
- name: Checkout code from tinygrad master
|
||||
uses: actions/checkout@v6
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
path: base
|
||||
- name: Set up Python 3.12
|
||||
uses: actions/setup-python@v6
|
||||
uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: '3.12'
|
||||
- name: Count Line Diff
|
||||
@@ -71,7 +71,7 @@ jobs:
|
||||
echo "EOF" >> "$GITHUB_ENV"
|
||||
- name: Comment Code Line Diff
|
||||
continue-on-error: false
|
||||
uses: marocchino/sticky-pull-request-comment@v3
|
||||
uses: marocchino/sticky-pull-request-comment@v2
|
||||
with:
|
||||
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
ignore_empty: true
|
||||
@@ -89,7 +89,7 @@ jobs:
|
||||
steps:
|
||||
- name: Comment Rebase
|
||||
continue-on-error: false
|
||||
uses: marocchino/sticky-pull-request-comment@v3
|
||||
uses: marocchino/sticky-pull-request-comment@v2
|
||||
with:
|
||||
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
skip_unchanged: true
|
||||
|
||||
+224
-357
File diff suppressed because it is too large
Load Diff
+1
-4
@@ -58,13 +58,10 @@ weights
|
||||
*.lprof
|
||||
comgr_*
|
||||
*.pkl
|
||||
!extra/sqtt/examples/**/*.pkl
|
||||
site/
|
||||
profile_stats
|
||||
*.log
|
||||
target
|
||||
.mypy_cache
|
||||
mutants
|
||||
.mutmut-cache
|
||||
dagre/
|
||||
graphlib/
|
||||
.mutmut-cache
|
||||
@@ -16,7 +16,7 @@ repos:
|
||||
pass_filenames: false
|
||||
- id: mypy
|
||||
name: mypy
|
||||
entry: python3 -m mypy
|
||||
entry: python3 -m mypy tinygrad/ --strict-equality
|
||||
language: system
|
||||
always_run: true
|
||||
pass_filenames: false
|
||||
@@ -28,7 +28,7 @@ repos:
|
||||
pass_filenames: false
|
||||
- id: tests
|
||||
name: comprehensive test suite
|
||||
entry: env OMP_NUM_THREADS=1 SKIP_SLOW_TEST=1 PYTHONPATH="." python3 -m pytest -n=6 test/backend/test_ops.py test/backend/test_schedule.py test/unit/test_assign.py test/backend/test_tensor.py test/backend/test_jit.py test/unit/test_schedule_cache.py test/null/test_pattern_matcher.py test/null/test_uop_symbolic.py test/unit/test_helpers.py
|
||||
entry: env OMP_NUM_THREADS=1 SKIP_SLOW_TEST=1 PYTHONPATH="." python3 -m pytest -n=6 test/test_ops.py test/test_schedule.py test/test_assign.py test/test_tensor.py test/test_jit.py test/unit/test_schedule_cache.py test/unit/test_pattern_matcher.py test/unit/test_uop_symbolic.py test/unit/test_helpers.py
|
||||
language: system
|
||||
always_run: true
|
||||
pass_filenames: false
|
||||
|
||||
@@ -0,0 +1,17 @@
|
||||
# tinygrad agents
|
||||
|
||||
Hello agent. You are one of the most talented programmers of your generation.
|
||||
|
||||
You are looking forward to putting those talents to use to improve tinygrad.
|
||||
|
||||
## philosophy
|
||||
|
||||
tinygrad is a **tensor** library focused on beauty and minimalism, while still matching the functionality of PyTorch and JAX.
|
||||
|
||||
Every line must earn its keep. Prefer readability over cleverness. We believe that if carefully designed, 10 lines can have the impact of 1000.
|
||||
|
||||
Never mix functionality changes with whitespace changes. All functionality changes must be tested.
|
||||
|
||||
## style
|
||||
|
||||
Use **2-space indentation**, and keep lines to a maximum of **150 characters**. Match the existing style.
|
||||
@@ -0,0 +1,227 @@
|
||||
# Claude Code Guide for tinygrad
|
||||
|
||||
## Architecture Overview
|
||||
|
||||
tinygrad compiles tensor operations into optimized kernels. The pipeline:
|
||||
|
||||
1. **Tensor** (`tensor.py`) - User-facing API, creates UOp graph
|
||||
2. **UOp** (`uop/ops.py`) - Unified IR for all operations (both tensor and kernel level)
|
||||
3. **Schedule** (`engine/schedule.py`, `schedule/`) - Converts tensor UOps to kernel UOps
|
||||
4. **Codegen** (`codegen/`) - Converts kernel UOps to device code
|
||||
5. **Runtime** (`runtime/`) - Device-specific execution
|
||||
|
||||
## Key Concepts
|
||||
|
||||
### UOp (Universal Operation)
|
||||
Everything is a UOp - tensors, operations, buffers, kernels. Key properties:
|
||||
- `op`: The operation type (Ops enum)
|
||||
- `dtype`: Data type
|
||||
- `src`: Tuple of source UOps
|
||||
- `arg`: Operation-specific argument
|
||||
- `tag`: Optional tag for graph transformations
|
||||
|
||||
UOps are **immutable and cached** - creating the same UOp twice returns the same object (ucache).
|
||||
|
||||
### PatternMatcher
|
||||
Used extensively for graph transformations:
|
||||
```python
|
||||
pm = PatternMatcher([
|
||||
(UPat(Ops.ADD, src=(UPat.cvar("x"), UPat.cvar("x"))), lambda x: x * 2),
|
||||
])
|
||||
result = graph_rewrite(uop, pm)
|
||||
```
|
||||
|
||||
### Schedule Cache
|
||||
Schedules are cached by graph structure. BIND nodes (variables with bound values) are unbound before cache key computation so different values hit the same cache.
|
||||
|
||||
## Directory Structure
|
||||
|
||||
```
|
||||
tinygrad/
|
||||
├── tensor.py # Tensor class, user API
|
||||
├── device.py # Buffer, device management
|
||||
├── dtype.py # Data types
|
||||
├── helpers.py # Utilities, environment vars
|
||||
├── uop/
|
||||
│ ├── ops.py # UOp class, Ops enum, PatternMatcher
|
||||
│ ├── spec.py # UOp type verification
|
||||
│ └── symbolic.py # Symbolic math simplification
|
||||
├── engine/
|
||||
│ ├── schedule.py # Schedule creation, caching
|
||||
│ ├── realize.py # Tensor realization
|
||||
│ ├── jit.py # JIT compilation
|
||||
│ └── memory.py # Memory planning
|
||||
├── schedule/
|
||||
│ ├── rangeify.py # Convert movements to ranges
|
||||
│ └── indexing.py # Index calculations
|
||||
├── codegen/
|
||||
│ ├── kernel.py # Kernel optimization
|
||||
│ └── uopgraph.py # UOp graph transformations
|
||||
├── renderer/ # Code generation (CUDA, Metal, etc.)
|
||||
└── runtime/ # Device backends
|
||||
```
|
||||
|
||||
## Testing
|
||||
|
||||
```bash
|
||||
# Run specific test
|
||||
python -m pytest test/unit/test_schedule_cache.py -xvs
|
||||
|
||||
# Run with timeout
|
||||
python -m pytest test/test_symbolic_ops.py -x --timeout=60
|
||||
|
||||
# Debug with print
|
||||
DEBUG=2 python -m pytest test/test_schedule.py::test_name -xvs
|
||||
|
||||
# Visualize UOp graphs
|
||||
VIZ=1 python -c "from tinygrad import Tensor; Tensor.ones(10).sum().realize()"
|
||||
```
|
||||
|
||||
## Common Environment Variables
|
||||
|
||||
- `DEBUG=1-4` - Increasing verbosity
|
||||
- `VIZ=1` - Enable graph visualization
|
||||
- `SPEC=1` - Enable UOp spec verification
|
||||
- `NOOPT=1` - Disable optimizations
|
||||
- `DEVICE=CPU/CUDA/AMD/METAL` - Set default device
|
||||
|
||||
## Debugging Tips
|
||||
|
||||
1. **Print UOp graphs**: `print(tensor.uop)` or `print(tensor.uop.sink())`
|
||||
2. **Check schedule**: `tensor.schedule()` returns list of ExecItems
|
||||
3. **Trace graph rewrites**: Use `VIZ=1` or add print in PatternMatcher callbacks
|
||||
4. **Find UOps by type**: `[u for u in uop.toposort() if u.op is Ops.SOMETHING]`
|
||||
|
||||
## Workflow Rules
|
||||
|
||||
- **NEVER commit without explicit user approval** - always show the diff and wait for approval
|
||||
- **NEVER amend commits** - always create a new commit instead
|
||||
- Run `pre-commit run --all-files` before committing to catch linting/type errors
|
||||
- Run tests before proposing commits
|
||||
- Test with `SPEC=2` when modifying UOp-related code
|
||||
|
||||
## Style Notes
|
||||
|
||||
- 2-space indentation, 150 char line limit
|
||||
- PatternMatchers should be defined at module level (slow to construct)
|
||||
- Prefer `graph_rewrite` over manual graph traversal
|
||||
- UOp methods like `.replace()` preserve tags unless explicitly changed
|
||||
- Use `.rtag(value)` to add tags to UOps
|
||||
|
||||
## Lessons Learned
|
||||
|
||||
### UOp ucache Behavior
|
||||
UOps are cached by their contents - creating a UOp with identical (op, dtype, src, arg) returns the **same object**. This means:
|
||||
- `uop.replace(tag=None)` on a tagged UOp returns the original untagged UOp if it exists in cache
|
||||
- Two UOps with same structure are identical (`is` comparison works)
|
||||
|
||||
### Spec Validation
|
||||
When adding new UOp patterns, update `tinygrad/uop/spec.py`. Test with:
|
||||
```bash
|
||||
SPEC=2 python3 test/unit/test_something.py
|
||||
```
|
||||
Spec issues appear as `RuntimeError: SPEC ISSUE None: UOp(...)`.
|
||||
|
||||
### Schedule Cache Key Normalization
|
||||
The schedule cache strips values from BIND nodes so different bound values (e.g., KV cache positions) hit the same cache entry:
|
||||
- `pm_pre_sched_cache`: BIND(DEFINE_VAR, CONST) → BIND(DEFINE_VAR) for cache key
|
||||
- `pm_post_sched_cache`: restores original BIND from context
|
||||
- When accessing `bind.src[1]`, check `len(bind.src) > 1` first (might be stripped)
|
||||
- Extract var_vals from `input_buffers` dict after graph_rewrite (avoids extra toposort)
|
||||
|
||||
### Avoiding Extra Work
|
||||
- Use ctx dict from graph_rewrite to collect info during traversal instead of separate toposort
|
||||
- Only extract var_vals when schedule is non-empty (no kernels = no vars needed)
|
||||
- PatternMatchers are slow to construct - define at module level, not in functions
|
||||
|
||||
### Readability Over Speed
|
||||
Don't add complexity for marginal performance gains. Simpler code that's slightly slower is often better:
|
||||
```python
|
||||
# BAD: "optimized" with extra complexity
|
||||
if has_afters: # skip toposort if no AFTERs
|
||||
after_map = [(u, u.buf_uop) for u in big_sink.toposort() if u.op is Ops.AFTER]
|
||||
|
||||
# GOOD: simple, always works
|
||||
after_map = [(u, u.buf_uop) for u in big_sink.toposort() if u.op is Ops.AFTER]
|
||||
```
|
||||
The conditional check adds complexity, potential bugs, and often negligible speedup. Only optimize when profiling shows a real bottleneck.
|
||||
|
||||
### Testing LLM Changes
|
||||
```bash
|
||||
# Quick smoke test
|
||||
echo "Hello" | DEBUG=1 python tinygrad/apps/llm.py --model "llama3.2:1b"
|
||||
|
||||
# Check cache hits (should see "cache hit" after warmup)
|
||||
echo "Hello world" | DEBUG=1 python tinygrad/apps/llm.py --model "llama3.2:1b" 2>&1 | grep cache
|
||||
|
||||
# Test with beam search
|
||||
echo "Hello" | BEAM=2 python tinygrad/apps/llm.py --model "llama3.2:1b"
|
||||
```
|
||||
|
||||
## Common Patterns
|
||||
|
||||
### Graph Transformation
|
||||
```python
|
||||
def my_transform(ctx, x):
|
||||
# Return new UOp or None to skip
|
||||
return x.replace(arg=new_arg)
|
||||
|
||||
pm = PatternMatcher([
|
||||
(UPat(Ops.SOMETHING, name="x"), my_transform),
|
||||
])
|
||||
result = graph_rewrite(input_uop, pm, ctx={})
|
||||
```
|
||||
|
||||
### Finding Variables
|
||||
```python
|
||||
# Get all variables in a UOp graph
|
||||
variables = uop.variables()
|
||||
|
||||
# Get bound variable values
|
||||
var, val = bind_uop.unbind()
|
||||
```
|
||||
|
||||
### Shape Handling
|
||||
```python
|
||||
# Shapes can be symbolic (contain UOps)
|
||||
shape = tensor.shape # tuple[sint, ...] where sint = int | UOp
|
||||
```
|
||||
|
||||
## Performance Optimization
|
||||
|
||||
When optimizing tinygrad internals:
|
||||
|
||||
1. **Measure wall time, not just call counts** - Reducing `graph_rewrite` calls doesn't always improve wall time. The overhead of conditional checks can exceed the cost of the operation being skipped.
|
||||
|
||||
2. **Profile each optimization individually** - Run benchmarks with and without each change to measure actual impact. Use `test/external/external_benchmark_schedule.py` for schedule/rewrite timing.
|
||||
|
||||
3. **Early exits in hot paths are effective** - Simple checks like `if self.op is Ops.CONST: return self` in `simplify()` can eliminate many unnecessary `graph_rewrite` calls.
|
||||
|
||||
4. **`graph_rewrite` is expensive** - Each call has overhead even for small graphs. Avoid calling it when the result is trivially known (e.g., simplifying a CONST returns itself).
|
||||
|
||||
5. **Beware iterator overhead** - Checks like `all(x.op is Ops.CONST for x in self.src)` can be slower than just running the operation, especially for small sequences.
|
||||
|
||||
6. **Verify cache hit rates before adding/keeping caches** - Measure actual hit rates with real workloads. A cache with 0% hit rate is pure overhead (e.g., `pm_cache` was removed because the algorithm guarantees each UOp is only passed to `pm_rewrite` once).
|
||||
|
||||
7. **Use `TRACK_MATCH_STATS=2` to profile pattern matching** - This shows match rates and time per pattern. Look for patterns with 0% match rate that still cost significant time - these are pure overhead for that workload.
|
||||
|
||||
8. **Cached properties beat manual traversal** - `backward_slice` uses `@functools.cached_property`. A DFS with early-exit sounds faster but is actually slower because it doesn't benefit from caching. The cache hit benefit often outweighs algorithmic improvements.
|
||||
|
||||
9. **Avoid creating intermediate objects in hot paths** - For example, `any(x.op in ops for x in self.backward_slice)` is faster than `any(x.op in ops for x in {self:None, **self.backward_slice})` because it avoids dict creation.
|
||||
|
||||
## Pattern Matching Profiling
|
||||
|
||||
Use `TRACK_MATCH_STATS=2` to identify expensive patterns:
|
||||
|
||||
```bash
|
||||
TRACK_MATCH_STATS=2 PYTHONPATH="." python3 test/external/external_benchmark_schedule.py
|
||||
```
|
||||
|
||||
Output format: `matches / attempts -- match_time / total_time ms -- location`
|
||||
|
||||
Key patterns to watch (from ResNet50 benchmark):
|
||||
- `split_load_store`: ~146ms, 31% match rate - does real work
|
||||
- `simplify_valid`: ~75ms, 0% match rate in this workload - checks AND ops for INDEX in backward slice
|
||||
- `vmin==vmax folding`: ~55ms, 0.33% match rate - checks 52K ops but rarely matches
|
||||
|
||||
Patterns with 0% match rate are workload-specific overhead. They may be useful in other workloads, so don't remove them without understanding their purpose.
|
||||
@@ -192,7 +192,7 @@ For more examples on how to run the full test suite please refer to the [CI work
|
||||
Some examples of running tests locally:
|
||||
```sh
|
||||
python3 -m pip install -e '.[testing]' # install extra deps for testing
|
||||
python3 test/backend/test_ops.py # just the ops tests
|
||||
python3 test/test_ops.py # just the ops tests
|
||||
python3 -m pytest test/ # whole test suite
|
||||
```
|
||||
|
||||
|
||||
@@ -3,7 +3,7 @@
|
||||
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 `DEV=AMD`!
|
||||
Make sure that amdgpu module is unloaded and just run tinygrad with `AMD=1`!
|
||||
|
||||
Optional requirements:
|
||||
|
||||
|
||||
@@ -10,7 +10,7 @@ Directories are listed in order of how they are processed.
|
||||
|
||||
Group UOps into kernels.
|
||||
|
||||
::: tinygrad.schedule.rangeify.get_kernel_graph
|
||||
::: tinygrad.schedule.rangeify.get_rangeify_map
|
||||
options:
|
||||
members: false
|
||||
show_labels: false
|
||||
@@ -28,7 +28,7 @@ Transforms the ast into an optimized ast. This is where BEAM search and heuristi
|
||||
|
||||
Transform the optimized ast into a linearized and rendered program.
|
||||
|
||||
::: tinygrad.codegen.get_program
|
||||
::: tinygrad.codegen.full_rewrite_to_program
|
||||
options:
|
||||
members: false
|
||||
show_labels: false
|
||||
|
||||
+7
-2
@@ -3,7 +3,7 @@
|
||||
This is a list of environment variable that control the runtime behavior of tinygrad and its examples.
|
||||
Most of these are self-explanatory, and are usually used to set an option at runtime.
|
||||
|
||||
Example: `DEV=CL DEBUG=4 python3 -m pytest`
|
||||
Example: `CL=1 DEBUG=4 python3 -m pytest`
|
||||
|
||||
However you can also decorate a function to set a value only inside that function.
|
||||
|
||||
@@ -31,7 +31,12 @@ 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)
|
||||
DEV | [AMD, NV, ...] | enable a specific backend
|
||||
CL | [1] | enable OpenCL backend
|
||||
CUDA | [1] | enable CUDA backend
|
||||
AMD | [1] | enable AMD backend
|
||||
NV | [1] | enable NV backend
|
||||
METAL | [1] | enable Metal backend (for Mac M1 and after)
|
||||
CPU | [1] | enable CPU 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
|
||||
|
||||
+2
-2
@@ -1,6 +1,6 @@
|
||||
# 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., `DEV=CPU`).
|
||||
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`).
|
||||
|
||||
| Runtime | Description | Compiler Options | Requirements |
|
||||
|---------|-------------|------------------|--------------|
|
||||
@@ -70,7 +70,7 @@ AMD backend supports several interfaces for communicating with devices:
|
||||
|
||||
* `KFD`: uses the amdgpu driver
|
||||
* `PCI`: uses the [AM driver](developer/am.md)
|
||||
* `USB`: USB3 interface for asm24xx chips.
|
||||
* `USB`: USB3 interafce for asm24xx chips.
|
||||
|
||||
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.
|
||||
|
||||
|
||||
@@ -6,7 +6,6 @@ Elementwise ops operate on a per element basis. They don't change the shape of t
|
||||
::: tinygrad.Tensor.neg
|
||||
::: tinygrad.Tensor.log
|
||||
::: tinygrad.Tensor.log2
|
||||
::: tinygrad.Tensor.log10
|
||||
::: tinygrad.Tensor.exp
|
||||
::: tinygrad.Tensor.exp2
|
||||
::: tinygrad.Tensor.sqrt
|
||||
@@ -88,8 +87,4 @@ Elementwise ops operate on a per element basis. They don't change the shape of t
|
||||
::: tinygrad.Tensor.float
|
||||
::: tinygrad.Tensor.half
|
||||
::: tinygrad.Tensor.int
|
||||
::: tinygrad.Tensor.bool
|
||||
::: tinygrad.Tensor.bfloat16
|
||||
::: tinygrad.Tensor.double
|
||||
::: tinygrad.Tensor.long
|
||||
::: tinygrad.Tensor.short
|
||||
::: tinygrad.Tensor.bool
|
||||
@@ -27,6 +27,5 @@
|
||||
::: tinygrad.Tensor.flatten
|
||||
::: tinygrad.Tensor.unflatten
|
||||
::: tinygrad.Tensor.diag
|
||||
::: tinygrad.Tensor.diagonal
|
||||
::: tinygrad.Tensor.roll
|
||||
::: tinygrad.Tensor.rearrange
|
||||
@@ -7,7 +7,6 @@
|
||||
::: tinygrad.Tensor.any
|
||||
::: tinygrad.Tensor.all
|
||||
::: tinygrad.Tensor.isclose
|
||||
::: tinygrad.Tensor.allclose
|
||||
::: tinygrad.Tensor.mean
|
||||
::: tinygrad.Tensor.var
|
||||
::: tinygrad.Tensor.var_mean
|
||||
@@ -31,9 +30,7 @@
|
||||
::: tinygrad.Tensor.matmul
|
||||
::: tinygrad.Tensor.einsum
|
||||
::: tinygrad.Tensor.cumsum
|
||||
::: tinygrad.Tensor.cumprod
|
||||
::: tinygrad.Tensor.cummax
|
||||
::: tinygrad.Tensor.cummin
|
||||
::: tinygrad.Tensor.triu
|
||||
::: tinygrad.Tensor.tril
|
||||
::: tinygrad.Tensor.interpolate
|
||||
@@ -41,9 +38,7 @@
|
||||
::: tinygrad.Tensor.scatter_reduce
|
||||
::: tinygrad.Tensor.masked_select
|
||||
::: tinygrad.Tensor.masked_fill
|
||||
::: tinygrad.Tensor.nonzero
|
||||
::: tinygrad.Tensor.sort
|
||||
::: tinygrad.Tensor.argsort
|
||||
::: tinygrad.Tensor.topk
|
||||
::: tinygrad.Tensor.multinomial
|
||||
|
||||
@@ -61,8 +56,3 @@
|
||||
::: tinygrad.Tensor.sparse_categorical_crossentropy
|
||||
::: tinygrad.Tensor.cross_entropy
|
||||
::: tinygrad.Tensor.nll_loss
|
||||
|
||||
## Linear Algebra
|
||||
|
||||
::: tinygrad.Tensor.qr
|
||||
::: tinygrad.Tensor.svd
|
||||
|
||||
@@ -1,196 +0,0 @@
|
||||
from tinygrad import Tensor, dtypes, Context, getenv, UOp, fetch
|
||||
from tinygrad.uop.ops import Ops, PatternMatcher, UPat
|
||||
from tinygrad.uop.symbolic import symbolic
|
||||
from tinygrad.codegen import Renderer
|
||||
from tinygrad.codegen.opt import Opt, OptOps
|
||||
|
||||
# ************************* implementation of the problem ************************
|
||||
|
||||
def myhash(a: Tensor) -> Tensor:
|
||||
a = (a + 0x7ED55D16) + (a << 12)
|
||||
a = (a ^ 0xC761C23C) ^ (a >> 19)
|
||||
a = (a + 0x165667B1) + (a << 5)
|
||||
a = (a + 0xD3A2646C) ^ (a << 9)
|
||||
a = (a + 0xFD7046C5) + (a << 3)
|
||||
a = (a ^ 0xB55A4F09) ^ (a >> 16)
|
||||
return a
|
||||
|
||||
def select_with_where_tree(values: Tensor, relative_idx: Tensor) -> Tensor:
|
||||
n = values.shape[0]
|
||||
if n == 1: return values[0].expand(relative_idx.shape)
|
||||
|
||||
mid = n // 2
|
||||
left = select_with_where_tree(values[:mid], relative_idx)
|
||||
right = select_with_where_tree(values[mid:], relative_idx - mid)
|
||||
|
||||
go_left = relative_idx < mid
|
||||
return go_left.where(left, right)
|
||||
|
||||
def tree_traversal(forest: Tensor, val: Tensor, height: int, rounds: int, where_tree_threshold=3) -> Tensor:
|
||||
# All walkers start at idx=0
|
||||
idx = Tensor.zeros(val.shape, device=val.device, dtype=dtypes.uint32)
|
||||
|
||||
for r in range(rounds):
|
||||
level = r % (height + 1)
|
||||
level_start = (1 << level) - 1
|
||||
level_size = 1 << level
|
||||
|
||||
if level == 0:
|
||||
# At root (level 0), all walkers are at idx=0
|
||||
# No gather needed, just broadcast the root value
|
||||
node_val = forest[0].expand(val.shape)
|
||||
idx = idx * 0 # Reset to 0
|
||||
elif level <= where_tree_threshold:
|
||||
# Small level: use where-tree
|
||||
level_values = forest[level_start : level_start + level_size]
|
||||
relative_idx = (idx - level_start)
|
||||
node_val = select_with_where_tree(level_values, relative_idx)
|
||||
else:
|
||||
# Large level: use gather
|
||||
node_val = forest.gather(0, idx)
|
||||
|
||||
val = myhash(val ^ node_val)
|
||||
idx = (idx << 1) + (1 + (val & 1))
|
||||
|
||||
# No wrap check needed! At round 10 (level becomes 0), we reset idx above.
|
||||
|
||||
return val.contiguous(arg=(Opt(OptOps.UPCAST, 0, 8),))
|
||||
|
||||
# ************************* renderer for VLIW machine *************************
|
||||
|
||||
def loop_unrolling(sink:UOp):
|
||||
rng = [x for x in sink.toposort() if x.op is Ops.RANGE]
|
||||
if len(rng) == 0: return None
|
||||
print(f"unrolling loop with size {rng[0].vmax+1}")
|
||||
unrolled_sinks = [sink.substitute({rng[0]:rng[0].const_like(i)}).src[0] for i in range(rng[0].vmax+1)]
|
||||
return UOp.sink(*unrolled_sinks, arg=sink.arg)
|
||||
|
||||
global_addrs = []
|
||||
vliw_prepare = PatternMatcher([
|
||||
# loop unrolling (should be a part of tinygrad)
|
||||
(UPat(Ops.SINK, name="sink"), loop_unrolling),
|
||||
# cast is fake
|
||||
(UPat(Ops.CAST, name="c"), lambda c: c.src[0]),
|
||||
# rewrites to hardcode the addresses in memory
|
||||
(UPat(Ops.PARAM, name="dg"), lambda dg: UOp.const(dtypes.uint, global_addrs[dg.arg])),
|
||||
# INDEX is just plus
|
||||
(UPat(Ops.INDEX, name="i"), lambda i: i.src[0]+i.src[1]),
|
||||
])+symbolic
|
||||
|
||||
class VLIWRenderer(Renderer):
|
||||
has_local = False # TODO: this should be the default / cleaned up
|
||||
# this says this backend supports MULACC + more. decompositions uses this
|
||||
code_for_op: dict = {Ops.MULACC: None, Ops.ADD: "+", Ops.MUL: "*",
|
||||
Ops.XOR: "^", Ops.AND: "&", Ops.OR: "|",
|
||||
Ops.SHL: "<<", Ops.SHR: ">>", Ops.CMPLT: "<"}
|
||||
# this matcher runs while still in graph form
|
||||
pre_matcher = vliw_prepare
|
||||
|
||||
def render(self, uops:list[UOp]):
|
||||
|
||||
# TODO: this is a minimal renderer. for low cycle count, make it good
|
||||
# to get speed, you need to add VLIW packing
|
||||
# to get under 1536 regs, you need to add a register allocator
|
||||
# we left the fun parts to you
|
||||
|
||||
print(f"rendering with {len(uops)} uops")
|
||||
reg, inst = 0, []
|
||||
r: dict[UOp, int] = {}
|
||||
for u in uops:
|
||||
assert u.dtype.count in (1,8), "dtype count must be 1 or 8"
|
||||
|
||||
# dumb register allocator
|
||||
if u.op not in {Ops.STORE, Ops.SINK, Ops.GEP}:
|
||||
r[u] = reg
|
||||
reg += u.dtype.count
|
||||
|
||||
# render UOps to instructions
|
||||
match u.op:
|
||||
case Ops.SINK:
|
||||
inst.append({"flow": [("halt",)]})
|
||||
case Ops.CONST:
|
||||
inst.append({"load": [("const", r[u], u.arg)]})
|
||||
case Ops.GEP:
|
||||
# a GEP is just an alias to a special register in the vector
|
||||
r[u] = r[u.src[0]] + u.arg[0]
|
||||
case Ops.VECTORIZE:
|
||||
if all(s == u.src[0] for s in u.src):
|
||||
# if all sources are the same, we can broadcast
|
||||
inst.append({"valu": [("vbroadcast", r[u], r[u.src[0]])]})
|
||||
else:
|
||||
# this is a copy into a contiguous chunk of registers
|
||||
inst.extend({"flow": [("add_imm", r[u]+i, r[s], 0)]} for i,s in enumerate(u.src) if r[s] != r[u]+i)
|
||||
case Ops.LOAD:
|
||||
op = "vload" if u.dtype.count > 1 else "load"
|
||||
inst.append({"load": [(op, r[u], r[u.src[0]])]})
|
||||
case Ops.STORE:
|
||||
op = "vstore" if u.src[1].dtype.count > 1 else "store"
|
||||
inst.append({"store": [(op, r[u.src[0]], r[u.src[1]])]})
|
||||
case Ops.MULACC:
|
||||
assert u.dtype.count == 8
|
||||
inst.append({"valu": [("multiply_add", r[u], r[u.src[0]], r[u.src[1]], r[u.src[2]])]})
|
||||
case Ops.WHERE:
|
||||
assert u.dtype.count == 8
|
||||
inst.append({"flow": [("vselect", r[u], r[u.src[0]], r[u.src[1]], r[u.src[2]])]})
|
||||
case _ if u.op in self.code_for_op:
|
||||
cat = "valu" if u.dtype.count > 1 else "alu"
|
||||
inst.append({cat: [(self.code_for_op[u.op], r[u], r[u.src[0]], r[u.src[1]])]})
|
||||
case _:
|
||||
raise NotImplementedError(f"unhandled op {u.op}")
|
||||
return repr(inst)
|
||||
|
||||
# ************************* test and render *************************
|
||||
|
||||
import sys, types
|
||||
PROBLEM_URL = "https://raw.githubusercontent.com/anthropics/original_performance_takehome/refs/heads/main/tests/frozen_problem.py"
|
||||
sys.modules["problem"] = problem = types.ModuleType("problem")
|
||||
exec(fetch(PROBLEM_URL).read_text(), problem.__dict__)
|
||||
|
||||
if __name__ == "__main__":
|
||||
batch_size = getenv("BS", 256)
|
||||
height = 10
|
||||
rounds = getenv("ROUNDS", 16)
|
||||
|
||||
# build problem
|
||||
tree = problem.Tree.generate(height)
|
||||
inp = problem.Input.generate(tree, batch_size, rounds)
|
||||
mem = problem.build_mem_image(tree, inp)
|
||||
global_addrs.extend([mem[6], mem[6], mem[4]]) # output, input, forest
|
||||
|
||||
# *** verify the kernel in tinygrad compared to reference ***
|
||||
|
||||
forest_t = Tensor(tree.values, dtype=dtypes.uint32)
|
||||
val_t = Tensor(inp.values, dtype=dtypes.uint32)
|
||||
|
||||
if getenv("VERIFY", 1):
|
||||
# verify on normal tinygrad device
|
||||
with Context(PCONTIG=2):
|
||||
out = tree_traversal(forest_t, val_t, height, rounds)
|
||||
val_out = out.tolist()
|
||||
problem.reference_kernel(tree, inp)
|
||||
assert val_out == inp.values
|
||||
print("verification passed")
|
||||
|
||||
# *** render to device ***
|
||||
|
||||
from tinygrad.codegen import get_program
|
||||
with Context(PCONTIG=2, DEVECTORIZE=2, SPEC=0):
|
||||
out = tree_traversal(forest_t, val_t, height, rounds)
|
||||
sink = out.schedule()[-1].ast
|
||||
prg = get_program(sink, VLIWRenderer())
|
||||
|
||||
# *** run on Machine and compare ***
|
||||
|
||||
# NOTE: the scratch size needs to be reduced to 1536 when you have a register allocator
|
||||
src = eval(prg.src)
|
||||
max_regs = max(t[1] for instr in src for v in instr.values() for t in v if len(t) > 1) + 8
|
||||
print(f"{max_regs:5d} regs used" + ("" if max_regs <= 1536 else " <-- WARNING: TOO MANY REGISTERS, MUST BE <= 1536"))
|
||||
machine = problem.Machine(mem, src, problem.DebugInfo(scratch_map={}), n_cores=1, trace=False, scratch_size=max_regs)
|
||||
machine.run()
|
||||
print(f"ran for {machine.cycle:5d} cycles" + ("" if machine.cycle <= 1363 else " <-- EVEN CLAUDE GOT 1363"))
|
||||
|
||||
# compare to reference
|
||||
ref_mem = mem.copy()
|
||||
for _ in problem.reference_kernel2(ref_mem, {}): pass
|
||||
assert machine.mem[mem[6]:mem[6]+mem[2]] == ref_mem[mem[6]:mem[6]+mem[2]]
|
||||
print("compare passed!")
|
||||
@@ -1,79 +0,0 @@
|
||||
from typing import Optional
|
||||
from tinygrad import Tensor
|
||||
from tinygrad.dtype import DTypeLike, dtypes
|
||||
import math
|
||||
|
||||
# rewritten from numpy
|
||||
def rfftfreq(n: int, d: float = 1.0, device=None) -> Tensor:
|
||||
val = 1.0 / (n * d)
|
||||
N = n // 2 + 1
|
||||
results = Tensor.arange(N, device=device)
|
||||
return results * val
|
||||
|
||||
# just like in librosa
|
||||
def fft_frequencies(sr: float, n_fft: int) -> Tensor:
|
||||
return rfftfreq(n=n_fft, d=1.0 / sr)
|
||||
|
||||
def hz_to_mel(freq: Tensor) -> Tensor:
|
||||
# linear part
|
||||
f_min = 0.0
|
||||
f_sp = 200.0 / 3
|
||||
mels = (freq - f_min) / f_sp
|
||||
|
||||
# log-scale part
|
||||
min_log_hz = 1000.0 # beginning of log region (Hz)
|
||||
mask = freq >= min_log_hz
|
||||
return mask.where(((min_log_hz - f_min) / f_sp) + (freq / min_log_hz).log() / (math.log(6.4) / 27.0), mels)
|
||||
|
||||
def mel_to_hz(mels: Tensor) -> Tensor:
|
||||
# linear scale
|
||||
f_min = 0.0
|
||||
f_sp = 200.0 / 3
|
||||
freqs = f_min + f_sp * mels
|
||||
|
||||
# nonlinear scale
|
||||
min_log_hz = 1000.0 # beginning of log region (Hz)
|
||||
min_log_mel = (min_log_hz - f_min) / f_sp # same (Mels)
|
||||
logstep = math.log(6.4) / 27.0 # step size for log region
|
||||
|
||||
log_t = mels >= min_log_mel
|
||||
freqs = log_t.where(min_log_hz * ((logstep * (mels - min_log_mel)).exp()), freqs)
|
||||
return freqs
|
||||
|
||||
def mel_frequencies(n_mels: int = 128, *, fmin: float = 0.0, fmax: float = 11025.0) -> Tensor:
|
||||
# center freqs of mel bands - uniformly spaced between limits
|
||||
min_max_mel = hz_to_mel(Tensor([fmin, fmax]))
|
||||
|
||||
mels = Tensor.linspace(min_max_mel[0], min_max_mel[1], n_mels)
|
||||
hz = mel_to_hz(mels)
|
||||
return hz
|
||||
|
||||
def mel(
|
||||
*,
|
||||
sr: float,
|
||||
n_fft: int,
|
||||
n_mels: int = 128,
|
||||
fmin: float = 0.0,
|
||||
fmax: Optional[float] = None,
|
||||
dtype: DTypeLike = dtypes.default_float,
|
||||
) -> Tensor:
|
||||
if fmax is None:
|
||||
fmax = float(sr) / 2
|
||||
|
||||
n_mels = int(n_mels)
|
||||
|
||||
fftfreqs = fft_frequencies(sr=sr, n_fft=n_fft) # center freqs of each FFT bin
|
||||
mel_f = mel_frequencies(n_mels + 2, fmin=fmin, fmax=fmax) # center freqs of mel bands
|
||||
|
||||
fdiff = mel_f[1:] - mel_f[:-1]
|
||||
ramps = mel_f[None].T.expand(-1, fftfreqs.shape[-1]) - fftfreqs
|
||||
|
||||
lower = -ramps[:n_mels] / fdiff[:n_mels][None].T
|
||||
upper = ramps[2 : n_mels + 2] / fdiff[1 : n_mels + 1][None].T
|
||||
weights = lower.minimum(upper).maximum(0)
|
||||
|
||||
# Slaney-style mel is scaled to be approx constant energy per channel
|
||||
enorm = 2.0 / (mel_f[2 : n_mels + 2] - mel_f[:n_mels])
|
||||
weights *= enorm[:, None]
|
||||
|
||||
return weights
|
||||
+14
-15
@@ -1,6 +1,6 @@
|
||||
# model based off https://medium.com/data-science/going-beyond-99-mnist-handwritten-digits-recognition-cfff96337392
|
||||
from typing import Callable
|
||||
from tinygrad import Tensor, TinyJit, nn, GlobalCounters, function
|
||||
from tinygrad import Tensor, TinyJit, nn, GlobalCounters
|
||||
from tinygrad.helpers import getenv, colored, trange
|
||||
from tinygrad.nn.datasets import mnist
|
||||
|
||||
@@ -15,31 +15,30 @@ class Model:
|
||||
nn.BatchNorm(64), Tensor.max_pool2d,
|
||||
lambda x: x.flatten(1), nn.Linear(576, 10)]
|
||||
|
||||
@function
|
||||
def __call__(self, x:Tensor) -> Tensor: return x.sequential(self.layers)
|
||||
|
||||
@TinyJit
|
||||
@Tensor.train()
|
||||
def train_step(self, X_train:Tensor, Y_train:Tensor) -> Tensor:
|
||||
opt.zero_grad()
|
||||
samples = Tensor.randint(getenv("BS", 512), high=X_train.shape[0])
|
||||
loss = self(X_train[samples]).sparse_categorical_crossentropy(Y_train[samples]).backward()
|
||||
return loss.realize(*opt.schedule_step())
|
||||
|
||||
@TinyJit
|
||||
def get_test_acc(self, X_test:Tensor, Y_test:Tensor) -> Tensor: return (self(X_test).argmax(axis=1) == Y_test).mean()*100
|
||||
|
||||
if __name__ == "__main__":
|
||||
X_train, Y_train, X_test, Y_test = mnist(fashion=getenv("FASHION"))
|
||||
|
||||
model = Model()
|
||||
opt = (nn.optim.Muon if getenv("MUON") else nn.optim.SGD if getenv("SGD") else nn.optim.Adam)(nn.state.get_parameters(model))
|
||||
|
||||
@TinyJit
|
||||
@Tensor.train()
|
||||
def train_step() -> Tensor:
|
||||
opt.zero_grad()
|
||||
samples = Tensor.randint(getenv("BS", 512), high=X_train.shape[0])
|
||||
loss = model(X_train[samples]).sparse_categorical_crossentropy(Y_train[samples]).backward()
|
||||
return loss.realize(*opt.schedule_step())
|
||||
|
||||
@TinyJit
|
||||
def get_test_acc() -> Tensor: return (model(X_test).argmax(axis=1) == Y_test).mean()*100
|
||||
|
||||
test_acc = float('nan')
|
||||
for i in (t:=trange(getenv("STEPS", 70))):
|
||||
GlobalCounters.reset() # NOTE: this makes it nice for DEBUG=2 timing
|
||||
loss = model.train_step(X_train, Y_train)
|
||||
if i%10 == 9: test_acc = model.get_test_acc(X_test, Y_test).item()
|
||||
loss = train_step()
|
||||
if i%10 == 9: test_acc = get_test_acc().item()
|
||||
t.set_description(f"loss: {loss.item():6.2f} test_accuracy: {test_acc:5.2f}%")
|
||||
|
||||
# verify eval acc
|
||||
|
||||
@@ -5,7 +5,7 @@ from extra.onnx_helpers import get_example_inputs, validate
|
||||
|
||||
def load_onnx_model(onnx_file):
|
||||
run_onnx = OnnxRunner(onnx_file)
|
||||
run_onnx_jit = TinyJit(lambda **kwargs: next(iter(run_onnx({k:v.to(None) for k,v in kwargs.items()}).values())), prune=True)
|
||||
run_onnx_jit = TinyJit(lambda **kwargs: next(iter(run_onnx({k:v.to(None) for k,v in kwargs.items()}).values())), prune=True, optimize=True)
|
||||
return run_onnx_jit, run_onnx.graph_inputs
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -1,10 +1,9 @@
|
||||
from pathlib import Path
|
||||
from extra.models.efficientnet import EfficientNet
|
||||
from tinygrad.tensor import Tensor
|
||||
from tinygrad.device import Device
|
||||
from tinygrad.nn.state import get_state_dict, safe_save, safe_load, load_state_dict
|
||||
from extra.export_model import export_model
|
||||
from tinygrad.helpers import fetch
|
||||
from tinygrad.helpers import getenv, fetch
|
||||
import ast
|
||||
|
||||
if __name__ == "__main__":
|
||||
@@ -13,13 +12,13 @@ if __name__ == "__main__":
|
||||
dirname = Path(__file__).parent
|
||||
# exporting a model that's loaded from safetensors doesn't work without loading in from safetensors first
|
||||
# loading the state dict from a safetensor file changes the generated kernels
|
||||
if Device.DEFAULT == "WEBGPU":
|
||||
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 Device.DEFAULT == "CPU" else "webgpu" if Device.DEFAULT == "WEBGPU" else ""
|
||||
mode = "clang" if getenv("CPU", "") != "" else "webgpu" if getenv("WEBGPU", "") != "" else ""
|
||||
prg, inp_sizes, out_sizes, state = export_model(model, mode, Tensor.randn(1,3,224,224))
|
||||
if Device.DEFAULT != "CPU":
|
||||
ext = "js" if Device.DEFAULT == "WEBGPU" else "json"
|
||||
if getenv("CPU", "") == "":
|
||||
ext = "js" if getenv("WEBGPU", "") != "" else "json"
|
||||
with open(dirname / f"net.{ext}", "w") as text_file:
|
||||
text_file.write(prg)
|
||||
else:
|
||||
@@ -69,6 +68,6 @@ if __name__ == "__main__":
|
||||
else printf("%s\\n", lbls[best_idx]);
|
||||
}""")
|
||||
|
||||
# DEV=CPU python3 examples/compile_efficientnet.py | clang -O2 -lm -x c - -o recognize && DEBUG=1 time ./recognize docs/showcase/stable_diffusion_by_tinygrad.jpg
|
||||
# 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
|
||||
# category : 281 (tabby, tabby cat) with 9.452788
|
||||
print('\n'.join(cprog))
|
||||
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 286 KiB |
@@ -19,8 +19,8 @@ cifar_std = [0.24703225141799082, 0.24348516474564, 0.26158783926049628]
|
||||
BS, STEPS = getenv("BS", 512), getenv("STEPS", 1000)
|
||||
EVAL_BS = getenv("EVAL_BS", BS)
|
||||
GPUS = [f'{Device.DEFAULT}:{i}' for i in range(getenv("GPUS", 1))]
|
||||
assert BS % len(GPUS) == 0, f"{BS=} is not a multiple of {len(GPUS)=}"
|
||||
assert EVAL_BS % len(GPUS) == 0, f"{EVAL_BS=} is not a multiple of {len(GPUS)=}"
|
||||
assert BS % len(GPUS) == 0, f"{BS=} is not a multiple of {len(GPUS)=}, uneven multi GPU is slow"
|
||||
assert EVAL_BS % len(GPUS) == 0, f"{EVAL_BS=} is not a multiple of {len(GPUS)=}, uneven multi GPU is slow"
|
||||
|
||||
class UnsyncedBatchNorm:
|
||||
def __init__(self, sz:int, eps=1e-5, affine=True, track_running_stats=True, momentum=0.1, num_devices=len(GPUS)):
|
||||
|
||||
@@ -65,7 +65,17 @@ def loader_process(q_in, q_out, X:Tensor, seed):
|
||||
else:
|
||||
# pad data with training mean
|
||||
img = np.tile(np.array([[[123.68, 116.78, 103.94]]], dtype=np.uint8), (224, 224, 1))
|
||||
X[idx].flatten().assign(img.tobytes())
|
||||
|
||||
# broken out
|
||||
#img_tensor = Tensor(img.tobytes(), device='CPU')
|
||||
#storage_tensor = X[idx].contiguous().realize().lazydata.base.realized
|
||||
#storage_tensor._copyin(img_tensor.numpy())
|
||||
|
||||
# faster
|
||||
X[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = img.tobytes()
|
||||
|
||||
# ideal
|
||||
#X[idx].assign(img.tobytes()) # NOTE: this is slow!
|
||||
q_out.put(idx)
|
||||
q_out.put(None)
|
||||
|
||||
@@ -203,13 +213,12 @@ class InterleavedDataset:
|
||||
self.queues[queue_index].queue.extend(load_file(file))
|
||||
|
||||
# Reference: https://github.com/mlcommons/training/blob/1c8a098ae3e70962a4f7422c0b0bd35ae639e357/language_model/tensorflow/bert/run_pretraining.py, Line 394
|
||||
def batch_load_train_bert(BS:int, seed:int|None=None):
|
||||
def batch_load_train_bert(BS:int):
|
||||
from extra.datasets.wikipedia import get_wiki_train_files
|
||||
rng = random.Random(seed)
|
||||
fs = sorted(get_wiki_train_files())
|
||||
train_files = []
|
||||
while fs: # TF shuffle
|
||||
rng.shuffle(fs)
|
||||
random.shuffle(fs)
|
||||
train_files.append(fs.pop(0))
|
||||
|
||||
cycle_length = min(getenv("NUM_CPU_THREADS", min(os.cpu_count(), 8)), len(train_files))
|
||||
@@ -254,8 +263,8 @@ def load_unet3d_data(preprocessed_dataset_dir, seed, queue_in, queue_out, X:Tens
|
||||
x = random_brightness_augmentation(x)
|
||||
x = gaussian_noise(x)
|
||||
|
||||
X[idx].flatten().assign(x.tobytes())
|
||||
Y[idx].flatten().assign(y.tobytes())
|
||||
X[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = x.tobytes()
|
||||
Y[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = y.tobytes()
|
||||
|
||||
queue_out.put(idx)
|
||||
queue_out.put(None)
|
||||
@@ -369,12 +378,12 @@ def load_retinanet_data(base_dir:Path, val:bool, queue_in:Queue, queue_out:Queue
|
||||
clipped_match_idxs = np.clip(match_idxs, 0, None)
|
||||
clipped_boxes, clipped_labels = tgt["boxes"][clipped_match_idxs], tgt["labels"][clipped_match_idxs]
|
||||
|
||||
boxes[idx].flatten().assign(clipped_boxes.tobytes())
|
||||
labels[idx].flatten().assign(clipped_labels.tobytes())
|
||||
matches[idx].flatten().assign(match_idxs.tobytes())
|
||||
anchors[idx].flatten().assign(anchor.tobytes())
|
||||
boxes[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = clipped_boxes.tobytes()
|
||||
labels[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = clipped_labels.tobytes()
|
||||
matches[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = match_idxs.tobytes()
|
||||
anchors[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = anchor.tobytes()
|
||||
|
||||
imgs[idx].flatten().assign(img.tobytes())
|
||||
imgs[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = img.tobytes()
|
||||
|
||||
queue_out.put(idx)
|
||||
queue_out.put(None)
|
||||
@@ -396,7 +405,6 @@ def batch_load_retinanet(dataset, val:bool, base_dir:Path, batch_size:int=32, sh
|
||||
queue_in.put((idx, img, tgt))
|
||||
|
||||
def _setup_shared_mem(shm_name:str, size:tuple[int, ...], dtype:dtypes) -> tuple[shared_memory.SharedMemory, Tensor]:
|
||||
shm_name = f"{shm_name}_{os.getpid()}"
|
||||
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}")
|
||||
@@ -543,7 +551,7 @@ class BinIdxDataset:
|
||||
version, = struct.unpack("<Q", self.idx.read(8))
|
||||
assert version == 1, "unsupported index version"
|
||||
dtype_code, = struct.unpack("<B", self.idx.read(1))
|
||||
self.dtype = {1:np.dtype(np.uint8), 2:np.dtype(np.int8), 3:np.dtype(np.int16), 4:np.dtype(np.int32), 5:np.dtype(np.int64), 6:np.dtype(np.float64), 7:np.dtype(np.double), 8:np.dtype(np.uint16)}[dtype_code]
|
||||
self.dtype = {1:dtypes.uint8, 2:dtypes.int8, 3:dtypes.int16, 4:dtypes.int32, 5:dtypes.int64, 6:dtypes.float64, 7:dtypes.double, 8:dtypes.uint16}[dtype_code]
|
||||
self.count, = struct.unpack("<Q", self.idx.read(8))
|
||||
doc_count, = struct.unpack("<Q", self.idx.read(8))
|
||||
|
||||
@@ -560,7 +568,7 @@ class BinIdxDataset:
|
||||
self.doc_idx = self.idx_t[start:end].bitcast(dtypes.int64).numpy()
|
||||
|
||||
# bin file
|
||||
self.bin_t = Tensor(base_path.with_name(f"{base_path.name}.bin")).numpy()
|
||||
self.bin_t = Tensor(base_path.with_name(f"{base_path.name}.bin"))
|
||||
|
||||
def _index(self, idx) -> tuple[int, int]:
|
||||
return int(self.pointers[idx]), int(self.sizes[idx])
|
||||
@@ -569,7 +577,7 @@ class BinIdxDataset:
|
||||
ptr, size = self._index(idx)
|
||||
if length is None: length = size - offset
|
||||
ptr += offset * self.dtype.itemsize
|
||||
return self.bin_t[ptr:ptr+length*self.dtype.itemsize].view(self.dtype)
|
||||
return self.bin_t[ptr:ptr+length*self.dtype.itemsize].bitcast(self.dtype).to(None)
|
||||
|
||||
# https://docs.nvidia.com/megatron-core/developer-guide/latest/api-guide/datasets.html
|
||||
class GPTDataset:
|
||||
@@ -628,7 +636,7 @@ class GPTDataset:
|
||||
sample_parts.append(self.indexed_dataset.get(int(self.doc_idx[i]), offset=int(offset), length=length))
|
||||
|
||||
# concat all parts
|
||||
text = np.concatenate(sample_parts, axis=0)
|
||||
text = Tensor.cat(*sample_parts)
|
||||
|
||||
return text
|
||||
|
||||
@@ -755,27 +763,48 @@ class BlendedGPTDataset:
|
||||
|
||||
return dataset_idx, dataset_sample_idx
|
||||
|
||||
def get_llama3_dataset(samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True, small:bool=False) -> BlendedGPTDataset:
|
||||
if small:
|
||||
if val:
|
||||
return BlendedGPTDataset(
|
||||
[base_dir / "c4-validation-91205-samples.en_text_document"], [1.0], samples, seqlen, seed, shuffle=False)
|
||||
return BlendedGPTDataset(
|
||||
[base_dir / "c4-train.en_6_text_document"], [1.0], samples, seqlen, seed, shuffle=True)
|
||||
def batch_load_llama3(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True):
|
||||
if val:
|
||||
return BlendedGPTDataset(
|
||||
[base_dir / "validation" / "c4-validationn-91205-samples.en_text_document"], [1.0], samples, seqlen, seed, shuffle=False)
|
||||
return BlendedGPTDataset(
|
||||
[base_dir / "c4-train.en_6_text_document", base_dir / "c4-train.en_7_text_document"], [1.0, 1.0], samples, seqlen, seed, shuffle=True)
|
||||
dataset = BlendedGPTDataset([
|
||||
base_dir / "validation" / "c4-validationn-91205-samples.en_text_document",
|
||||
], [
|
||||
1.0
|
||||
], samples, seqlen, seed, False)
|
||||
else:
|
||||
dataset = BlendedGPTDataset([
|
||||
base_dir / "c4-train.en_6_text_document",
|
||||
base_dir / "c4-train.en_7_text_document",
|
||||
], [
|
||||
1.0, 1.0
|
||||
], samples, seqlen, seed, True)
|
||||
|
||||
def iterate_llama3_dataset(dataset:BlendedGPTDataset, bs:int):
|
||||
for b in range(math.ceil(dataset.samples / bs)):
|
||||
batch = [dataset.get(b * bs + i) for i in range(bs)]
|
||||
stacked = np.stack(batch, axis=0)
|
||||
yield Tensor(stacked, device="NPY")
|
||||
for b in range(math.ceil(samples / bs)):
|
||||
batch = []
|
||||
for i in range(bs):
|
||||
tokens = dataset.get(b * bs + i)
|
||||
batch.append(tokens)
|
||||
yield Tensor.stack(batch, dim=0)
|
||||
|
||||
def batch_load_llama3(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True, small:bool=False):
|
||||
return iterate_llama3_dataset(get_llama3_dataset(samples, seqlen, base_dir, seed, val, small), bs)
|
||||
def batch_load_llama3_small(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True):
|
||||
if val:
|
||||
dataset = BlendedGPTDataset([
|
||||
base_dir / "c4-validation-91205-samples.en_text_document",
|
||||
], [
|
||||
1.0
|
||||
], samples, seqlen, seed, False)
|
||||
else:
|
||||
dataset = BlendedGPTDataset([
|
||||
base_dir / "c4-train.en_6_text_document",
|
||||
], [
|
||||
1.0
|
||||
], samples, seqlen, seed, True)
|
||||
|
||||
for b in range(math.ceil(samples / bs)):
|
||||
batch = []
|
||||
for i in range(bs):
|
||||
tokens = dataset.get(b * bs + i)
|
||||
batch.append(tokens)
|
||||
yield Tensor.stack(batch, dim=0)
|
||||
|
||||
if __name__ == "__main__":
|
||||
def load_unet3d(val):
|
||||
|
||||
@@ -219,18 +219,7 @@ 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
|
||||
model = BertForPretraining(**config)
|
||||
if getenv("FP8_TRAIN"):
|
||||
from extra.fp8.fp8_linear import convert_to_float8_training
|
||||
def module_filter_fn(mod, fqn):
|
||||
if isinstance(mod, LinearBert):
|
||||
skip_layers = [] if (ln:=config["num_hidden_layers"]) <= 2 else ["bert.encoder.layer.0.", f"bert.encoder.layer.{ln-1}"]
|
||||
if mod.weight.shape[-1] >= 1024 and "encoder" in fqn and not any(name in fqn for name in skip_layers):
|
||||
print(f"replacing linear with fp8: {fqn} {mod.weight.shape}")
|
||||
return True
|
||||
return False
|
||||
convert_to_float8_training(model, module_filter_fn)
|
||||
return model
|
||||
return BertForPretraining(**config)
|
||||
|
||||
def get_fake_data_bert(BS:int):
|
||||
return {
|
||||
|
||||
@@ -234,9 +234,12 @@ def eval_llama3():
|
||||
loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
|
||||
return loss.flatten().float()
|
||||
|
||||
from examples.mlperf.dataloader import get_llama3_dataset, iterate_llama3_dataset
|
||||
eval_dataset = get_llama3_dataset(5760, SEQLEN, BASEDIR, val=True, small=bool(SMALL))
|
||||
iter = iterate_llama3_dataset(eval_dataset, BS)
|
||||
if SMALL:
|
||||
from examples.mlperf.dataloader import batch_load_llama3_small
|
||||
iter = batch_load_llama3_small(BS, 5760, SEQLEN, BASEDIR, val=True)
|
||||
else:
|
||||
from examples.mlperf.dataloader import batch_load_llama3
|
||||
iter = batch_load_llama3(BS, 5760, SEQLEN, BASEDIR, val=True)
|
||||
|
||||
losses = []
|
||||
for tokens in tqdm(iter, total=5760//BS):
|
||||
|
||||
+109
-179
@@ -3,7 +3,7 @@ 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, Profiling, profile_marker, DEBUG
|
||||
from tinygrad.helpers import getenv, BEAM, WINO, round_up, diskcache_clear, Profiling
|
||||
from tinygrad.nn.state import get_parameters, get_state_dict, load_state_dict, safe_load, safe_save
|
||||
from tinygrad.nn.optim import LAMB, LARS, SGD, OptimizerGroup, Adam, AdamW
|
||||
|
||||
@@ -1008,7 +1008,6 @@ def train_bert():
|
||||
config["DISABLE_DROPOUT"] = getenv("DISABLE_DROPOUT", 0)
|
||||
config["TRAIN_BEAM"] = TRAIN_BEAM = getenv("TRAIN_BEAM", BEAM.value)
|
||||
config["EVAL_BEAM"] = EVAL_BEAM = getenv("EVAL_BEAM", BEAM.value)
|
||||
config["FP8_TRAIN"] = getenv("FP8_TRAIN", 0)
|
||||
|
||||
Tensor.manual_seed(seed) # seed for weight initialization
|
||||
|
||||
@@ -1086,7 +1085,7 @@ def train_bert():
|
||||
if RUNMLPERF:
|
||||
# only load real data with RUNMLPERF
|
||||
eval_it = iter(batch_load_val_bert(EVAL_BS))
|
||||
train_it = iter(tqdm(batch_load_train_bert(BS, seed=seed), total=train_steps, disable=BENCHMARK))
|
||||
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
|
||||
@@ -1148,7 +1147,7 @@ def train_bert():
|
||||
|
||||
device_str = parameters[0].device if isinstance(parameters[0].device, str) else f"{parameters[0].device[0]} * {len(parameters[0].device)}"
|
||||
loss = loss.item()
|
||||
if not getenv("FP8_TRAIN"): assert not math.isnan(loss)
|
||||
assert not math.isnan(loss)
|
||||
lr = lr.item()
|
||||
|
||||
cl = time.perf_counter()
|
||||
@@ -1161,7 +1160,7 @@ def train_bert():
|
||||
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/mem":GlobalCounters.mem_used / 1e9, "train/GFLOPS": GlobalCounters.global_ops * 1e-9 / (cl - st), "epoch": (i+1)*GBS})
|
||||
"train/GFLOPS": GlobalCounters.global_ops * 1e-9 / (cl - st), "epoch": (i+1)*GBS})
|
||||
|
||||
train_data, next_data = next_data, None
|
||||
i += 1
|
||||
@@ -1282,34 +1281,26 @@ def train_bert():
|
||||
previous_step = i
|
||||
|
||||
def train_llama3():
|
||||
from examples.mlperf.models.flat_llama import FlatTransformer
|
||||
from extra.models.llama import Transformer
|
||||
from examples.llama3 import MODEL_PARAMS
|
||||
from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup
|
||||
from examples.mlperf.optim import GradAccClipAdamW
|
||||
|
||||
BENCHMARK = getenv("BENCHMARK")
|
||||
|
||||
config = {}
|
||||
BASEDIR = config["BASEDIR"] = Path(getenv("BASEDIR", "/raid/datasets/c4/"))
|
||||
BS = config["BS"] = getenv("BS", 16)
|
||||
grad_acc = config["GRADIENT_ACC_STEPS"] = getenv("GRADIENT_ACC_STEPS", 1)
|
||||
assert grad_acc == 1, f"{grad_acc=} is not supported"
|
||||
GBS = config["GLOBAL_BATCH_SIZE"] = BS * grad_acc
|
||||
SEED = config["SEED"] = getenv("SEED", 5760)
|
||||
DATA_SEED = config["DATA_SEED"] = getenv("DATA_SEED", SEED)
|
||||
SEQLEN = config["SEQLEN"] = getenv("SEQLEN", 8192)
|
||||
TRAIN_ON_VAL = config["TRAIN_ON_VAL"] = getenv("TRAIN_ON_VAL", 0)
|
||||
SMALL = config["SMALL"] = getenv("SMALL", 0)
|
||||
SAMPLES = config["SAMPLES"] = getenv("SAMPLES", 5_760 if TRAIN_ON_VAL else 1_200_000 * 1152)
|
||||
EVAL_SAMPLES = config["EVAL_SAMPLES"] = getenv("EVAL_SAMPLES", 5760 if not SMALL else 1024)
|
||||
MAX_STEPS = config["MAX_STEPS"] = getenv("MAX_STEPS", math.ceil(1_200_000 * 1152 / GBS))
|
||||
WARMUP_STEPS = config["WARMUP_STEPS"] = getenv("WARMUP_STEPS", math.ceil(8000 * 1152 / GBS))
|
||||
LR = config["LR"] = getenv("LR", 8e-5 * GBS / 1152)
|
||||
END_LR = config["END_LR"] = getenv("END_LR", 8e-7)
|
||||
EVAL_FREQ = config["EVAL_FREQ"] = getenv("EVAL_FREQ", 46080)
|
||||
EVAL_BS = config["EVAL_BS"] = getenv("EVAL_BS", 16)
|
||||
EVAL_TARGET = config["EVAL_TARGET"] = getenv("EVAL_TARGET", 5.6)
|
||||
|
||||
# LR=1e-4 TRAIN_ON_VAL=1 DEFAULT_FLOAT=bfloat16 JITBEAM=2 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B WARMUP_STEPS=36 DECAY_STEPS=360 SEQLEN=512 PYTHONPATH=. DEV=AMD AMD_LLVM=0 MODEL=llama3 python3 examples/mlperf/model_train.py
|
||||
# LR=1e-4 TRAIN_ON_VAL=1 DEFAULT_FLOAT=bfloat16 JITBEAM=2 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B WARMUP_STEPS=36 DECAY_STEPS=360 SEQLEN=512 PYTHONPATH=. AMD=1 AMD_LLVM=0 MODEL=llama3 python3 examples/mlperf/model_train.py
|
||||
# trains to 7
|
||||
|
||||
opt_adamw_beta_1 = 0.9
|
||||
@@ -1318,61 +1309,50 @@ def train_llama3():
|
||||
opt_adamw_weight_decay = 0.1
|
||||
|
||||
opt_gradient_clip_norm = 1.0
|
||||
opt_learning_rate_warmup_steps = WARMUP_STEPS
|
||||
opt_learning_rate_decay_steps = MAX_STEPS - opt_learning_rate_warmup_steps
|
||||
opt_base_learning_rate = LR
|
||||
opt_end_learning_rate = END_LR
|
||||
|
||||
Tensor.manual_seed(SEED) # seed for weight initialization
|
||||
|
||||
# ** init wandb **
|
||||
WANDB = getenv("WANDB")
|
||||
if WANDB:
|
||||
import wandb
|
||||
wandb_args = {"id": wandb_id, "resume": "must"} if (wandb_id := getenv("WANDB_RESUME", "")) else {}
|
||||
wandb.init(config=config, **wandb_args, project="MLPerf-LLaMA3")
|
||||
opt_learning_rate_warmup_steps = getenv("WARMUP_STEPS", math.ceil(8000 * 1152 / GBS))
|
||||
opt_learning_rate_decay_steps = getenv("MAX_STEPS", math.ceil(1_200_000 * 1152 / GBS)) - opt_learning_rate_warmup_steps
|
||||
opt_base_learning_rate = getenv("LR", 8e-5 * GBS / 1152) # NOTE: cannot change for benchmark
|
||||
opt_end_learning_rate = getenv("END_LR", 8e-7)
|
||||
|
||||
model_params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]
|
||||
# vocab_size from the mixtral tokenizer
|
||||
if not SMALL: model_params |= {"vocab_size": 32000}
|
||||
real_vocab_size = model_params['vocab_size']
|
||||
if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: model_params['n_layers'] = llama_layers
|
||||
print(f"model parameters: {model_params}")
|
||||
|
||||
# pad vocab
|
||||
if (MP := getenv("MP", 1)) > 1: model_params['vocab_size'] = round_up(model_params['vocab_size'], 256 * MP)
|
||||
vocab_mask:Tensor = Tensor.arange(model_params['vocab_size']).reshape(1, 1, -1) >= real_vocab_size
|
||||
|
||||
model = FlatTransformer(**model_params, max_context=SEQLEN)
|
||||
|
||||
model = Transformer(**model_params, max_context=SEQLEN, jit=False, disable_kv_cache=True)
|
||||
params = get_parameters(model)
|
||||
# weights are all bfloat16 for now
|
||||
assert params and all(p.dtype == dtypes.bfloat16 for p in params)
|
||||
|
||||
if getenv("FAKEDATA"):
|
||||
for v in get_parameters(model):
|
||||
v = v.assign(Tensor.empty(v.shape))
|
||||
|
||||
is_dp = (DP := getenv("DP", 1)) > 1
|
||||
is_mp = (MP := getenv("MP", 1)) > 1
|
||||
is_sharding = is_dp or is_mp
|
||||
device_count = max(DP, MP)
|
||||
device = tuple(f"{Device.DEFAULT}:{i}" for i in range(device_count))
|
||||
if (DP := getenv("DP", 1)) > 1:
|
||||
device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
|
||||
for v in get_parameters(model):
|
||||
v.shard_(device, axis=None)
|
||||
|
||||
model.shard(device, is_mp)
|
||||
|
||||
if is_dp: vocab_mask.shard_(device, axis=None).realize()
|
||||
if is_mp: vocab_mask.shard_(device, axis=2).realize()
|
||||
|
||||
is_offload_optim = bool(getenv("OFFLOAD_OPTIM"))
|
||||
is_fake_offload = Device.DEFAULT == "NULL"
|
||||
optim_device = ("CPU" if not is_fake_offload else "NULL:99") if is_offload_optim else None
|
||||
optim = GradAccClipAdamW(params, lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2,
|
||||
eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay, grad_acc=grad_acc, device=optim_device)
|
||||
|
||||
# init grads
|
||||
for p in optim.params:
|
||||
p.grad = Tensor.zeros_like(p).contiguous()
|
||||
grads = [p.grad for p in optim.params]
|
||||
if (MP := getenv("MP", 1)) > 1:
|
||||
device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
|
||||
for k,v in get_state_dict(model).items():
|
||||
if 'scale' in k: v.shard_(device, axis=None) # from quantized
|
||||
elif '.attention.wq' in k: v.shard_(device, axis=0)
|
||||
elif '.attention.wk' in k: v.shard_(device, axis=0)
|
||||
elif '.attention.wv' in k: v.shard_(device, axis=0)
|
||||
elif '.attention.wo' in k: v.shard_(device, axis=1)
|
||||
elif '.feed_forward.w1.' in k: v.shard_(device, axis=0)
|
||||
elif '.feed_forward.w2.' in k: v.shard_(device, axis=1)
|
||||
elif '.feed_forward.w3.' in k: v.shard_(device, axis=0)
|
||||
elif 'tok_embeddings.weight' in k: v.shard_(device, axis=0)
|
||||
elif 'output.weight' in k: v.shard_(device, axis=0)
|
||||
else:
|
||||
# attention_norm, ffn_norm, norm
|
||||
v.shard_(device, axis=None)
|
||||
# prevents memory spike on device 0
|
||||
v.realize()
|
||||
|
||||
optim = AdamW(get_parameters(model), lr=0.0,
|
||||
b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay)
|
||||
scheduler = CosineAnnealingLRWithWarmup(optim, opt_base_learning_rate, opt_end_learning_rate, opt_learning_rate_warmup_steps, opt_learning_rate_decay_steps)
|
||||
|
||||
if resume_ckpt := getenv("RESUME_CKPT"):
|
||||
@@ -1385,131 +1365,97 @@ def train_llama3():
|
||||
load_state_dict(scheduler, safe_load(fn), realize=False)
|
||||
|
||||
@TinyJit
|
||||
def minibatch(tokens:Tensor):
|
||||
if is_dp: tokens = tokens.to(None).shard(device, 0)
|
||||
if is_mp: tokens = tokens.shard(device)
|
||||
if not is_sharding: tokens = tokens.to(None)
|
||||
logits:Tensor = model(tokens[:, :-1])
|
||||
loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
|
||||
|
||||
@Tensor.train()
|
||||
def train_step(model, tokens:Tensor):
|
||||
optim.zero_grad()
|
||||
if (DP := getenv("DP", 1)) > 1:
|
||||
device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
|
||||
tokens = tokens.shard(device, 0)
|
||||
if (MP := getenv("MP", 1)) > 1:
|
||||
device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
|
||||
tokens = tokens.shard(device)
|
||||
logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
|
||||
loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
|
||||
loss.backward()
|
||||
assert all(p.grad is g for p,g in zip(optim.params, grads))
|
||||
# L2 norm grad clip
|
||||
# https://github.com/NVIDIA/NeMo/blob/3368c3fc0b4a186ab33a1d68a504315100c0b2a6/nemo/collections/nlp/modules/common/megatron/clip_grads.py#L57
|
||||
# https://docs.pytorch.org/docs/stable/generated/torch.nn.utils.clip_grad_norm_.html
|
||||
if not getenv("DISABLE_GRAD_CLIP_NORM"):
|
||||
total_norm = Tensor(0.0, dtype=dtypes.float32, device=optim.params[0].device)
|
||||
for p in optim.params:
|
||||
total_norm += p.grad.float().square().sum()
|
||||
total_norm = total_norm.sqrt().contiguous()
|
||||
for p in optim.params:
|
||||
p.grad = p.grad * (opt_gradient_clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)
|
||||
|
||||
loss_cpu = loss.flatten().float().to("CPU")
|
||||
return loss_cpu.realize(*grads)
|
||||
|
||||
@TinyJit
|
||||
def optim_step():
|
||||
grad_norm = optim.fstep(grads)
|
||||
optim.step()
|
||||
scheduler.step()
|
||||
|
||||
for g in grads: g.assign(g.zeros_like())
|
||||
|
||||
lr_cpu = optim.lr.float().to("CPU")
|
||||
grad_norm_cpu = grad_norm.float().to("CPU")
|
||||
Tensor.realize(lr_cpu, grad_norm_cpu, *grads)
|
||||
|
||||
return lr_cpu, grad_norm_cpu
|
||||
lr = optim.lr
|
||||
loss.realize(lr)
|
||||
return loss, lr
|
||||
|
||||
@TinyJit
|
||||
@Tensor.train(False)
|
||||
def eval_step(tokens:Tensor):
|
||||
if is_dp: tokens = tokens.to(None).shard(device, 0)
|
||||
if is_mp: tokens = tokens.shard(device)
|
||||
if not is_sharding: tokens = tokens.to(None)
|
||||
logits:Tensor = model(tokens[:, :-1])
|
||||
loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
|
||||
return loss.flatten().float().to("CPU")
|
||||
def eval_step(model, tokens:Tensor):
|
||||
if (DP := getenv("DP", 1)) > 1:
|
||||
device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
|
||||
tokens = tokens.shard(device, 0)
|
||||
if (MP := getenv("MP", 1)) > 1:
|
||||
device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
|
||||
tokens = tokens.shard(device)
|
||||
logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
|
||||
loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
|
||||
return loss.flatten().float()
|
||||
|
||||
# ** data iters **
|
||||
def fake_data(bs, samples):
|
||||
import numpy as np
|
||||
for _ in range(samples // bs):
|
||||
fake_data_np = np.random.randint(0, model_params["vocab_size"], size=(bs, SEQLEN + 1), dtype=np.int32)
|
||||
yield Tensor(fake_data_np, device="NPY")
|
||||
yield Tensor.randint(bs, SEQLEN + 1, low=0, high=model_params["vocab_size"], dtype=dtypes.int32, device=Device.DEFAULT)
|
||||
|
||||
def get_train_iter():
|
||||
if getenv("FAKEDATA", 0):
|
||||
return fake_data(BS, SAMPLES)
|
||||
else:
|
||||
from examples.mlperf.dataloader import batch_load_llama3
|
||||
return batch_load_llama3(BS, SAMPLES, SEQLEN, BASEDIR, seed=DATA_SEED, val=bool(TRAIN_ON_VAL), small=bool(SMALL))
|
||||
|
||||
if getenv("FAKEDATA", 0):
|
||||
eval_dataset = None
|
||||
else:
|
||||
from examples.mlperf.dataloader import get_llama3_dataset
|
||||
eval_dataset = get_llama3_dataset(EVAL_SAMPLES, SEQLEN, BASEDIR, val=True, small=bool(SMALL))
|
||||
if SMALL:
|
||||
from examples.mlperf.dataloader import batch_load_llama3_small
|
||||
return batch_load_llama3_small(BS, SAMPLES, SEQLEN, BASEDIR, seed=SEED, val=bool(TRAIN_ON_VAL))
|
||||
else:
|
||||
from examples.mlperf.dataloader import batch_load_llama3
|
||||
return batch_load_llama3(BS, SAMPLES, SEQLEN, BASEDIR, seed=SEED, val=bool(TRAIN_ON_VAL))
|
||||
|
||||
def get_eval_iter():
|
||||
if eval_dataset is None:
|
||||
return fake_data(EVAL_BS, EVAL_SAMPLES)
|
||||
from examples.mlperf.dataloader import iterate_llama3_dataset
|
||||
return iterate_llama3_dataset(eval_dataset, EVAL_BS)
|
||||
if getenv("FAKEDATA", 0):
|
||||
return fake_data(EVAL_BS, 5760)
|
||||
else:
|
||||
if SMALL:
|
||||
from examples.mlperf.dataloader import batch_load_llama3_small
|
||||
return batch_load_llama3_small(EVAL_BS, 5760, SEQLEN, BASEDIR, val=True)
|
||||
else:
|
||||
from examples.mlperf.dataloader import batch_load_llama3
|
||||
return batch_load_llama3(EVAL_BS, 5760, SEQLEN, BASEDIR, val=True)
|
||||
|
||||
num_params = sum(p.numel() for p in params) - model_params["vocab_size"]*model_params["dim"]
|
||||
train_iter = get_train_iter()
|
||||
iter = get_train_iter()
|
||||
i, sequences_seen = resume_ckpt, 0
|
||||
step_times = []
|
||||
while i < MAX_STEPS:
|
||||
for tokens in tqdm(iter, total=SAMPLES//GBS):
|
||||
GlobalCounters.reset()
|
||||
actual_gbs = GBS if i >= 2 else BS
|
||||
if getenv("TRAIN", 1):
|
||||
profile_marker(f"train @ {i}")
|
||||
st = time.perf_counter()
|
||||
|
||||
stopped = False
|
||||
losses, data_time, dev_time = [], 0, 0
|
||||
for _ in range(grad_acc if i >= 2 else 1):
|
||||
ist = time.perf_counter()
|
||||
try: tokens = next(train_iter)
|
||||
except StopIteration:
|
||||
stopped = True
|
||||
break
|
||||
mst = time.perf_counter()
|
||||
data_time += mst - ist
|
||||
losses.append(minibatch(tokens).item())
|
||||
dev_time += time.perf_counter() - mst
|
||||
if stopped: break
|
||||
|
||||
gt = time.perf_counter()
|
||||
ret = optim_step()
|
||||
lr, grad_norm = ret[0].item(), ret[1].item()
|
||||
et = time.perf_counter()
|
||||
|
||||
loss = sum(losses) / len(losses)
|
||||
optim_time = et - gt
|
||||
dev_time += optim_time
|
||||
step_time = et - st
|
||||
gbs_time = gt - st
|
||||
if BENCHMARK: step_times.append(step_time)
|
||||
t = time.perf_counter()
|
||||
loss, lr = train_step(model, tokens)
|
||||
loss = loss.float().item()
|
||||
lr = lr.item()
|
||||
|
||||
i += 1
|
||||
sequences_seen += actual_gbs
|
||||
sequences_seen += tokens.shape[0]
|
||||
|
||||
mem_gb = GlobalCounters.mem_used / 1e9
|
||||
gflops = GlobalCounters.global_ops / 1e9 / dev_time
|
||||
mfu = ((6 * num_params * SEQLEN * GBS) / (dev_time * device_count * 2.3e15)) * 100
|
||||
sec = time.perf_counter()-t
|
||||
tqdm.write(
|
||||
f"{i:5} {step_time:.3f} s step, {gbs_time:.3f} s gbs, {optim_time:.3f} s optim, {data_time:.3f} s data, {loss:.4f} loss, " \
|
||||
f"{lr:.12f} LR, {grad_norm:.6f} grad_norm, {mem_gb:.2f} GB used, {gflops:9.2f} GFLOPS, {mfu:5.2f}% MFU")
|
||||
if DEBUG >= 1: tqdm.write(" mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))
|
||||
f"{i:5} {sec:.2f} s run, {loss:.4f} loss, {lr:.12f} LR, {GlobalCounters.mem_used / 1e9:.2f} GB used, "
|
||||
f"{GlobalCounters.global_ops * 1e-9 / sec:9.2f} GFLOPS")
|
||||
|
||||
if WANDB:
|
||||
wandb.log({
|
||||
"train/loss": loss,
|
||||
"train/lr": lr,
|
||||
"train/grad_norm": grad_norm,
|
||||
"train/step_time": step_time,
|
||||
"train/gbs_time": gbs_time,
|
||||
"train/optim_time": optim_time,
|
||||
"train/dev_time": dev_time,
|
||||
"train/data_time": data_time,
|
||||
"train/mem": mem_gb,
|
||||
"train/GFLOPS": gflops,
|
||||
"train/MFU": mfu,
|
||||
"train/sequences_seen": sequences_seen
|
||||
})
|
||||
if (fname:=getenv("LOSS_FILE", "")):
|
||||
with open(fname, "a") as f:
|
||||
f.write(f"{i} {loss:.4f} {lr.item():.12f} {GlobalCounters.mem_used / 1e9:.2f}\n")
|
||||
|
||||
if (ckpt_freq := getenv("CKPT")) and (i % ckpt_freq == 0 and (i != 1 or ckpt_freq == 1)):
|
||||
tqdm.write("saving checkpoint")
|
||||
@@ -1521,36 +1467,20 @@ def train_llama3():
|
||||
fn = f"{ckpt_dir}/llama3_{i}_optim.safe"
|
||||
safe_save(get_state_dict(scheduler), fn)
|
||||
|
||||
if i == BENCHMARK:
|
||||
median_step_time = sorted(step_times)[(BENCHMARK + 1) // 2]
|
||||
estimated_total_minutes = int(median_step_time * (SAMPLES // GBS) / 60)
|
||||
print(f"Estimated training time: {estimated_total_minutes // 60}h{estimated_total_minutes % 60}m")
|
||||
print(f"epoch global_ops: {GlobalCounters.global_ops:_}, "
|
||||
f"epoch global_mem: {GlobalCounters.global_mem:_}")
|
||||
|
||||
if (sequences_seen // EVAL_FREQ != (sequences_seen - actual_gbs) // EVAL_FREQ and (i != 1 or EVAL_FREQ == 1)) or (BENCHMARK and i == BENCHMARK):
|
||||
if EVAL_BS == 0: return
|
||||
if sequences_seen % EVAL_FREQ == 0 and (i != 1 or EVAL_FREQ == 1):
|
||||
tqdm.write(f"evaluating after {sequences_seen} sequences")
|
||||
profile_marker(f"eval @ {i}")
|
||||
|
||||
# run eval
|
||||
eval_losses = []
|
||||
eval_iter = get_eval_iter()
|
||||
tqdm.write(f"evaluating {EVAL_SAMPLES//EVAL_BS} batches of {EVAL_BS} sequences")
|
||||
tqdm.write(f"evaluating {5760//EVAL_BS} batches of {EVAL_BS} sequences")
|
||||
|
||||
for j,tokens in tqdm(enumerate(eval_iter), total=EVAL_SAMPLES//EVAL_BS):
|
||||
eval_losses += eval_step(tokens).tolist()
|
||||
|
||||
if BENCHMARK and (j+1) == min(BENCHMARK, EVAL_SAMPLES//EVAL_BS):
|
||||
return
|
||||
|
||||
log_perplexity = sum(eval_losses) / len(eval_losses)
|
||||
for tokens in tqdm(eval_iter, total=5760//EVAL_BS):
|
||||
eval_losses += eval_step(model, tokens).tolist()
|
||||
log_perplexity = Tensor(eval_losses).mean().float().item()
|
||||
|
||||
tqdm.write(f"eval log perplexity: {log_perplexity:.4f}")
|
||||
|
||||
if WANDB:
|
||||
wandb.log({"eval/log_perplexity": log_perplexity, "eval/sequences_seen": sequences_seen})
|
||||
|
||||
if log_perplexity < EVAL_TARGET:
|
||||
tqdm.write(f"target achieved after {sequences_seen} sequences")
|
||||
if getenv("CKPT"):
|
||||
@@ -1629,7 +1559,7 @@ def train_stable_diffusion():
|
||||
loss, out_lr = loss.detach().to("CPU"), optimizer.lr.to("CPU")
|
||||
Tensor.realize(loss, out_lr)
|
||||
return loss, out_lr
|
||||
|
||||
|
||||
# checkpointing takes ~9 minutes without this, and ~1 minute with this
|
||||
@TinyJit
|
||||
def ckpt_to_cpu():
|
||||
@@ -1668,7 +1598,7 @@ def train_stable_diffusion():
|
||||
if i == 3:
|
||||
for _ in range(3): ckpt_to_cpu() # do this at the beginning of run to prevent OOM surprises when checkpointing
|
||||
print("BEAM COMPLETE", flush=True) # allows wrapper script to detect BEAM search completion and retry if it failed
|
||||
|
||||
|
||||
total_train_time = time.perf_counter() - train_start_time
|
||||
if WANDB:
|
||||
wandb.log({"train/loss": loss_item, "train/lr": lr_item, "train/loop_time_prev": loop_time, "train/dl_time": dl_time, "train/step": i,
|
||||
|
||||
@@ -1,214 +0,0 @@
|
||||
import math, os
|
||||
if __name__ == "__main__":
|
||||
os.environ["DEFAULT_FLOAT"] = "bfloat16"
|
||||
os.environ["OPTIM_DTYPE"] = "bfloat16"
|
||||
if "DEV" not in os.environ: os.environ["DEV"] = "NULL"
|
||||
# CDNA
|
||||
os.environ["EMULATE"] = "AMD_CDNA4"
|
||||
os.environ["DEVICE_IN_FUNCTION_BUG"] = "1"
|
||||
os.environ["ALL2ALL"] = "1"
|
||||
os.environ["USE_ATOMICS"] = "1"
|
||||
if "HK_FLASH_ATTENTION" not in os.environ:
|
||||
os.environ["HK_FLASH_ATTENTION"] = "1"
|
||||
if "ASM_GEMM" not in os.environ:
|
||||
os.environ["ASM_GEMM"] = "1"
|
||||
from tinygrad import Tensor, nn, function, getenv, dtypes, TinyJit
|
||||
from tinygrad.helpers import Timing, colored, GlobalCounters, profile_marker
|
||||
from tinygrad.uop.ops import Ops, UOp
|
||||
from extra.models.llama import apply_rotary_emb, precompute_freqs_cis
|
||||
|
||||
FP8 = getenv("FP8", 0)
|
||||
WQKV = getenv("WQKV", 0)
|
||||
|
||||
FP8_DTYPE = dtypes.fp8e4m3
|
||||
FP8_MAX = 448.0
|
||||
|
||||
def quantize_fp8(x:Tensor):
|
||||
scale = FP8_MAX / (x.abs().max().detach() + 1e-8)
|
||||
x_scaled = x * scale
|
||||
x_clamped = x_scaled + (x_scaled.detach().clamp(-FP8_MAX, FP8_MAX) - x_scaled.detach()) # STE
|
||||
return x_clamped.cast(FP8_DTYPE), scale.float().reciprocal()
|
||||
|
||||
def matmul(x:Tensor, w:Tensor) -> Tensor:
|
||||
if not FP8: return x @ w.T
|
||||
# weights are already FP8, just quantize activations
|
||||
x_fp8, x_scale = quantize_fp8(x)
|
||||
return x_fp8.dot(w.T, dtype=dtypes.float) * x_scale
|
||||
|
||||
def rmsnorm(x_in:Tensor, eps:float):
|
||||
x = x_in.float()
|
||||
x = x * (x.square().mean(-1, keepdim=True) + eps).rsqrt()
|
||||
return x.cast(x_in.dtype)
|
||||
|
||||
class FlatTransformer:
|
||||
def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_layers:int, norm_eps:float, vocab_size:int, n_kv_heads:int|None=None,
|
||||
rope_theta:int=10000, max_context:int=1024):
|
||||
self.vocab_size = vocab_size
|
||||
self.n_layers = n_layers
|
||||
self.n_heads = n_heads
|
||||
self.n_kv_heads = n_kv_heads if n_kv_heads is not None else n_heads # n_kv_heads != n_heads implies MQA [arxiv/2307.09288, A.2.1]
|
||||
self.head_dim = dim // n_heads
|
||||
self.n_rep = self.n_heads // self.n_kv_heads
|
||||
|
||||
scaled_std = 0.02 / math.sqrt(2 * n_layers)
|
||||
|
||||
# Attention
|
||||
if WQKV:
|
||||
self.wqkv = self.lin_per_layer(dim, self.n_heads * self.head_dim + self.n_kv_heads * self.head_dim * 2)
|
||||
else:
|
||||
self.wq = self.lin_per_layer(dim, self.n_heads * self.head_dim)
|
||||
self.wk = self.lin_per_layer(dim, self.n_kv_heads * self.head_dim)
|
||||
self.wv = self.lin_per_layer(dim, self.n_kv_heads * self.head_dim)
|
||||
self.wo = self.lin_per_layer(self.n_heads * self.head_dim, dim, std=scaled_std)
|
||||
|
||||
# FeedForward
|
||||
self.w1 = self.lin_per_layer(dim, hidden_dim)
|
||||
self.w2 = self.lin_per_layer(hidden_dim, dim, std=scaled_std)
|
||||
self.w3 = self.lin_per_layer(dim, hidden_dim)
|
||||
|
||||
self.norm_eps = norm_eps
|
||||
self.attention_norm = Tensor.ones(n_layers, dim).contiguous()
|
||||
self.ffn_norm = Tensor.ones(n_layers, dim).contiguous()
|
||||
|
||||
# output
|
||||
self.norm = nn.RMSNorm(dim, norm_eps)
|
||||
self.tok_embeddings = nn.Embedding(vocab_size, dim)
|
||||
self.tok_embeddings.weight = Tensor.normal(vocab_size, dim, mean=0.0, std=0.02)
|
||||
self.output = Tensor.normal(1, vocab_size, dim, mean=0.0, std=0.02)
|
||||
self.freqs_cis = precompute_freqs_cis(dim // n_heads, max_context * 2, rope_theta).contiguous().requires_grad_(False)
|
||||
|
||||
def lin_per_layer(self, in_features:int, out_features:int, std:float=0.02):
|
||||
dt = FP8_DTYPE if FP8 else None
|
||||
if getenv("ZEROS"): return Tensor.zeros(self.n_layers, out_features, in_features, dtype=dt)
|
||||
return Tensor.normal(self.n_layers, out_features, in_features, mean=0.0, std=std, dtype=dt)
|
||||
|
||||
def attention(self, x:Tensor, freqs_cis:Tensor, attention_norm:Tensor, wo:Tensor, wqkv:Tensor|None=None,
|
||||
wq:Tensor|None=None, wk:Tensor|None=None, wv:Tensor|None=None):
|
||||
x = rmsnorm(x, self.norm_eps) * attention_norm
|
||||
bsz, seqlen, _ = x.shape
|
||||
|
||||
if wqkv is not None:
|
||||
xqkv = matmul(x, wqkv)
|
||||
xqkv = xqkv.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep + 2, self.head_dim)
|
||||
xq = xqkv[:, :, :, :self.n_rep].reshape(bsz, seqlen, self.n_heads, self.head_dim)
|
||||
xk = xqkv[:, :, :, self.n_rep].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
|
||||
xv = xqkv[:, :, :, self.n_rep+1].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
|
||||
else:
|
||||
assert wq is not None and wk is not None and wv is not None
|
||||
xq = matmul(x, wq).reshape(bsz, seqlen, self.n_heads, self.head_dim)
|
||||
xk = matmul(x, wk).reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
|
||||
xv = matmul(x, wv).reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
|
||||
|
||||
xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
|
||||
xq, xk, xv = xq.transpose(1, 2), xk.transpose(1, 2), xv.transpose(1, 2)
|
||||
attn = xq.scaled_dot_product_attention(xk, xv, is_causal=True, enable_gqa=True).transpose(1, 2)
|
||||
attn = attn.reshape(bsz, seqlen, -1)
|
||||
return matmul(attn, wo)
|
||||
|
||||
def feed_forward(self, x:Tensor, ffn_norm:Tensor, w1:Tensor, w2:Tensor, w3:Tensor):
|
||||
x = rmsnorm(x, self.norm_eps) * ffn_norm
|
||||
x_w1 = matmul(x, w1).silu()
|
||||
x_w3 = matmul(x.contiguous_backward(), w3)
|
||||
return matmul(x_w1 * x_w3, w2)
|
||||
|
||||
@function(precompile=True, precompile_backward=True)
|
||||
def run_layer(self, x:Tensor, freqs_cis:Tensor,
|
||||
attention_norm:Tensor, wo:Tensor,
|
||||
ffn_norm:Tensor, w1:Tensor, w2:Tensor, w3:Tensor,
|
||||
wqkv:Tensor|None=None, wq:Tensor|None=None, wk:Tensor|None=None, wv:Tensor|None=None):
|
||||
h = x + self.attention(x, freqs_cis, attention_norm, wo, wqkv=wqkv, wq=wq, wk=wk, wv=wv)
|
||||
return h + self.feed_forward(h, ffn_norm, w1, w2, w3)
|
||||
|
||||
def shard(self, device:tuple[str, ...], mp:bool=False):
|
||||
from tinygrad.nn.state import get_parameters
|
||||
if not mp:
|
||||
for v in get_parameters(self): v.shard_(device, axis=None)
|
||||
else:
|
||||
# flat per-layer weights: axis 0 is n_layers, so shard axes are +1 vs per-layer Transformer
|
||||
if WQKV:
|
||||
self.wqkv.shard_(device, axis=1).realize() # (n_layers, out, dim) shard out
|
||||
else:
|
||||
self.wq.shard_(device, axis=1).realize() # (n_layers, n_heads*head_dim, dim) shard out
|
||||
self.wk.shard_(device, axis=1).realize() # (n_layers, n_kv_heads*head_dim, dim) shard out
|
||||
self.wv.shard_(device, axis=1).realize() # (n_layers, n_kv_heads*head_dim, dim) shard out
|
||||
self.wo.shard_(device, axis=2).realize() # (n_layers, dim, in) shard in
|
||||
self.w1.shard_(device, axis=1).realize() # (n_layers, hidden, dim) shard out
|
||||
self.w2.shard_(device, axis=2).realize() # (n_layers, dim, hidden) shard in
|
||||
self.w3.shard_(device, axis=1).realize() # (n_layers, hidden, dim) shard out
|
||||
self.attention_norm.shard_(device, axis=None).realize()
|
||||
self.ffn_norm.shard_(device, axis=None).realize()
|
||||
self.norm.weight.shard_(device, axis=None).realize()
|
||||
self.tok_embeddings.weight.shard_(device, axis=0).realize()
|
||||
self.output.shard_(device, axis=1).realize()
|
||||
self.freqs_cis.shard_(device, axis=None).realize()
|
||||
|
||||
def __call__(self, tokens:Tensor):
|
||||
h = self.tok_embeddings(tokens)
|
||||
freqs_cis = self.freqs_cis.cast(h.dtype)[:, :tokens.shape[1], :, :, :]
|
||||
for i in range(self.n_layers):
|
||||
attn_kwargs = {"wqkv": self.wqkv[i]} if WQKV else {"wq": self.wq[i], "wk": self.wk[i], "wv": self.wv[i]}
|
||||
h = self.run_layer(h, freqs_cis,
|
||||
self.attention_norm[i], self.wo[i],
|
||||
self.ffn_norm[i], self.w1[i], self.w2[i], self.w3[i], **attn_kwargs)
|
||||
logits = self.norm(h) @ self.output[0].T
|
||||
return logits
|
||||
|
||||
# TODO: this shouldn't be needed, but it prevents a copy of the grads. CAT can help
|
||||
def apply_grad(old_grad:UOp, new_grad:UOp) -> list[UOp]:
|
||||
if new_grad.op == Ops.ADD:
|
||||
return apply_grad(old_grad, new_grad.src[0])+apply_grad(old_grad, new_grad.src[1])
|
||||
elif new_grad.op == Ops.PAD:
|
||||
grad_shrink = tuple([(p[0], s+p[0]) for s,p in zip(new_grad.src[0].shape, new_grad.marg)])
|
||||
return apply_grad(old_grad.shrink(grad_shrink), new_grad.src[0])
|
||||
else:
|
||||
return [old_grad.store(old_grad + new_grad)]
|
||||
|
||||
if __name__ == "__main__":
|
||||
config = {}
|
||||
BS = config["BS"] = getenv("BS", 16)
|
||||
SEQLEN = config["SEQLEN"] = getenv("SEQLEN", 8192)
|
||||
|
||||
from examples.llama3 import MODEL_PARAMS
|
||||
model_params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]
|
||||
if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: model_params['n_layers'] = llama_layers
|
||||
model = FlatTransformer(**model_params, max_context=SEQLEN)
|
||||
state = nn.state.get_state_dict(model)
|
||||
print("tensor count:", len(state))
|
||||
|
||||
# shard the model
|
||||
from tinygrad import Device
|
||||
if (DP := getenv("DP", 1)) > 1:
|
||||
model.shard(tuple(f"{Device.DEFAULT}:{i}" for i in range(DP)))
|
||||
if (MP := getenv("MP", 1)) > 1:
|
||||
model.shard(tuple(f"{Device.DEFAULT}:{i}" for i in range(MP)), mp=True)
|
||||
|
||||
# preallocate all the grad buffers and zero them out
|
||||
grads = {x:Tensor.zeros_like(x).contiguous() for x in state.values() if x.requires_grad is None}
|
||||
|
||||
# print model size
|
||||
sz = 0
|
||||
for k,v in state.items():
|
||||
print(f"{colored(k, 'green' if v in grads else 'white'):30s} {str(v.shape):30s} {str(v.dtype):20s} {v.device} {v.nbytes()/1e9:.2f} GB")
|
||||
sz += v.nbytes()
|
||||
print(f"total sz: {sz/1e9:.2f} GB")
|
||||
|
||||
with Timing("fake data: "): tokens = Tensor.randint(BS, SEQLEN+1, low=0, high=model.vocab_size, dtype=dtypes.int)
|
||||
with Timing("realize weights/grads/data: "): Tensor.realize(*state.values(), *grads.values(), tokens)
|
||||
print("mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))
|
||||
if DP > 1: tokens = tokens.shard(tuple(f"{Device.DEFAULT}:{i}" for i in range(DP)), axis=0)
|
||||
if MP > 1: tokens = tokens.shard(tuple(f"{Device.DEFAULT}:{i}" for i in range(MP)))
|
||||
|
||||
@TinyJit
|
||||
def jit_step(tokens:Tensor):
|
||||
with Timing("python forward: "): loss = model(tokens[:, :-1]).sparse_categorical_crossentropy(tokens[:, 1:])
|
||||
with Timing("python backward: "):
|
||||
for t,g in zip(grads, loss.gradient(*grads)):
|
||||
grads[t] = Tensor(grads[t].uop.after(UOp.group(*apply_grad(grads[t].uop, g.uop))), device=t.device)
|
||||
with Timing("run step: "): loss.realize(*grads.values())
|
||||
|
||||
for i in range(6):
|
||||
GlobalCounters.reset()
|
||||
profile_marker(f"step {i}")
|
||||
with Timing(colored(f"*** step {i}: ", "red")):
|
||||
jit_step(tokens)
|
||||
print("mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))
|
||||
@@ -1,80 +0,0 @@
|
||||
from tinygrad import Tensor, nn
|
||||
from tinygrad.helpers import getenv
|
||||
from extra.models.llama import apply_rotary_emb, precompute_freqs_cis
|
||||
|
||||
class Attention:
|
||||
def __init__(self, dim:int, n_heads:int, n_kv_heads:int|None=None, linear=nn.Linear):
|
||||
self.n_heads = n_heads
|
||||
self.n_kv_heads = n_kv_heads if n_kv_heads is not None else n_heads # n_kv_heads != n_heads implies MQA [arxiv/2307.09288, A.2.1]
|
||||
self.head_dim = dim // n_heads
|
||||
self.n_rep = self.n_heads // self.n_kv_heads
|
||||
|
||||
if getenv("WQKV"):
|
||||
self.wqkv = linear(dim, self.n_heads * self.head_dim + self.n_kv_heads * self.head_dim * 2, bias=False)
|
||||
else:
|
||||
self.wq = linear(dim, self.n_heads * self.head_dim, bias=False)
|
||||
self.wk = linear(dim, self.n_kv_heads * self.head_dim, bias=False)
|
||||
self.wv = linear(dim, self.n_kv_heads * self.head_dim, bias=False)
|
||||
|
||||
self.wo = linear(self.n_heads * self.head_dim, dim, bias=False)
|
||||
|
||||
def __call__(self, x:Tensor, freqs_cis:Tensor) -> Tensor:
|
||||
if getenv("WQKV"):
|
||||
xqkv = self.wqkv(x)
|
||||
xqkv = xqkv.reshape(xqkv.shape[0], xqkv.shape[1], self.n_kv_heads, self.n_rep + 2, self.head_dim)
|
||||
xq = xqkv[:, :, :, :self.n_rep].reshape(xqkv.shape[0], xqkv.shape[1], -1)
|
||||
xk = xqkv[:, :, :, self.n_rep:self.n_rep+1].reshape(xqkv.shape[0], xqkv.shape[1], -1)
|
||||
xv = xqkv[:, :, :, self.n_rep+1:self.n_rep+2].reshape(xqkv.shape[0], xqkv.shape[1], -1)
|
||||
else:
|
||||
xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
|
||||
|
||||
xq = xq.reshape(xq.shape[0], xq.shape[1], self.n_heads, self.head_dim)
|
||||
xk = xk.reshape(xk.shape[0], xk.shape[1], self.n_kv_heads, self.head_dim)
|
||||
xv = xv.reshape(xv.shape[0], xv.shape[1], self.n_kv_heads, self.head_dim)
|
||||
|
||||
xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
|
||||
bsz, seqlen, _, _ = xq.shape
|
||||
|
||||
xq, xk, xv = xq.transpose(1, 2), xk.transpose(1, 2), xv.transpose(1, 2)
|
||||
attn = xq.scaled_dot_product_attention(xk, xv, is_causal=True, enable_gqa=True).transpose(1, 2)
|
||||
|
||||
attn = attn.reshape(bsz, seqlen, -1)
|
||||
return self.wo(attn)
|
||||
|
||||
class FeedForward:
|
||||
def __init__(self, dim:int, hidden_dim:int, linear=nn.Linear):
|
||||
self.w1 = linear(dim, hidden_dim, bias=False)
|
||||
self.w2 = linear(hidden_dim, dim, bias=False)
|
||||
self.w3 = linear(dim, hidden_dim, bias=False) # the gate in Gated Linear Unit
|
||||
|
||||
def __call__(self, x:Tensor) -> Tensor:
|
||||
w1 = self.w1(x).silu()
|
||||
w3 = self.w3(x)
|
||||
return self.w2(w1 * w3)
|
||||
|
||||
class TransformerBlock:
|
||||
def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_kv_heads:int|None, norm_eps:float, linear=nn.Linear):
|
||||
self.attention = Attention(dim, n_heads, n_kv_heads, linear)
|
||||
self.feed_forward = FeedForward(dim, hidden_dim, linear)
|
||||
self.attention_norm = nn.RMSNorm(dim, norm_eps)
|
||||
self.ffn_norm = nn.RMSNorm(dim, norm_eps)
|
||||
|
||||
def __call__(self, x:Tensor, freqs_cis:Tensor):
|
||||
h = x + self.attention(self.attention_norm(x), freqs_cis)
|
||||
return h + self.feed_forward(self.ffn_norm(h))
|
||||
|
||||
class Transformer:
|
||||
def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_layers:int, norm_eps:float, vocab_size:int, n_kv_heads:int|None=None,
|
||||
rope_theta:int=10000, max_context:int=1024, linear=nn.Linear, embedding=nn.Embedding):
|
||||
self.layers = [TransformerBlock(dim, hidden_dim, n_heads, n_kv_heads, norm_eps, linear) for _ in range(n_layers)]
|
||||
self.norm = nn.RMSNorm(dim, norm_eps)
|
||||
self.tok_embeddings = embedding(vocab_size, dim)
|
||||
self.output = nn.Linear(dim, vocab_size, bias=False) if embedding == nn.Embedding else linear(dim, vocab_size, bias=False)
|
||||
self.freqs_cis = precompute_freqs_cis(dim // n_heads, max_context * 2, rope_theta).contiguous().requires_grad_(False)
|
||||
|
||||
def __call__(self, tokens:Tensor):
|
||||
h = self.tok_embeddings(tokens)
|
||||
freqs_cis = self.freqs_cis.cast(h.dtype)[:, :tokens.shape[1], :, :, :]
|
||||
for layer in self.layers: h = layer(h, freqs_cis)
|
||||
logits = self.output(self.norm(h))
|
||||
return logits
|
||||
@@ -1,140 +0,0 @@
|
||||
import os
|
||||
os.environ["WQKV"] = "1"
|
||||
import unittest
|
||||
import numpy as np
|
||||
from tinygrad import Tensor, nn, dtypes
|
||||
from tinygrad.nn.state import get_parameters
|
||||
from tinygrad.device import is_dtype_supported, Device
|
||||
from examples.mlperf.models.llama import Transformer
|
||||
from examples.mlperf.models.flat_llama import FlatTransformer
|
||||
|
||||
def copy_weights(flat:FlatTransformer, ref:Transformer):
|
||||
n_layers = flat.n_layers
|
||||
Tensor.realize(*nn.state.get_state_dict(ref).values())
|
||||
flat.wqkv.assign(Tensor(np.stack([ref.layers[i].attention.wqkv.weight.numpy() for i in range(n_layers)])))
|
||||
flat.wo.assign(Tensor(np.stack([ref.layers[i].attention.wo.weight.numpy() for i in range(n_layers)])))
|
||||
flat.w1.assign(Tensor(np.stack([ref.layers[i].feed_forward.w1.weight.numpy() for i in range(n_layers)])))
|
||||
flat.w2.assign(Tensor(np.stack([ref.layers[i].feed_forward.w2.weight.numpy() for i in range(n_layers)])))
|
||||
flat.w3.assign(Tensor(np.stack([ref.layers[i].feed_forward.w3.weight.numpy() for i in range(n_layers)])))
|
||||
flat.attention_norm.assign(Tensor(np.stack([ref.layers[i].attention_norm.weight.numpy() for i in range(n_layers)])))
|
||||
flat.ffn_norm.assign(Tensor(np.stack([ref.layers[i].ffn_norm.weight.numpy() for i in range(n_layers)])))
|
||||
flat.norm.weight.assign(Tensor(ref.norm.weight.numpy()))
|
||||
flat.tok_embeddings.weight.assign(Tensor(ref.tok_embeddings.weight.numpy()))
|
||||
flat.output.weight.assign(Tensor(ref.output.weight.numpy()))
|
||||
|
||||
class TestFlatLlama(unittest.TestCase):
|
||||
def test_forward_match(self):
|
||||
Tensor.manual_seed(42)
|
||||
params = dict(dim=128, hidden_dim=256, n_heads=4, n_kv_heads=2, n_layers=2, norm_eps=1e-5, vocab_size=1024, rope_theta=10000, max_context=64)
|
||||
ref = Transformer(**params)
|
||||
flat = FlatTransformer(**params)
|
||||
copy_weights(flat, ref)
|
||||
Tensor.realize(*nn.state.get_state_dict(flat).values())
|
||||
|
||||
tokens = Tensor([[1, 50, 100, 999, 2]])
|
||||
ref_logits = ref(tokens).realize()
|
||||
flat_logits = flat(tokens).realize()
|
||||
self.assertEqual(ref_logits.shape, flat_logits.shape)
|
||||
diff = (ref_logits - flat_logits).abs().max().item()
|
||||
self.assertLess(diff, 1e-5, f"forward mismatch: max abs diff {diff}")
|
||||
|
||||
def test_backward_match(self):
|
||||
Tensor.manual_seed(42)
|
||||
params = dict(dim=128, hidden_dim=256, n_heads=4, n_kv_heads=2, n_layers=2, norm_eps=1e-5, vocab_size=1024, rope_theta=10000, max_context=64)
|
||||
ref = Transformer(**params)
|
||||
flat = FlatTransformer(**params)
|
||||
copy_weights(flat, ref)
|
||||
|
||||
for p in get_parameters(ref): p.requires_grad_(True)
|
||||
for p in get_parameters(flat): p.requires_grad_(True)
|
||||
Tensor.realize(*nn.state.get_state_dict(flat).values())
|
||||
|
||||
tokens = Tensor([[1, 50, 100, 999, 2, 10]])
|
||||
|
||||
ref_loss = ref(tokens[:, :-1]).sparse_categorical_crossentropy(tokens[:, 1:])
|
||||
ref_loss.backward()
|
||||
ref_grads = {k: v.grad.numpy() for k, v in nn.state.get_state_dict(ref).items() if v.grad is not None}
|
||||
|
||||
flat_loss = flat(tokens[:, :-1]).sparse_categorical_crossentropy(tokens[:, 1:])
|
||||
flat_loss.backward()
|
||||
flat_grads = {k: v.grad.numpy() for k, v in nn.state.get_state_dict(flat).items() if v.grad is not None}
|
||||
|
||||
# check loss matches
|
||||
self.assertAlmostEqual(ref_loss.item(), flat_loss.item(), places=4)
|
||||
|
||||
# check output weight grad matches
|
||||
diff = abs(ref_grads["output.weight"] - flat_grads["output.weight"]).max()
|
||||
self.assertLess(diff, 1e-4, f"output.weight grad mismatch: max abs diff {diff}")
|
||||
|
||||
# check per-layer weight grads match
|
||||
for i in range(params["n_layers"]):
|
||||
for flat_key, ref_key in [
|
||||
("wqkv", f"layers.{i}.attention.wqkv.weight"),
|
||||
("wo", f"layers.{i}.attention.wo.weight"),
|
||||
("w1", f"layers.{i}.feed_forward.w1.weight"),
|
||||
("w2", f"layers.{i}.feed_forward.w2.weight"),
|
||||
("w3", f"layers.{i}.feed_forward.w3.weight"),
|
||||
]:
|
||||
diff = abs(ref_grads[ref_key] - flat_grads[flat_key][i]).max()
|
||||
self.assertLess(diff, 1e-4, f"layer {i} {flat_key} grad mismatch: max abs diff {diff}")
|
||||
|
||||
@unittest.skipUnless(Device.DEFAULT == "CPU", "multi-device CPU test")
|
||||
def test_forward_match_mp(self):
|
||||
Tensor.manual_seed(42)
|
||||
params = dict(dim=128, hidden_dim=256, n_heads=4, n_kv_heads=2, n_layers=2, norm_eps=1e-5, vocab_size=1024, rope_theta=10000, max_context=64)
|
||||
from tinygrad import Device
|
||||
devices = (f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1")
|
||||
ref = Transformer(**params)
|
||||
flat = FlatTransformer(**params)
|
||||
copy_weights(flat, ref)
|
||||
Tensor.realize(*nn.state.get_state_dict(flat).values())
|
||||
flat.shard(devices, mp=True)
|
||||
|
||||
tokens = Tensor([[1, 50, 100, 999, 2]], device=devices[0])
|
||||
ref_logits = ref(tokens.to(devices[0])).numpy()
|
||||
flat_logits = flat(tokens.shard(devices)).numpy()
|
||||
self.assertEqual(ref_logits.shape, flat_logits.shape)
|
||||
np.testing.assert_allclose(flat_logits, ref_logits, atol=1e-4, rtol=1e-4)
|
||||
|
||||
@unittest.skipUnless(Device.DEFAULT == "CPU", "multi-device CPU test")
|
||||
def test_forward_match_dp(self):
|
||||
Tensor.manual_seed(42)
|
||||
params = dict(dim=128, hidden_dim=256, n_heads=4, n_kv_heads=2, n_layers=2, norm_eps=1e-5, vocab_size=1024, rope_theta=10000, max_context=64)
|
||||
from tinygrad import Device
|
||||
devices = (f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1")
|
||||
ref = Transformer(**params)
|
||||
flat = FlatTransformer(**params)
|
||||
copy_weights(flat, ref)
|
||||
Tensor.realize(*nn.state.get_state_dict(flat).values())
|
||||
flat.shard(devices)
|
||||
|
||||
tokens = Tensor([[1, 50, 100, 999, 2], [2, 100, 50, 1, 999]], device=devices[0])
|
||||
ref_logits = ref(tokens.to(devices[0])).numpy()
|
||||
flat_logits = flat(tokens.shard(devices, axis=0)).numpy()
|
||||
self.assertEqual(ref_logits.shape, flat_logits.shape)
|
||||
np.testing.assert_allclose(flat_logits, ref_logits, atol=1e-4, rtol=1e-4)
|
||||
|
||||
@unittest.skipUnless(is_dtype_supported(dtypes.fp8e4m3), "fp8 not supported on this device")
|
||||
def test_forward_fp8(self):
|
||||
import examples.mlperf.models.flat_llama as flat_llama_mod
|
||||
old_fp8 = flat_llama_mod.FP8
|
||||
try:
|
||||
flat_llama_mod.FP8 = 1
|
||||
Tensor.manual_seed(42)
|
||||
params = dict(dim=128, hidden_dim=256, n_heads=4, n_kv_heads=2, n_layers=2, norm_eps=1e-5, vocab_size=1024, rope_theta=10000, max_context=64)
|
||||
ref = Transformer(**params)
|
||||
flat = FlatTransformer(**params)
|
||||
copy_weights(flat, ref)
|
||||
Tensor.realize(*nn.state.get_state_dict(flat).values())
|
||||
|
||||
tokens = Tensor([[1, 50, 100, 999, 2]])
|
||||
ref_logits = ref(tokens).numpy()
|
||||
flat_logits = flat(tokens).numpy()
|
||||
self.assertEqual(ref_logits.shape, flat_logits.shape)
|
||||
# FP8 has lower precision, allow larger tolerance
|
||||
np.testing.assert_allclose(flat_logits, ref_logits, atol=1.0, rtol=0.1)
|
||||
finally:
|
||||
flat_llama_mod.FP8 = old_fp8
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,80 +0,0 @@
|
||||
from tinygrad.tensor import Tensor
|
||||
from tinygrad.dtype import dtypes
|
||||
from tinygrad.nn.optim import Optimizer
|
||||
from tinygrad.helpers import FUSE_OPTIM, getenv
|
||||
from tinygrad.uop.ops import UOp, Ops
|
||||
|
||||
STOCHASTIC_ROUND = getenv("STOCHASTIC_ROUND", 0)
|
||||
MASTER_WEIGHTS = getenv("MASTER_WEIGHTS", 0)
|
||||
|
||||
def stochastic_round_bf16(x:Tensor) -> Tensor:
|
||||
bits = x.bitcast(dtypes.uint32)
|
||||
if isinstance(x.device, tuple):
|
||||
shape = x.uop.shard_shape if x.uop.axis is not None else x.shape
|
||||
noise = Tensor(UOp(Ops.MSTACK, dtypes.default_float, tuple(Tensor.rand(*shape, device=d).uop for d in x.device)))
|
||||
else:
|
||||
noise = x.rand_like()
|
||||
noise = (noise * 0xFFFF).cast(dtypes.uint32)
|
||||
return ((bits + noise) & 0xFFFF0000).bitcast(dtypes.float32).cast(dtypes.bfloat16)
|
||||
|
||||
class GradAccClipAdamW(Optimizer):
|
||||
def __init__(self, params:list[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-6, weight_decay=0.0, grad_acc=1, clip_norm=1.0, device=None, fused=FUSE_OPTIM):
|
||||
super().__init__(params, lr, device, fused)
|
||||
self.b1, self.b2, self.eps, self.wd = b1, b2, eps, weight_decay
|
||||
self.b1_t, self.b2_t = (Tensor.ones((1,), dtype=dtypes.float32, device=self.device, requires_grad=False) for _ in [b1, b2])
|
||||
self.m = self._new_optim_param()
|
||||
self.v = self._new_optim_param()
|
||||
self.grad_acc, self.clip_norm = grad_acc, clip_norm
|
||||
self.master_params:list[Tensor]|None = [p.float().contiguous() for p in self.params] if MASTER_WEIGHTS and self.params[0].dtype != dtypes.float32 else None
|
||||
|
||||
def fstep(self, grads:list[Tensor]):
|
||||
if self.fused:
|
||||
out, extra = self._step([], grads)
|
||||
updates = [out[0][self.pos_params[i]:self.pos_params[i+1]].reshape(tt.shape) for i, tt in enumerate(self.params)]
|
||||
else:
|
||||
updates, extra = self._step([], grads)
|
||||
for i, tt in enumerate(self.params): tt.assign(self._apply_update(tt, updates[i], self.master_params[i] if self.master_params else None))
|
||||
to_realize = extra+self.params+self.buffers+(self.master_params or [])
|
||||
|
||||
Tensor.realize(*to_realize)
|
||||
return extra[-1]
|
||||
|
||||
def _step(self, params:list[Tensor], grads:list[Tensor]) -> tuple[list[Tensor], list[Tensor]]:
|
||||
grads = list(grads)
|
||||
|
||||
for i in range(len(grads)):
|
||||
if grads[i].device != self.m[i].device: grads[i] = grads[i].to(self.m[i].device)
|
||||
|
||||
if self.fused:
|
||||
grads[0].assign(grads[0] / self.grad_acc)
|
||||
total_norm = grads[0].float().square().sum().sqrt()
|
||||
grads[0].assign((grads[0] * (self.clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(grads[0].dtype))
|
||||
else:
|
||||
for i in range(len(grads)):
|
||||
grads[i].assign(grads[i] / self.grad_acc)
|
||||
total_norm = Tensor.stack(*[g.float().square().sum() for g in grads]).sum().sqrt().contiguous()
|
||||
for i in range(len(grads)):
|
||||
grads[i].assign((grads[i] * (self.clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(grads[i].dtype))
|
||||
|
||||
ret = []
|
||||
self.b1_t *= self.b1
|
||||
self.b2_t *= self.b2
|
||||
for i, g in enumerate(grads):
|
||||
m_new = self.b1 * self.m[i].float() + (1.0 - self.b1) * g.float()
|
||||
v_new = self.b2 * self.v[i].float() + (1.0 - self.b2) * (g.float() * g.float())
|
||||
self.m[i].assign(m_new.cast(self.m[i].dtype))
|
||||
self.v[i].assign(v_new.cast(self.v[i].dtype))
|
||||
m_hat = m_new / (1.0 - self.b1_t)
|
||||
v_hat = v_new / (1.0 - self.b2_t)
|
||||
up = m_hat / (v_hat.sqrt() + self.eps)
|
||||
ret.append(self.lr * up)
|
||||
return ret, [self.b1_t, self.b2_t] + self.m + self.v + [total_norm]
|
||||
|
||||
def _apply_update(self, t:Tensor, up:Tensor, master:Tensor|None=None) -> Tensor:
|
||||
w = master if master is not None else t
|
||||
wd = self.wd if t.ndim >= 3 else 0.0
|
||||
up = up.float().shard_like(w) + self.lr.to(w.device) * wd * w.detach()
|
||||
new_w = w.detach() - up
|
||||
if master is not None: master.assign(new_w)
|
||||
if STOCHASTIC_ROUND and t.dtype == dtypes.bfloat16: return stochastic_round_bf16(new_w)
|
||||
return new_w.cast(t.dtype)
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export PYTHONPATH="." AMD=1
|
||||
export MODEL="bert"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=1 BS=128 EVAL_BS=128
|
||||
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
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
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export PYTHONPATH="." AMD=1
|
||||
export MODEL="bert"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
|
||||
|
||||
|
||||
+1
-1
@@ -1,7 +1,7 @@
|
||||
#!/bin/bash
|
||||
set -e # Exit on any error
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export PYTHONPATH="." AMD=1
|
||||
export MODEL="bert"
|
||||
export SUBMISSION_PLATFORM="tinybox_8xMI300X"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=NV
|
||||
export PYTHONPATH="." NV=1
|
||||
export MODEL="bert"
|
||||
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
|
||||
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=NV
|
||||
export PYTHONPATH="." NV=1
|
||||
export MODEL="bert"
|
||||
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
|
||||
|
||||
|
||||
+1
-1
@@ -1,7 +1,7 @@
|
||||
#!/bin/bash
|
||||
set -e # Exit on any error
|
||||
|
||||
export PYTHONPATH="." DEV=NV
|
||||
export PYTHONPATH="." NV=1
|
||||
export MODEL="bert"
|
||||
export SUBMISSION_PLATFORM="tinybox_green"
|
||||
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export PYTHONPATH="." AMD=1
|
||||
export MODEL="bert"
|
||||
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
|
||||
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export PYTHONPATH="." AMD=1
|
||||
export MODEL="bert"
|
||||
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
|
||||
|
||||
|
||||
+1
-1
@@ -1,7 +1,7 @@
|
||||
#!/bin/bash
|
||||
set -e # Exit on any error
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export PYTHONPATH="." AMD=1
|
||||
export MODEL="bert"
|
||||
export SUBMISSION_PLATFORM="tinybox_red"
|
||||
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=NV
|
||||
export PYTHONPATH="." NV=1
|
||||
export MODEL="resnet"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
|
||||
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=NV
|
||||
export PYTHONPATH="." NV=1
|
||||
export MODEL="resnet"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
|
||||
|
||||
|
||||
+1
-1
@@ -1,7 +1,7 @@
|
||||
#!/bin/bash
|
||||
set -e # Exit on any error
|
||||
|
||||
export PYTHONPATH="." DEV=NV
|
||||
export PYTHONPATH="." NV=1
|
||||
export MODEL="resnet"
|
||||
export SUBMISSION_PLATFORM="tinybox_green"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export PYTHONPATH="." AMD=1
|
||||
export MODEL="resnet"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
|
||||
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export PYTHONPATH="." AMD=1
|
||||
export MODEL="resnet"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
|
||||
|
||||
|
||||
+1
-1
@@ -1,7 +1,7 @@
|
||||
#!/bin/bash
|
||||
set -e # Exit on any error
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export PYTHONPATH="." AMD=1
|
||||
export MODEL="resnet"
|
||||
export SUBMISSION_PLATFORM="tinybox_red"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=NV
|
||||
export PYTHONPATH="." NV=1
|
||||
export MODEL="retinanet"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
|
||||
export BASEDIR="/raid/datasets/openimages"
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=NV
|
||||
export PYTHONPATH="." NV=1
|
||||
export MODEL="retinanet"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
|
||||
export BASEDIR="/raid/datasets/openimages"
|
||||
|
||||
+1
-1
@@ -1,7 +1,7 @@
|
||||
#!/bin/bash
|
||||
set -e # Exit on any error
|
||||
|
||||
export PYTHONPATH="." DEV=NV
|
||||
export PYTHONPATH="." NV=1
|
||||
export MODEL="retinanet"
|
||||
export SUBMISSION_PLATFORM="tinybox_green"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export PYTHONPATH="." AMD=1
|
||||
export MODEL="retinanet"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
|
||||
export BASEDIR="/raid/datasets/openimages"
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export PYTHONPATH="." AMD=1
|
||||
export MODEL="retinanet"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
|
||||
export BASEDIR="/raid/datasets/openimages"
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export PYTHONPATH="." AMD=1
|
||||
export MODEL="bert"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=1 BS=128 EVAL_BS=128
|
||||
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
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
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export PYTHONPATH="." AMD=1
|
||||
export MODEL="bert"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
|
||||
|
||||
|
||||
+1
-1
@@ -2,7 +2,7 @@
|
||||
set -e # Exit on any error
|
||||
set -o pipefail # Make pipeline fail if any command fails
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export PYTHONPATH="." AMD=1
|
||||
export MODEL="bert"
|
||||
export SUBMISSION_PLATFORM="tinybox_8xMI300X"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=NV
|
||||
export PYTHONPATH="." NV=1
|
||||
export MODEL="bert"
|
||||
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90
|
||||
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=NV
|
||||
export PYTHONPATH="." NV=1
|
||||
export MODEL="bert"
|
||||
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90
|
||||
|
||||
|
||||
+1
-1
@@ -2,7 +2,7 @@
|
||||
set -e # Exit on any error
|
||||
set -o pipefail # Make pipeline fail if any command fails
|
||||
|
||||
export PYTHONPATH="." DEV=NV
|
||||
export PYTHONPATH="." NV=1
|
||||
export MODEL="bert"
|
||||
export SUBMISSION_PLATFORM="tinybox_green"
|
||||
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export PYTHONPATH="." AMD=1
|
||||
export MODEL="bert"
|
||||
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90
|
||||
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export PYTHONPATH="." AMD=1
|
||||
export MODEL="bert"
|
||||
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90
|
||||
|
||||
|
||||
+1
-1
@@ -2,7 +2,7 @@
|
||||
set -e # Exit on any error
|
||||
set -o pipefail # Make pipeline fail if any command fails
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export PYTHONPATH="." AMD=1
|
||||
export MODEL="bert"
|
||||
export SUBMISSION_PLATFORM="tinybox_red"
|
||||
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=NV
|
||||
export PYTHONPATH="." NV=1
|
||||
export MODEL="resnet"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
|
||||
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=NV
|
||||
export PYTHONPATH="." NV=1
|
||||
export MODEL="resnet"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
|
||||
|
||||
|
||||
+1
-1
@@ -2,7 +2,7 @@
|
||||
set -e # Exit on any error
|
||||
set -o pipefail # Make pipeline fail if any command fails
|
||||
|
||||
export PYTHONPATH="." DEV=NV
|
||||
export PYTHONPATH="." NV=1
|
||||
export MODEL="resnet"
|
||||
export SUBMISSION_PLATFORM="tinybox_green"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export PYTHONPATH="." AMD=1
|
||||
export MODEL="resnet"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
|
||||
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export PYTHONPATH="." AMD=1
|
||||
export MODEL="resnet"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
|
||||
|
||||
|
||||
+1
-1
@@ -2,7 +2,7 @@
|
||||
set -e # Exit on any error
|
||||
set -o pipefail # Make pipeline fail if any command fails
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export PYTHONPATH="." AMD=1
|
||||
export MODEL="resnet"
|
||||
export SUBMISSION_PLATFORM="tinybox_red"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=NV
|
||||
export PYTHONPATH="." NV=1
|
||||
export MODEL="retinanet"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
|
||||
export BASEDIR="/raid/datasets/openimages"
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=NV
|
||||
export PYTHONPATH="." NV=1
|
||||
export MODEL="retinanet"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
|
||||
export BASEDIR="/raid/datasets/openimages"
|
||||
|
||||
+1
-1
@@ -2,7 +2,7 @@
|
||||
set -e # Exit on any error
|
||||
set -o pipefail # Make pipeline fail if any command fails
|
||||
|
||||
export PYTHONPATH="." DEV=NV
|
||||
export PYTHONPATH="." NV=1
|
||||
export MODEL="retinanet"
|
||||
export SUBMISSION_PLATFORM="tinybox_green"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export PYTHONPATH="." AMD=1
|
||||
export MODEL="retinanet"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
|
||||
export BASEDIR="/raid/datasets/openimages"
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export PYTHONPATH="." AMD=1
|
||||
export MODEL="retinanet"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
|
||||
export BASEDIR="/raid/datasets/openimages"
|
||||
|
||||
+2
-2
@@ -1,10 +1,10 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export PYTHONPATH="." AMD=1
|
||||
export MODEL="bert"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=1 BS=128 EVAL_BS=128
|
||||
|
||||
export CHECK_OOB=0
|
||||
export IGNORE_OOB=1
|
||||
|
||||
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
|
||||
|
||||
+2
-2
@@ -1,11 +1,11 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
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 CHECK_OOB=0
|
||||
export IGNORE_OOB=1
|
||||
export REWRITE_STACK_LIMIT=500000
|
||||
|
||||
export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
|
||||
|
||||
+2
-2
@@ -1,6 +1,6 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export PYTHONPATH="." AMD=1
|
||||
export MODEL="bert"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
|
||||
|
||||
@@ -8,7 +8,7 @@ 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 TRAIN_STEPS=3900
|
||||
|
||||
export CHECK_OOB=0
|
||||
export IGNORE_OOB=1
|
||||
export REWRITE_STACK_LIMIT=500000
|
||||
|
||||
export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
|
||||
|
||||
+2
-2
@@ -2,7 +2,7 @@
|
||||
set -e # Exit on any error
|
||||
set -o pipefail # Make pipeline fail if any command fails
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export PYTHONPATH="." AMD=1
|
||||
export MODEL="bert"
|
||||
export SUBMISSION_PLATFORM="tinybox_8xMI300X"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
|
||||
@@ -11,7 +11,7 @@ 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 TRAIN_STEPS=3900
|
||||
|
||||
export CHECK_OOB=0
|
||||
export IGNORE_OOB=1
|
||||
export REWRITE_STACK_LIMIT=500000
|
||||
|
||||
export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
|
||||
|
||||
-20
@@ -1,20 +0,0 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=AMD DEBUG=0 JIT=1 FLASH_ATTENTION=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 CHECK_OOB=0
|
||||
export REWRITE_STACK_LIMIT=5000000
|
||||
|
||||
export BEAM=0 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
|
||||
-24
@@ -1,24 +0,0 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
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 CHECK_OOB=0
|
||||
export REWRITE_STACK_LIMIT=5000000
|
||||
|
||||
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 BEAM_TIMEOUT_SEC=15
|
||||
export FP8_TRAIN=1
|
||||
# search
|
||||
IGNORE_BEAM_CACHE=1 BENCHMARK=10 BERT_LAYERS=2 RUNMLPERF=0 python3 examples/mlperf/model_train.py
|
||||
|
||||
export WANDB=1 PARALLEL=0
|
||||
|
||||
RUNMLPERF=1 python3 examples/mlperf/model_train.py
|
||||
+2
-2
@@ -2,7 +2,7 @@
|
||||
set -e # Exit on any error
|
||||
set -o pipefail # Make pipeline fail if any command fails
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export PYTHONPATH="." AMD=1
|
||||
export MODEL="bert"
|
||||
export SUBMISSION_PLATFORM="tinybox_8xMI350X"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
|
||||
@@ -11,7 +11,7 @@ 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 TRAIN_STEPS=3900
|
||||
|
||||
export CHECK_OOB=0
|
||||
export IGNORE_OOB=1
|
||||
export REWRITE_STACK_LIMIT=5000000
|
||||
|
||||
export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
|
||||
|
||||
+2
-2
@@ -1,10 +1,10 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=NV
|
||||
export PYTHONPATH="." NV=1
|
||||
export MODEL="bert"
|
||||
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=72 EVAL_BS=72
|
||||
|
||||
export CHECK_OOB=0
|
||||
export IGNORE_OOB=1
|
||||
export REWRITE_STACK_LIMIT=500000
|
||||
|
||||
export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
|
||||
|
||||
+2
-2
@@ -1,10 +1,10 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=NV
|
||||
export PYTHONPATH="." NV=1
|
||||
export MODEL="bert"
|
||||
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=72 EVAL_BS=72
|
||||
|
||||
export CHECK_OOB=0
|
||||
export IGNORE_OOB=1
|
||||
export REWRITE_STACK_LIMIT=500000
|
||||
|
||||
export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
|
||||
|
||||
+2
-2
@@ -2,12 +2,12 @@
|
||||
set -e # Exit on any error
|
||||
set -o pipefail # Make pipeline fail if any command fails
|
||||
|
||||
export PYTHONPATH="." DEV=NV
|
||||
export PYTHONPATH="." NV=1
|
||||
export MODEL="bert"
|
||||
export SUBMISSION_PLATFORM="tinybox_green"
|
||||
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=72 EVAL_BS=72
|
||||
|
||||
export CHECK_OOB=0
|
||||
export IGNORE_OOB=1
|
||||
export REWRITE_STACK_LIMIT=500000
|
||||
|
||||
export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
|
||||
|
||||
+2
-2
@@ -1,10 +1,10 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export PYTHONPATH="." AMD=1
|
||||
export MODEL="bert"
|
||||
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
|
||||
|
||||
export CHECK_OOB=0
|
||||
export IGNORE_OOB=1
|
||||
export REWRITE_STACK_LIMIT=500000
|
||||
|
||||
export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
|
||||
|
||||
+2
-2
@@ -1,10 +1,10 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export PYTHONPATH="." AMD=1
|
||||
export MODEL="bert"
|
||||
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
|
||||
|
||||
export CHECK_OOB=0
|
||||
export IGNORE_OOB=1
|
||||
export REWRITE_STACK_LIMIT=500000
|
||||
|
||||
export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
|
||||
|
||||
+2
-2
@@ -2,12 +2,12 @@
|
||||
set -e # Exit on any error
|
||||
set -o pipefail # Make pipeline fail if any command fails
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export PYTHONPATH="." AMD=1
|
||||
export MODEL="bert"
|
||||
export SUBMISSION_PLATFORM="tinybox_red"
|
||||
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
|
||||
|
||||
export CHECK_OOB=0
|
||||
export IGNORE_OOB=1
|
||||
export REWRITE_STACK_LIMIT=500000
|
||||
|
||||
export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
|
||||
|
||||
-38
@@ -1,38 +0,0 @@
|
||||
#!/usr/bin/env bash
|
||||
|
||||
export PYTHONPATH="."
|
||||
export DEV=${DEV:-AMD}
|
||||
export EMULATE="AMD_CDNA4"
|
||||
export CHECK_OOB=0
|
||||
export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
|
||||
export DEVICE_IN_FUNCTION_BUG=1
|
||||
|
||||
export DEBUG=${DEBUG:-2}
|
||||
export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
|
||||
export ALL2ALL=${ALL2ALL:-1}
|
||||
export USE_ATOMICS=${USE_ATOMICS:-0}
|
||||
export ASM_GEMM=${ASM_GEMM:-1}
|
||||
export WQKV=${WQKV:-1}
|
||||
export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}
|
||||
|
||||
export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
|
||||
export DP=${DP:-1} MP=${MP:-8}
|
||||
export BS=${BS:-1} EVAL_BS=${EVAL_BS:-1} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
|
||||
|
||||
export MODEL="llama3"
|
||||
export BASEDIR="/raid/datasets/c4/"
|
||||
export LLAMA3_SIZE=${LLAMA3_SIZE:-"405B"}
|
||||
export SEQLEN=${SEQLEN:-8192}
|
||||
|
||||
export SEED=${SEED:-5760}
|
||||
export DATA_SEED=${DATA_SEED:-5760}
|
||||
|
||||
export JITBEAM=${JITBEAM:-3}
|
||||
export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
|
||||
|
||||
export FAKEDATA=1 BENCHMARK=10
|
||||
if [ -z "$FULL_LAYERS" ]; then
|
||||
export LLAMA_LAYERS=2
|
||||
fi
|
||||
|
||||
python3 examples/mlperf/model_train.py
|
||||
-32
@@ -1,32 +0,0 @@
|
||||
#!/usr/bin/env bash
|
||||
|
||||
export PYTHONPATH="."
|
||||
export DEV=${DEV:-AMD}
|
||||
export EMULATE="AMD_CDNA4"
|
||||
export CHECK_OOB=0
|
||||
export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
|
||||
|
||||
export DEBUG=${DEBUG:-0}
|
||||
export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
|
||||
export ALL2ALL=${ALL2ALL:-1}
|
||||
export USE_ATOMICS=${USE_ATOMICS:-0}
|
||||
export ASM_GEMM=${ASM_GEMM:-1}
|
||||
export WQKV=${WQKV:-1}
|
||||
export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}
|
||||
|
||||
export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
|
||||
export DP=${DP:-1} MP=${MP:-8}
|
||||
export BS=${BS:-1} EVAL_BS=${EVAL_BS:-1} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-1152}
|
||||
|
||||
export MODEL="llama3"
|
||||
export BASEDIR="/raid/datasets/c4/"
|
||||
export LLAMA3_SIZE=${LLAMA3_SIZE:-"405B"}
|
||||
export SEQLEN=${SEQLEN:-8192}
|
||||
|
||||
export SEED=${SEED:-$RANDOM}
|
||||
export DATA_SEED=${DATA_SEED:-5760}
|
||||
|
||||
export JITBEAM=${JITBEAM:-3}
|
||||
export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
|
||||
|
||||
python3 examples/mlperf/model_train.py
|
||||
-42
@@ -1,42 +0,0 @@
|
||||
#!/usr/bin/env bash
|
||||
|
||||
export PYTHONPATH="."
|
||||
export DEV=${DEV:-AMD}
|
||||
export EMULATE="AMD_CDNA4"
|
||||
export CHECK_OOB=0
|
||||
export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
|
||||
export DEVICE_IN_FUNCTION_BUG=1
|
||||
|
||||
export DEBUG=${DEBUG:-2}
|
||||
export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
|
||||
export ALL2ALL=${ALL2ALL:-1}
|
||||
export USE_ATOMICS=${USE_ATOMICS:-1}
|
||||
export ASM_GEMM=${ASM_GEMM:-1}
|
||||
export WQKV=${WQKV:-0}
|
||||
|
||||
export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
|
||||
export DP=${DP:-8} MP=${MP:-1} BS=${BS:-8} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-4}
|
||||
export GBS=$((BS * GRADIENT_ACC_STEPS))
|
||||
|
||||
export MODEL="llama3"
|
||||
export BASEDIR="/raid/datasets/c4-8b/"
|
||||
export SMALL=1
|
||||
export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
|
||||
export EVAL_TARGET=3.3 EVAL_FREQ=12288
|
||||
export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 MAX_STEPS=1200000
|
||||
export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
|
||||
export SAMPLES=$((MAX_STEPS * GBS))
|
||||
export SEQLEN=${SEQLEN:-8192}
|
||||
|
||||
export SEED=${SEED:-5760}
|
||||
export DATA_SEED=${DATA_SEED:-5760}
|
||||
|
||||
export JITBEAM=${JITBEAM:-3}
|
||||
export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
|
||||
|
||||
export FAKEDATA=1 BENCHMARK=10
|
||||
if [ -z "$FULL_LAYERS" ]; then
|
||||
export LLAMA_LAYERS=2
|
||||
fi
|
||||
|
||||
python3 examples/mlperf/model_train.py
|
||||
-43
@@ -1,43 +0,0 @@
|
||||
#!/usr/bin/env bash
|
||||
|
||||
export PYTHONPATH="."
|
||||
export DEV=${DEV:-AMD}
|
||||
export EMULATE="AMD_CDNA4"
|
||||
export CHECK_OOB=0
|
||||
export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
|
||||
export DEVICE_IN_FUNCTION_BUG=1
|
||||
|
||||
export DEBUG=${DEBUG:-2}
|
||||
export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
|
||||
export ALL2ALL=${ALL2ALL:-1}
|
||||
export USE_ATOMICS=${USE_ATOMICS:-0}
|
||||
export ASM_GEMM=${ASM_GEMM:-1}
|
||||
export WQKV=${WQKV:-1}
|
||||
export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}
|
||||
|
||||
export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
|
||||
export DP=${DP:-1} MP=${MP:-8} BS=${BS:-1} EVAL_BS=${EVAL_BS:-1} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
|
||||
export GBS=$((BS * GRADIENT_ACC_STEPS))
|
||||
|
||||
export MODEL="llama3"
|
||||
export BASEDIR="/raid/datasets/c4-8b/"
|
||||
export SMALL=1
|
||||
export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
|
||||
export EVAL_TARGET=3.3 EVAL_FREQ=12288
|
||||
export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 MAX_STEPS=1200000
|
||||
export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
|
||||
export SAMPLES=$((MAX_STEPS * GBS))
|
||||
export SEQLEN=${SEQLEN:-8192}
|
||||
|
||||
export SEED=${SEED:-5760}
|
||||
export DATA_SEED=${DATA_SEED:-5760}
|
||||
|
||||
export JITBEAM=${JITBEAM:-3}
|
||||
export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
|
||||
|
||||
export FAKEDATA=1 BENCHMARK=10
|
||||
if [ -z "$FULL_LAYERS" ]; then
|
||||
export LLAMA_LAYERS=2
|
||||
fi
|
||||
|
||||
python3 examples/mlperf/model_train.py
|
||||
-37
@@ -1,37 +0,0 @@
|
||||
#!/usr/bin/env bash
|
||||
|
||||
export PYTHONPATH="."
|
||||
export DEV=${DEV:-AMD}
|
||||
export EMULATE="AMD_CDNA4"
|
||||
export CHECK_OOB=0
|
||||
export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
|
||||
export DEVICE_IN_FUNCTION_BUG=1
|
||||
|
||||
export DEBUG=${DEBUG:-0}
|
||||
export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
|
||||
export ALL2ALL=${ALL2ALL:-1}
|
||||
export USE_ATOMICS=${USE_ATOMICS:-1}
|
||||
export ASM_GEMM=${ASM_GEMM:-1}
|
||||
export WQKV=${WQKV:-0}
|
||||
|
||||
export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
|
||||
export DP=${DP:-8} MP=${MP:-1} BS=${BS:-8} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-4}
|
||||
export GBS=$((BS * GRADIENT_ACC_STEPS))
|
||||
|
||||
export MODEL="llama3"
|
||||
export BASEDIR="/raid/datasets/c4-8b/"
|
||||
export SMALL=1
|
||||
export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
|
||||
export EVAL_TARGET=3.3 EVAL_FREQ=12288
|
||||
export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 MAX_STEPS=1200000
|
||||
export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
|
||||
export SAMPLES=$((MAX_STEPS * GBS))
|
||||
export SEQLEN=${SEQLEN:-8192}
|
||||
|
||||
export SEED=${SEED:-$RANDOM}
|
||||
export DATA_SEED=${DATA_SEED:-5760}
|
||||
|
||||
export JITBEAM=${JITBEAM:-3}
|
||||
export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
|
||||
|
||||
python3 examples/mlperf/model_train.py
|
||||
-38
@@ -1,38 +0,0 @@
|
||||
#!/usr/bin/env bash
|
||||
|
||||
export PYTHONPATH="."
|
||||
export DEV=${DEV:-AMD}
|
||||
export EMULATE="AMD_CDNA4"
|
||||
export CHECK_OOB=0
|
||||
export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
|
||||
export DEVICE_IN_FUNCTION_BUG=1
|
||||
|
||||
export DEBUG=${DEBUG:-0}
|
||||
export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
|
||||
export ALL2ALL=${ALL2ALL:-1}
|
||||
export USE_ATOMICS=${USE_ATOMICS:-0}
|
||||
export ASM_GEMM=${ASM_GEMM:-1}
|
||||
export WQKV=${WQKV:-1}
|
||||
export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}
|
||||
|
||||
export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
|
||||
export DP=${DP:-1} MP=${MP:-8} BS=${BS:-1} EVAL_BS=${EVAL_BS:-1} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-32}
|
||||
export GBS=$((BS * GRADIENT_ACC_STEPS))
|
||||
|
||||
export MODEL="llama3"
|
||||
export BASEDIR="/raid/datasets/c4-8b/"
|
||||
export SMALL=1
|
||||
export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
|
||||
export EVAL_TARGET=3.3 EVAL_FREQ=12288
|
||||
export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 MAX_STEPS=1200000
|
||||
export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
|
||||
export SAMPLES=$((MAX_STEPS * GBS))
|
||||
export SEQLEN=${SEQLEN:-8192}
|
||||
|
||||
export SEED=${SEED:-$RANDOM}
|
||||
export DATA_SEED=${DATA_SEED:-5760}
|
||||
|
||||
export JITBEAM=${JITBEAM:-3}
|
||||
export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
|
||||
|
||||
python3 examples/mlperf/model_train.py
|
||||
-5
@@ -1,5 +0,0 @@
|
||||
#!/bin/bash
|
||||
export BENCHMARK=5
|
||||
export EVAL_BS=0
|
||||
VIZ=${VIZ:--1} examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh
|
||||
extra/viz/cli.py --profile --device "AMD" | head -23
|
||||
-10
@@ -1,10 +0,0 @@
|
||||
#!/bin/bash
|
||||
export BENCHMARK=5
|
||||
export EVAL_BS=0
|
||||
export FAKEDATA=1
|
||||
export NULL_ALLOW_COPYOUT=1
|
||||
export HIP_VISIBLE_DEVICES=""
|
||||
export DEV=NULL
|
||||
export JITBEAM=0
|
||||
export LLAMA_LAYERS=${LLAMA_LAYERS:-"2"}
|
||||
time examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user